From 706b6787d340495d894d6a59db4049a18e5b7884 Mon Sep 17 00:00:00 2001
From: Daniel Lundin <dln@arity.se>
Date: Wed, 25 Dec 2024 00:11:50 +0100
Subject: [PATCH] chore: clean up tracing subscriber setup

---
 agent/src/main.rs          |  9 ++++--
 instrumentation/src/lib.rs | 61 +++++++++++++++++++++++---------------
 2 files changed, 43 insertions(+), 27 deletions(-)

diff --git a/agent/src/main.rs b/agent/src/main.rs
index aa60347..d232061 100644
--- a/agent/src/main.rs
+++ b/agent/src/main.rs
@@ -6,12 +6,15 @@ mod patagia_api;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
-struct Cli {}
+struct Cli {
+    #[arg(long = "telemetry-otlp-endpoint", default_value = "http://localhost:4317", value_name = "OTEL_EXPORTER_OTLP_ENDPOINT")]
+    otlp_endpoint: Option<String>,
+}
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    let _args = Cli::parse();
-    let _tracing = instrumentation::init_tracing_subscriber()?;
+    let args = Cli::parse();
+    let _tracing = instrumentation::init_tracing(args.otlp_endpoint.as_ref())?;
 
     tracing::info!("Patagia Agent");
 
diff --git a/instrumentation/src/lib.rs b/instrumentation/src/lib.rs
index b05e0fc..d6e1df4 100644
--- a/instrumentation/src/lib.rs
+++ b/instrumentation/src/lib.rs
@@ -4,6 +4,7 @@ use opentelemetry::{trace::TracerProvider as _, KeyValue};
 use opentelemetry_otlp::{WithExportConfig, WithTonicConfig};
 use opentelemetry_sdk::{
     metrics::{MeterProviderBuilder, PeriodicReader, SdkMeterProvider},
+    propagation::TraceContextPropagator,
     runtime,
     trace::{RandomIdGenerator, Sampler, TracerProvider},
     Resource,
@@ -74,32 +75,44 @@ fn init_tracer_provider(otel_endpoint: &String) -> Result<TracerProvider> {
 
 // Initialize tracing-subscriber and return TracingGuard for opentelemetry-related termination processing
 pub fn init_tracing(otel_endpoint: Option<&String>) -> Result<TracingGuard> {
-    let sub = tracing_subscriber::registry()
-        .with(tracing_subscriber::EnvFilter::from_default_env())
-        .with(tracing_subscriber::fmt::layer());
+    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
+        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(""));
 
-    match otel_endpoint {
-        None => {
-            sub.init();
-            Ok(TracingGuard {
-                meter_provider: None,
-                tracer_provider: None,
-            })
-        }
+    let fmt_layer = tracing_subscriber::fmt::layer().with_writer(std::io::stderr);
 
-        Some(otel_endpoint) => {
-            let meter_provider = init_meter_provider(otel_endpoint)?;
-            let tracer_provider = init_tracer_provider(otel_endpoint)?;
-            let tracer = tracer_provider.tracer("tracing-otel-subscriber");
-            sub.with(MetricsLayer::new(meter_provider.clone()))
-                .with(OpenTelemetryLayer::new(tracer))
-                .init();
-            Ok(TracingGuard {
-                meter_provider: Some(meter_provider),
-                tracer_provider: Some(tracer_provider),
-            })
-        }
-    }
+    let meter_provider = match otel_endpoint {
+        Some(endpoint) => Some(init_meter_provider(endpoint)?),
+        None => None,
+    };
+
+    let metrics_layer = match meter_provider {
+        Some(ref p) => Some(MetricsLayer::new(p.to_owned())),
+        None => None,
+    };
+
+    let tracer_provider = match otel_endpoint {
+        Some(endpoint) => Some(init_tracer_provider(endpoint)?),
+        None => None,
+    };
+
+    let trace_layer = match tracer_provider {
+        Some(ref p) => Some(OpenTelemetryLayer::new(p.tracer("tracing-otel-subscriber"))),
+        None => None,
+    };
+
+    opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());
+
+    tracing_subscriber::registry()
+        .with(env_filter)
+        .with(fmt_layer)
+        .with(metrics_layer)
+        .with(trace_layer)
+        .init();
+
+    Ok(TracingGuard {
+        meter_provider,
+        tracer_provider,
+    })
 }
 
 pub struct TracingGuard {