From a2b9ff0d40cb3904652aa3c839440dcf6691210f Mon Sep 17 00:00:00 2001
From: Vlad Lazar <vlad@neon.tech>
Date: Tue, 11 Feb 2025 11:33:41 +0100
Subject: [PATCH] tracing-utils: allow for explicit export config

---
 compute_tools/src/logger.rs   |  3 ++-
 libs/tracing-utils/src/lib.rs | 27 +++++++++++++++++++++------
 proxy/src/logging.rs          |  3 ++-
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/compute_tools/src/logger.rs b/compute_tools/src/logger.rs
index a65614e94e..c36f302f99 100644
--- a/compute_tools/src/logger.rs
+++ b/compute_tools/src/logger.rs
@@ -24,7 +24,8 @@ pub async fn init_tracing_and_logging(default_log_level: &str) -> anyhow::Result
         .with_writer(std::io::stderr);
 
     // Initialize OpenTelemetry
-    let otlp_layer = tracing_utils::init_tracing("compute_ctl").await;
+    let otlp_layer =
+        tracing_utils::init_tracing("compute_ctl", tracing_utils::ExportConfig::default()).await;
 
     // Put it all together
     tracing_subscriber::registry()
diff --git a/libs/tracing-utils/src/lib.rs b/libs/tracing-utils/src/lib.rs
index 72f94d61e4..0ee97f5f79 100644
--- a/libs/tracing-utils/src/lib.rs
+++ b/libs/tracing-utils/src/lib.rs
@@ -21,7 +21,7 @@
 //!         .with_writer(std::io::stderr);
 //!
 //!     // Initialize OpenTelemetry. Exports tracing spans as OpenTelemetry traces
-//!     let otlp_layer = tracing_utils::init_tracing("my_application").await;
+//!     let otlp_layer = tracing_utils::init_tracing("my_application", tracing_utils::ExportConfig::default()).await;
 //!
 //!     // Put it all together
 //!     tracing_subscriber::registry()
@@ -38,6 +38,8 @@ pub mod http;
 
 use opentelemetry::KeyValue;
 use opentelemetry::trace::TracerProvider;
+use opentelemetry_otlp::WithExportConfig;
+pub use opentelemetry_otlp::{ExportConfig, Protocol};
 use tracing::Subscriber;
 use tracing_subscriber::Layer;
 use tracing_subscriber::registry::LookupSpan;
@@ -69,19 +71,28 @@ use tracing_subscriber::registry::LookupSpan;
 ///
 /// This doesn't block, but is marked as 'async' to hint that this must be called in
 /// asynchronous execution context.
-pub async fn init_tracing<S>(service_name: &str) -> Option<impl Layer<S>>
+pub async fn init_tracing<S>(
+    service_name: &str,
+    export_config: ExportConfig,
+) -> Option<impl Layer<S>>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
     if std::env::var("OTEL_SDK_DISABLED") == Ok("true".to_string()) {
         return None;
     };
-    Some(init_tracing_internal(service_name.to_string()))
+    Some(init_tracing_internal(
+        service_name.to_string(),
+        export_config,
+    ))
 }
 
 /// Like `init_tracing`, but creates a separate tokio Runtime for the tracing
 /// tasks.
-pub fn init_tracing_without_runtime<S>(service_name: &str) -> Option<impl Layer<S>>
+pub fn init_tracing_without_runtime<S>(
+    service_name: &str,
+    export_config: ExportConfig,
+) -> Option<impl Layer<S>>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
@@ -112,16 +123,20 @@ where
     ));
     let _guard = runtime.enter();
 
-    Some(init_tracing_internal(service_name.to_string()))
+    Some(init_tracing_internal(
+        service_name.to_string(),
+        export_config,
+    ))
 }
 
-fn init_tracing_internal<S>(service_name: String) -> impl Layer<S>
+fn init_tracing_internal<S>(service_name: String, export_config: ExportConfig) -> impl Layer<S>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
     // Sets up exporter from the OTEL_EXPORTER_* environment variables.
     let exporter = opentelemetry_otlp::SpanExporter::builder()
         .with_http()
+        .with_export_config(export_config)
         .build()
         .expect("could not initialize opentelemetry exporter");
 
diff --git a/proxy/src/logging.rs b/proxy/src/logging.rs
index 6f9845fd6e..454fe81357 100644
--- a/proxy/src/logging.rs
+++ b/proxy/src/logging.rs
@@ -46,7 +46,8 @@ pub async fn init() -> anyhow::Result<LoggingGuard> {
                 .expect("this should be a valid filter directive"),
         );
 
-    let otlp_layer = tracing_utils::init_tracing("proxy").await;
+    let otlp_layer =
+        tracing_utils::init_tracing("proxy", tracing_utils::ExportConfig::default()).await;
 
     let json_log_layer = if logfmt == LogFormat::Json {
         Some(JsonLoggingLayer::new(