From 7a154f74166c1309ca5752149e02bbe44cd91431 Mon Sep 17 00:00:00 2001
From: Lambda <csgl@xiretza.xyz>
Date: Thu, 23 May 2024 18:53:54 +0000
Subject: [PATCH] observability: allow combining multiple tracing layers

Previously, jaeger, tracing_flame and stderr output were mutually
exclusive. There's no good reason for this.
---
 src/observability.rs | 84 +++++++++++++++++++++++++-------------------
 1 file changed, 48 insertions(+), 36 deletions(-)

diff --git a/src/observability.rs b/src/observability.rs
index 9b7a9c89..0459a994 100644
--- a/src/observability.rs
+++ b/src/observability.rs
@@ -6,7 +6,7 @@ use std::{fs::File, io::BufWriter};
 use opentelemetry::KeyValue;
 use opentelemetry_sdk::Resource;
 use tracing_flame::{FlameLayer, FlushGuard};
-use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};
+use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Layer, Registry};
 
 use crate::{config::Config, error, utils::error::Result};
 
@@ -25,47 +25,59 @@ impl Drop for Guard {
 
 /// Initialize observability
 pub(crate) fn init(config: &Config) -> Result<Guard, error::Observability> {
-    let mut flame_guard = None;
-    if config.allow_jaeger {
-        opentelemetry::global::set_text_map_propagator(
-            opentelemetry_jaeger_propagator::Propagator::new(),
-        );
-        let tracer = opentelemetry_otlp::new_pipeline()
-            .tracing()
-            .with_trace_config(
-                opentelemetry_sdk::trace::config().with_resource(
-                    Resource::new(vec![KeyValue::new(
-                        "service.name",
-                        env!("CARGO_PKG_NAME"),
-                    )]),
-                ),
+    let config_filter_layer = || EnvFilter::try_new(&config.log);
+
+    let jaeger_layer = config
+        .allow_jaeger
+        .then(|| {
+            opentelemetry::global::set_text_map_propagator(
+                opentelemetry_jaeger_propagator::Propagator::new(),
+            );
+            let tracer = opentelemetry_otlp::new_pipeline()
+                .tracing()
+                .with_trace_config(
+                    opentelemetry_sdk::trace::config().with_resource(
+                        Resource::new(vec![KeyValue::new(
+                            "service.name",
+                            env!("CARGO_PKG_NAME"),
+                        )]),
+                    ),
+                )
+                .with_exporter(opentelemetry_otlp::new_exporter().tonic())
+                .install_batch(opentelemetry_sdk::runtime::Tokio)?;
+            let telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
+
+            Ok::<_, error::Observability>(
+                telemetry.with_filter(config_filter_layer()?),
             )
-            .with_exporter(opentelemetry_otlp::new_exporter().tonic())
-            .install_batch(opentelemetry_sdk::runtime::Tokio)?;
-        let telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
+        })
+        .transpose()?;
 
-        let filter_layer = EnvFilter::try_new(&config.log)?;
+    let (flame_layer, flame_guard) = config
+        .tracing_flame
+        .then(|| {
+            let (flame_layer, guard) =
+                FlameLayer::with_file("./tracing.folded")?;
+            let flame_layer = flame_layer.with_empty_samples(false);
 
-        let subscriber = Registry::default().with(filter_layer).with(telemetry);
-        tracing::subscriber::set_global_default(subscriber)?;
-    } else if config.tracing_flame {
-        let registry = Registry::default();
-        let (flame_layer, guard) = FlameLayer::with_file("./tracing.folded")?;
-        flame_guard = Some(guard);
-        let flame_layer = flame_layer.with_empty_samples(false);
+            let filter_layer = EnvFilter::new("trace,h2=off");
 
-        let filter_layer = EnvFilter::new("trace,h2=off");
+            Ok::<_, error::Observability>((
+                flame_layer.with_filter(filter_layer),
+                guard,
+            ))
+        })
+        .transpose()?
+        .unzip();
 
-        let subscriber = registry.with(filter_layer).with(flame_layer);
-        tracing::subscriber::set_global_default(subscriber)?;
-    } else {
-        let registry = Registry::default();
-        let fmt_layer = tracing_subscriber::fmt::Layer::new();
-        let filter_layer = EnvFilter::try_new(&config.log)?;
+    let fmt_layer = tracing_subscriber::fmt::Layer::new()
+        .with_filter(config_filter_layer()?);
 
-        let subscriber = registry.with(filter_layer).with(fmt_layer);
-        tracing::subscriber::set_global_default(subscriber)?;
-    }
+    let subscriber = Registry::default()
+        .with(jaeger_layer)
+        .with(flame_layer)
+        .with(fmt_layer);
+    tracing::subscriber::set_global_default(subscriber)?;
 
     Ok(Guard {
         flame_guard,