krabka_telemetry/
subscriber.rs

1//! Install of the global `tracing` subscriber and the guard that flushes it.
2//!
3//! The module wires the stdout JSON layer, the optional OTLP span layer, and
4//! the optional OTLP log bridge into one registry, and it owns the shutdown
5//! path for the providers that those layers export through.
6
7use {
8    crate::{config::OtlpConfig, error::TelemetryError, log_levels::LogLevelController},
9    opentelemetry_sdk::{logs::SdkLoggerProvider, trace::SdkTracerProvider},
10    std::{env, io},
11    tokio::task::JoinHandle,
12    tracing_subscriber::{
13        EnvFilter, Layer, Registry, layer::SubscriberExt as _, util::SubscriberInitExt as _,
14    },
15};
16
17/// The stdout layer, boxed as `krabka_logfmt::layer` returns it.
18type StdoutLayer = Box<dyn Layer<Registry> + Send + Sync>;
19
20/// Owns the OTLP `SdkTracerProvider`, so shutdown flushes the spans.
21///
22/// Dropping the guard stops the heartbeat and shuts down both exporters,
23/// including when another provider clone remains alive. Call
24/// [`TelemetryGuard::shutdown`] before exit to deliver the final batch.
25#[must_use = "hold the guard for the process lifetime and call shutdown() before exit"]
26pub struct TelemetryGuard {
27    provider: Option<SdkTracerProvider>,
28    logger_provider: Option<SdkLoggerProvider>,
29    heartbeat_task: Option<JoinHandle<()>>,
30    log_levels: LogLevelController,
31}
32
33impl TelemetryGuard {
34    /// The handle that reads and retargets the stdout layer's filter while
35    /// the process runs.
36    ///
37    /// A Kafka broker exposes this through the `BROKER_LOGGER` config
38    /// resource, so `kafka-configs --entity-type broker-loggers` can raise a
39    /// level without a restart. Clone it out before the guard is parked for
40    /// the process lifetime.
41    #[must_use]
42    pub fn log_levels(&self) -> LogLevelController { self.log_levels.clone() }
43
44    /// Flush and shut down the OTLP exporters for traces and logs.
45    ///
46    /// This function does nothing when OTLP is disabled.
47    pub fn shutdown(mut self) { self.shutdown_providers(); }
48
49    fn shutdown_providers(&mut self) {
50        if let Some(task) = self.heartbeat_task.take() {
51            task.abort();
52        }
53        if let Some(provider) = self.provider.take()
54            && let Err(e) = provider.shutdown()
55        {
56            tracing::warn!(error = %e, "OTLP tracer provider shutdown error");
57        }
58        if let Some(logger_provider) = self.logger_provider.take()
59            && let Err(e) = logger_provider.shutdown()
60        {
61            tracing::warn!(error = %e, "OTLP logger provider shutdown error");
62        }
63    }
64}
65
66impl Drop for TelemetryGuard {
67    fn drop(&mut self) { self.shutdown_providers(); }
68}
69
70/// Install the global `tracing` subscriber.
71///
72/// The subscriber has a stdout JSON `fmt` layer. When `otlp` is `Some`, the
73/// subscriber also has a batch OTLP export layer.
74///
75/// - `fmt_default_filter`: the `fmt` layer's filter when `RUST_LOG` is unset.
76/// - `otel_default_filter`: the OTLP layer's filter when `KRABKA_OTLP_FILTER`
77///   is unset.
78/// - `tracer_name`: the name that this function passes to
79///   `TracerProvider::tracer(...)`.
80///
81/// You must call this function exactly one time, and from inside the tokio
82/// runtime. The gRPC exporter captures the current runtime handle.
83///
84/// # Errors
85///
86/// Returns `TelemetryError::Exporter` when an OTLP exporter cannot be built.
87/// On `wasm32-wasip1`, returns `TelemetryError::Unsupported` for any OTLP
88/// configuration and installs nothing.
89pub fn init(
90    otlp: Option<OtlpConfig>,
91    fmt_default_filter: &str,
92    otel_default_filter: &str,
93    tracer_name: &str,
94) -> Result<TelemetryGuard, TelemetryError> {
95    // The stdout layer's filter is the one an operator retargets at runtime,
96    // so it is a `LogLevelFilter` and not a bare `EnvFilter`. It starts from
97    // `RUST_LOG` when that is set and from `fmt_default_filter` otherwise.
98    let fmt_spec = env::var("RUST_LOG")
99        .ok()
100        .filter(|spec| EnvFilter::try_new(spec).is_ok())
101        .unwrap_or_else(|| fmt_default_filter.to_owned());
102    let (log_levels, fmt_filter) = LogLevelController::new(&fmt_spec);
103    // Structured Cloud Logging-friendly JSON to stdout (see `krabka_logfmt`),
104    // so GKE / Cloud Logging ingests fields rather than ANSI-coloured text.
105    let fmt_layer = krabka_logfmt::layer(fmt_filter, io::stdout);
106
107    let Some(cfg) = otlp else {
108        tracing_subscriber::registry().with(fmt_layer).init();
109        return Ok(TelemetryGuard {
110            provider: None,
111            logger_provider: None,
112            heartbeat_task: None,
113            log_levels,
114        });
115    };
116    otlp::install(
117        &cfg,
118        fmt_layer,
119        log_levels,
120        otel_default_filter,
121        tracer_name,
122    )
123}
124
125/// The OTLP span and log layers, installed beside the stdout layer.
126#[cfg(not(target_family = "wasm"))]
127mod otlp {
128    use {
129        super::{StdoutLayer, TelemetryGuard},
130        crate::{
131            config::OtlpConfig, error::TelemetryError, heartbeat::spawn_heartbeat_task,
132            log_levels::LogLevelController,
133        },
134        krabka_units::prelude::TimeExt,
135        opentelemetry::{global, trace::TracerProvider as _},
136        opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge,
137        opentelemetry_sdk::{
138            logs::SdkLoggerProvider, propagation::TraceContextPropagator, trace::SdkTracerProvider,
139        },
140        std::env,
141        tracing_subscriber::{
142            EnvFilter, Layer as _, layer::SubscriberExt as _, util::SubscriberInitExt as _,
143        },
144    };
145
146    /// Per-layer filter for the OTLP layer.
147    ///
148    /// Operators who want more control can override the filter with
149    /// `KRABKA_OTLP_FILTER`. Without that variable, the filter falls back to
150    /// `default`.
151    fn otel_filter(default: &str, get: impl Fn(&str) -> Option<String>) -> EnvFilter {
152        get("KRABKA_OTLP_FILTER")
153            .and_then(|s| EnvFilter::try_new(s).ok())
154            .unwrap_or_else(|| EnvFilter::new(default))
155    }
156
157    pub(super) fn install(
158        cfg: &OtlpConfig,
159        fmt_layer: StdoutLayer,
160        log_levels: LogLevelController,
161        otel_default_filter: &str,
162        tracer_name: &str,
163    ) -> Result<TelemetryGuard, TelemetryError> {
164        let exporter = cfg.build_exporter()?;
165        let provider = SdkTracerProvider::builder()
166            .with_batch_exporter(exporter)
167            .with_resource(cfg.resource())
168            .with_sampler(cfg.sampler())
169            .build();
170        let tracer = provider.tracer(tracer_name.to_owned());
171
172        global::set_text_map_propagator(TraceContextPropagator::new());
173        global::set_tracer_provider(provider.clone());
174
175        let otel_layer = tracing_opentelemetry::layer()
176            .with_tracer(tracer)
177            .with_location(false)
178            .with_filter(otel_filter(otel_default_filter, |k| env::var(k).ok()));
179
180        // OTLP LOGS: a `tracing` → OTLP-log-record bridge so services ship their
181        // logs over OTLP (alloy:4317 → logs-distributor) — Docker-stdout tailing
182        // doesn't stream on Docker Desktop. Quiet the OTLP/transport stack on THIS
183        // layer so the exporter's own events can't feed back into more log records.
184        let log_exporter = cfg.build_log_exporter()?;
185        let logger_provider = SdkLoggerProvider::builder()
186            .with_batch_exporter(log_exporter)
187            .with_resource(cfg.resource())
188            .build();
189        let logs_default = format!(
190            "{otel_default_filter},opentelemetry=warn,opentelemetry_sdk=warn,opentelemetry_otlp=warn,hyper_util=warn,tonic=warn,h2=warn,tower=warn"
191        );
192        let otel_logs_layer = OpenTelemetryTracingBridge::new(&logger_provider)
193            .with_filter(otel_filter(&logs_default, |k| env::var(k).ok()));
194
195        tracing_subscriber::registry()
196            .with(fmt_layer)
197            .with(otel_layer)
198            .with(otel_logs_layer)
199            .init();
200
201        let heartbeat_task = cfg.heartbeat_interval.map(|interval| {
202            spawn_heartbeat_task(
203                provider.clone(),
204                cfg.service_name.clone(),
205                cfg.service_instance_id.clone(),
206                interval,
207            )
208        });
209
210        tracing::info!(
211            endpoint = %cfg.endpoint,
212            protocol = ?cfg.protocol,
213            sample_ratio = cfg.sample_ratio,
214            heartbeat_interval_secs = cfg.heartbeat_interval.map(TimeExt::secs_f64),
215            "OTLP distributed tracing + logs enabled"
216        );
217
218        Ok(TelemetryGuard {
219            provider: Some(provider),
220            logger_provider: Some(logger_provider),
221            heartbeat_task,
222            log_levels,
223        })
224    }
225
226    #[cfg(test)]
227    mod tests {
228        use {super::otel_filter, crate::test_support::env_from};
229
230        #[test]
231        fn otel_filter_uses_override_or_default() {
232            let filter = otel_filter("warn", env_from(&[]));
233            assert2::assert!(filter.to_string() == "warn");
234
235            let filter = otel_filter(
236                "warn",
237                env_from(&[("KRABKA_OTLP_FILTER", "krabka_telemetry=debug")]),
238            );
239            assert2::assert!(filter.to_string() == "krabka_telemetry=debug");
240
241            let filter = otel_filter("warn", env_from(&[("KRABKA_OTLP_FILTER", "[")]));
242            assert2::assert!(filter.to_string() == "warn");
243        }
244    }
245}
246
247// OTLP export needs a gRPC or HTTP client stack, which `wasm32-wasip1` does
248// not have.
249#[cfg(target_family = "wasm")]
250mod otlp {
251    use {
252        super::{StdoutLayer, TelemetryGuard},
253        crate::{config::OtlpConfig, error::TelemetryError, log_levels::LogLevelController},
254    };
255
256    pub(super) fn install(
257        _cfg: &OtlpConfig,
258        _fmt_layer: StdoutLayer,
259        _log_levels: LogLevelController,
260        _otel_default_filter: &str,
261        _tracer_name: &str,
262    ) -> Result<TelemetryGuard, TelemetryError> {
263        Err(TelemetryError::Unsupported {
264            what: "OTLP export",
265        })
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use {
272        super::*,
273        opentelemetry_sdk::{
274            error::OTelSdkResult,
275            trace::{SpanData, SpanExporter},
276        },
277        std::{
278            future::{self, Future},
279            sync::{
280                Arc,
281                atomic::{AtomicUsize, Ordering::SeqCst},
282            },
283            time::Duration,
284        },
285        tokio::sync::oneshot,
286    };
287
288    #[derive(Clone, Debug)]
289    struct ShutdownCountingSpanExporter {
290        calls: Arc<AtomicUsize>,
291    }
292
293    impl SpanExporter for ShutdownCountingSpanExporter {
294        fn export(&self, _batch: Vec<SpanData>) -> impl Future<Output = OTelSdkResult> + Send {
295            future::ready(Ok(()))
296        }
297
298        fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult {
299            self.calls.fetch_add(1, SeqCst);
300            Ok(())
301        }
302    }
303
304    /// A tracer provider whose exporter counts its shutdowns into the
305    /// returned counter.
306    fn counting_provider() -> (Arc<AtomicUsize>, SdkTracerProvider) {
307        let calls = Arc::new(AtomicUsize::new(0));
308        let provider = SdkTracerProvider::builder()
309            .with_simple_exporter(ShutdownCountingSpanExporter {
310                calls: Arc::clone(&calls),
311            })
312            .build();
313        (calls, provider)
314    }
315
316    #[test]
317    fn telemetry_guard_shutdown_flushes_held_provider_clone() {
318        let (calls, provider) = counting_provider();
319        let held_clone = provider.clone();
320
321        TelemetryGuard {
322            provider: Some(provider),
323            logger_provider: None,
324            heartbeat_task: None,
325            log_levels: LogLevelController::new("info").0,
326        }
327        .shutdown();
328
329        assert2::assert!(calls.load(SeqCst) == 1);
330        drop(held_clone);
331    }
332
333    #[tokio::test]
334    async fn telemetry_guard_drop_shuts_down_providers_and_stops_the_heartbeat() {
335        let (calls, provider) = counting_provider();
336        let held_clone = provider.clone();
337        let (started, ready) = oneshot::channel();
338        let (finished, stopped) = oneshot::channel::<()>();
339        let heartbeat_task = tokio::spawn(async move {
340            let _finished = finished;
341            started.send(()).unwrap();
342            future::pending::<()>().await;
343        });
344        ready.await.unwrap();
345        drop(TelemetryGuard {
346            provider: Some(provider),
347            logger_provider: None,
348            heartbeat_task: Some(heartbeat_task),
349            log_levels: LogLevelController::new("info").0,
350        });
351        assert2::assert!(calls.load(SeqCst) == 1);
352        assert2::assert!(stopped.await.is_err());
353        drop(held_clone);
354        assert2::assert!(calls.load(SeqCst) == 1);
355    }
356}
krabka_telemetry/subscriber.rs

krabka_telemetry/
subscriber.rs
