From 850f2516d76b65df1060f65c771e36ed1a74096a Mon Sep 17 00:00:00 2001
From: Francesco Ceccon <francesco@ceccon.me>
Date: Fri, 21 Jul 2023 23:48:25 +0200
Subject: [PATCH] o11y: add option for json logger

**Summary**

This option enables json logging in production by setting the
`RUST_LOG_FORMAT=json` environment variable.
---
 Cargo.lock               | 13 +++++++
 Cargo.toml               |  2 +-
 observability/src/lib.rs | 74 ++++++++++++++++++++++++----------------
 3 files changed, 58 insertions(+), 31 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4efaf515..65fb68e8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -7826,6 +7826,16 @@ dependencies = [
  "tracing-subscriber",
 ]
 
+[[package]]
+name = "tracing-serde"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc6b213177105856957181934e4920de57730fc69bf42c37ee5bb664d406d9e1"
+dependencies = [
+ "serde",
+ "tracing-core",
+]
+
 [[package]]
 name = "tracing-subscriber"
 version = "0.3.17"
@@ -7836,12 +7846,15 @@ dependencies = [
  "nu-ansi-term",
  "once_cell",
  "regex",
+ "serde",
+ "serde_json",
  "sharded-slab",
  "smallvec 1.10.0",
  "thread_local",
  "tracing",
  "tracing-core",
  "tracing-log",
+ "tracing-serde",
 ]
 
 [[package]]
diff --git a/Cargo.toml b/Cargo.toml
index b995ecbe..f48dc4c9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -68,7 +68,7 @@ tower = "0.4.13"
 tracing = { version = "0.1.36", features = ["max_level_trace", "release_max_level_debug"] }
 tracing-futures = { version = "0.2.5", features = ["tokio", "futures-03"] }
 tracing-opentelemetry = "0.18.0"
-tracing-subscriber = { version = "0.3.15", features = ["std", "env-filter"] }
+tracing-subscriber = { version = "0.3.15", features = ["std", "env-filter", "json"] }
 tracing-tree = "0.2.2"
 quickcheck = "1.0.3"
 quickcheck_macros = "1.0.0"
diff --git a/observability/src/lib.rs b/observability/src/lib.rs
index af9c4e5e..e5921609 100644
--- a/observability/src/lib.rs
+++ b/observability/src/lib.rs
@@ -12,17 +12,19 @@ use opentelemetry::{
     trace::TraceError,
 };
 use opentelemetry_otlp::WithExportConfig;
-use tracing::dispatcher::SetGlobalDefaultError;
+use tracing::{dispatcher::SetGlobalDefaultError, Subscriber};
 
 pub use opentelemetry::metrics::{ObservableCounter, ObservableGauge};
 pub use opentelemetry::{Context, Key, KeyValue};
 use tracing_opentelemetry::MetricsLayer;
-use tracing_subscriber::{filter, prelude::*, EnvFilter};
+use tracing_subscriber::{prelude::*, registry::LookupSpan, EnvFilter, Layer};
 
 pub use opentelemetry::metrics::{Counter, Meter};
 
 const OTEL_SDK_DISABLED: &str = "OTEL_SDK_DISABLED";
 
+pub type BoxedLayer<S> = Box<dyn Layer<S> + Send + Sync>;
+
 #[derive(Debug, thiserror::Error)]
 pub enum OpenTelemetryInitError {
     #[error("error setting global default subscriber")]
@@ -49,30 +51,26 @@ pub fn init_opentelemetry() -> Result<(), OpenTelemetryInitError> {
         std::env::set_var("RUST_LOG", "info");
     }
 
-    if sdk_disabled {
-        init_opentelemetry_no_sdk()
-    } else {
-        init_opentelemetry_with_sdk()
+    let mut layers = vec![stdout()];
+
+    if !sdk_disabled {
+        let otel_layer = otel()?;
+        layers.push(otel_layer);
     }
-}
 
-fn init_opentelemetry_no_sdk() -> Result<(), OpenTelemetryInitError> {
-    let log_env_filter =
-        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("INFO"));
-    let logtree_layer = tracing_tree::HierarchicalLayer::new(2).and_then(log_env_filter);
+    tracing_subscriber::registry().with(layers).init();
 
-    tracing_subscriber::Registry::default()
-        .with(logtree_layer)
-        .init();
     Ok(())
 }
 
-fn init_opentelemetry_with_sdk() -> Result<(), OpenTelemetryInitError> {
+fn otel<S>() -> Result<BoxedLayer<S>, OpenTelemetryInitError>
+where
+    S: Subscriber + Send + Sync,
+    for<'a> S: LookupSpan<'a>,
+{
     // filter traces by crate/level
     let otel_env_filter =
         EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("INFO"));
-    let log_env_filter =
-        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("INFO"));
 
     // Both tracer and meter are configured with environment variables.
     let meter = opentelemetry_otlp::new_pipeline()
@@ -96,19 +94,35 @@ fn init_opentelemetry_with_sdk() -> Result<(), OpenTelemetryInitError> {
     let otel_metrics_layer = MetricsLayer::new(meter);
     let otel_layer = otel_trace_layer
         .and_then(otel_metrics_layer)
-        .and_then(otel_env_filter);
-
-    // display traces on stdout
-    let logtree_layer = tracing_tree::HierarchicalLayer::new(2)
-        .and_then(log_env_filter)
-        .with_filter(filter::filter_fn(|metadata| {
-            metadata.fields().field("data.is_metrics").is_none()
-        }));
+        .and_then(otel_env_filter)
+        .boxed();
+    Ok(otel_layer)
+}
 
-    tracing_subscriber::Registry::default()
-        .with(otel_layer)
-        .with(logtree_layer)
-        .init();
+fn stdout<S>() -> BoxedLayer<S>
+where
+    S: Subscriber,
+    for<'a> S: LookupSpan<'a>,
+{
+    let log_env_filter =
+        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("INFO"));
 
-    Ok(())
+    let json_fmt = std::env::var("RUST_LOG_FORMAT")
+        .map(|val| val == "json")
+        .unwrap_or(false);
+
+    if json_fmt {
+        tracing_subscriber::fmt::layer()
+            .with_ansi(false)
+            .with_target(true)
+            .json()
+            .with_filter(log_env_filter)
+            .boxed()
+    } else {
+        tracing_subscriber::fmt::layer()
+            .with_ansi(true)
+            .with_target(true)
+            .with_filter(log_env_filter)
+            .boxed()
+    }
 }