From d9922b370ab680602e7669a1480b6fa6694aaa1e Mon Sep 17 00:00:00 2001
From: Adam Singer <adam-singer@users.noreply.github.com>
Date: Fri, 12 Jul 2024 23:58:35 -0700
Subject: [PATCH] Introduce NL_LOG to control logging format (#1154)

`NL_LOG` envionment variable is used to control
the logging format that `tracing_subscriber`
emits out. We currently want to allow for better
ingestion of formatted logs in json.

`NL_LOG` can be configured with `pretty`, `json`,
or `compact`.
---
 nativelink-util/Cargo.toml |  2 +-
 nativelink-util/src/lib.rs | 52 +++++++++++++++++++++++++++-----------
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/nativelink-util/Cargo.toml b/nativelink-util/Cargo.toml
index 679bc51fc..c4737cc4a 100644
--- a/nativelink-util/Cargo.toml
+++ b/nativelink-util/Cargo.toml
@@ -30,7 +30,7 @@ tokio-stream = { version = "0.1.15", features = ["sync"] }
 tokio-util = { version = "0.7.11" }
 tonic = { version = "0.11.0", features = ["tls"] }
 tracing = "0.1.40"
-tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+tracing-subscriber = { version = "0.3.18", features = ["env-filter", "json"] }
 uuid = { version = "1.8.0", features = ["v4", "serde"] }
 
 [dev-dependencies]
diff --git a/nativelink-util/src/lib.rs b/nativelink-util/src/lib.rs
index 717985274..2811a4d68 100644
--- a/nativelink-util/src/lib.rs
+++ b/nativelink-util/src/lib.rs
@@ -53,22 +53,44 @@ pub fn init_tracing() -> Result<(), nativelink_error::Error> {
         .with_default_directive(tracing::metadata::LevelFilter::WARN.into())
         .from_env_lossy();
 
+    // Setup tracing logger for multiple format types, compact, json, and pretty as a single layer.
+    // Configuration for log format comes from environment variable NL_LOG_FMT due to subscribers
+    // being configured before config parsing.
+    let nl_log_fmt = std::env::var("NL_LOG").unwrap_or_else(|_| "pretty".to_string());
+    // Layers vector is used for due to how tracing_subscriber::fmt::layer builds type signature
+    // not being able to unify a single trait type before being boxed. For example see
+    // https://docs.rs/tracing-subscriber/0.3.18/tracing_subscriber/layer/index.html
+    let mut layers = Vec::new();
+    match nl_log_fmt.as_str() {
+        "compact" => layers.push(
+            tracing_subscriber::fmt::layer()
+                .compact()
+                .with_timer(tracing_subscriber::fmt::time::time())
+                .with_filter(env_filter)
+                .boxed(),
+        ),
+        "json" => layers.push(
+            tracing_subscriber::fmt::layer()
+                .json()
+                .with_timer(tracing_subscriber::fmt::time::time())
+                .with_filter(env_filter)
+                .boxed(),
+        ),
+        _ => layers.push(
+            tracing_subscriber::fmt::layer()
+                .pretty()
+                .with_timer(tracing_subscriber::fmt::time::time())
+                .with_filter(env_filter)
+                .boxed(),
+        ),
+    };
+
+    // Add a console subscriber if the feature is enabled, see tokio-console for a client console.
+    // https://crates.io/crates/tokio-console
     if cfg!(feature = "enable_tokio_console") {
-        tracing_subscriber::registry()
-            .with(console_subscriber::spawn())
-            .with(
-                tracing_subscriber::fmt::layer()
-                    .pretty()
-                    .with_timer(tracing_subscriber::fmt::time::time())
-                    .with_filter(env_filter),
-            )
-            .init();
-    } else {
-        tracing_subscriber::fmt()
-            .pretty()
-            .with_timer(tracing_subscriber::fmt::time::time())
-            .with_env_filter(env_filter)
-            .init();
+        layers.push(console_subscriber::spawn().boxed());
     }
+
+    tracing_subscriber::registry().with(layers).init();
     Ok(())
 }