tracing

rust/src/telemetry.rs

@@ -1,12 +1,15 @@
 use std::fmt;
+use std::future::Future;
 use std::io;
+use std::pin::Pin;
 use std::time::Duration;
 
 use axum::Router;
 use http::Request;
 use tower_http::{classify::ServerErrorsFailureClass, trace::TraceLayer};
-use tracing::{Level, Span};
+use tracing::Span;
 
+use tracing_log::LogTracer;
 use tracing_subscriber::{
     filter::{LevelFilter, Targets},
     fmt::{format::Format, Layer},
@@ -16,13 +19,31 @@ use tracing_subscriber::{
 };
 use uuid::Uuid;
 
-use opentelemetry::global;
+use opentelemetry::{global, runtime::Tokio};
 
-pub fn otel_init(f: impl FnOnce() -> ()) {
-    f()
+pub enum OpenTelemetryConfig {
+    Enabled,
+    Disabled,
 }
 
-pub fn init_tracing() {
+pub enum TokioConsoleConfig {
+    Enabled,
+    Disabled,
+}
+
+pub async fn init_tracing<Func, T>(
+    opentelemetry_config: OpenTelemetryConfig,
+    tokio_console_config: TokioConsoleConfig,
+    args: crate::cmd::Args,
+    f: Func,
+) -> T
+where
+    Func: FnOnce(crate::cmd::Args) -> Pin<Box<dyn Future<Output = T>>>,
+    T: std::process::Termination,
+{
+    let mut shutdown_functions: Vec<Box<dyn FnOnce() -> Result<(), Box<dyn std::error::Error>>>> =
+        vec![];
+
     // default is the Full format, there is no way to specify this, but it can be
     // overridden via builder methods
     let stdout_format = Format::default()
@@ -37,53 +58,81 @@ pub fn init_tracing() {
         .with_writer(io::stdout);
 
     let stdout_filter = Targets::new()
-        .with_default(LevelFilter::OFF)
+        .with_default(LevelFilter::WARN)
         .with_targets(vec![
-            (env!("CARGO_PKG_NAME"), Level::DEBUG),
-            // this is for axum requests
-            ("request", Level::DEBUG),
-            // required for tokio-console as by the docs
-            // ("tokio", Level::TRACE),
-            // ("runtime", Level::TRACE),
+            (env!("CARGO_PKG_NAME"), LevelFilter::DEBUG),
+            ("request", LevelFilter::DEBUG),
+            ("runtime", LevelFilter::OFF),
+            ("sqlx", LevelFilter::TRACE),
         ]);
 
     let stdout_layer = stdout_layer.with_filter(stdout_filter);
 
-    let console_layer = console_subscriber::Builder::default().spawn();
+    let console_layer = match tokio_console_config {
+        TokioConsoleConfig::Enabled => Some(console_subscriber::Builder::default().spawn()),
+        TokioConsoleConfig::Disabled => None,
+    };
 
-    global::set_text_map_propagator(opentelemetry_jaeger::Propagator::new());
-    // Sets up the machinery needed to export data to Jaeger
-    // There are other OTel crates that provide pipelines for the vendors
-    // mentioned earlier.
-    let tracer = opentelemetry_jaeger::new_agent_pipeline()
-        .with_service_name(env!("CARGO_PKG_NAME"))
-        .install_simple()
-        .unwrap();
+    let opentelemetry_layer = match opentelemetry_config {
+        OpenTelemetryConfig::Enabled => {
+            global::set_text_map_propagator(opentelemetry_jaeger::Propagator::new());
+            // Sets up the machinery needed to export data to Jaeger
+            // There are other OTel crates that provide pipelines for the vendors
+            // mentioned earlier.
+            let tracer = opentelemetry_jaeger::new_agent_pipeline()
+                .with_service_name(env!("CARGO_PKG_NAME"))
+                .with_max_packet_size(20_000)
+                .with_auto_split_batch(true)
+                .install_batch(Tokio)
+                .unwrap();
 
-    let opentelemetry_filter = Targets::new()
-        .with_default(LevelFilter::OFF)
-        .with_targets(vec![
-            (env!("CARGO_PKG_NAME"), Level::DEBUG),
-            // this is for axum requests
-            ("request", Level::DEBUG),
-            // required for tokio-console as by the docs
-            // ("tokio", Level::TRACE),
-            // ("runtime", Level::TRACE),
-        ]);
+            let opentelemetry_filter = {
+                Targets::new()
+                    .with_default(LevelFilter::DEBUG)
+                    .with_targets(vec![
+                        (env!("CARGO_PKG_NAME"), LevelFilter::DEBUG),
+                        ("request", LevelFilter::DEBUG),
+                        ("runtime", LevelFilter::OFF),
+                        ("sqlx", LevelFilter::DEBUG),
+                    ])
+            };
 
-    let opentelemetry = tracing_opentelemetry::layer()
-        .with_tracer(tracer)
-        .with_filter(opentelemetry_filter);
+            let opentelemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer);
+            // .with_filter(opentelemetry_filter);
+
+            shutdown_functions.push(Box::new(|| {
+                println!("shutting down otel");
+                global::shutdown_tracer_provider();
+                Ok(())
+            }));
+
+            println!("set up otel");
+
+            Some(opentelemetry_layer)
+        }
+        OpenTelemetryConfig::Disabled => None,
+    };
 
     let registry = Registry::default()
         .with(console_layer)
-        .with(opentelemetry)
+        .with(opentelemetry_layer)
         // just an example, you can actuall pass Options here for layers that might be
         // set/unset at runtime
         .with(Some(stdout_layer))
         .with(None::<Layer<_>>);
 
     tracing::subscriber::set_global_default(registry).unwrap();
+
+    tracing::debug!("tracing setup finished");
+
+    tracing_log::log_tracer::Builder::new().init().unwrap();
+
+    let result = f(args).await;
+
+    for shutdown_func in shutdown_functions {
+        shutdown_func().unwrap();
+    }
+    result
 }
 
 struct Latency(Duration);