diode_base/
tracing.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::str::FromStr as _;
4use std::sync::Arc;
5use std::time::Duration;
6
7use diode::{App, AppBuilder, StdError};
8use duration_str::deserialize_option_duration;
9use opentelemetry::trace::{SpanKind, TracerProvider as _};
10use opentelemetry::{Key, KeyValue};
11use opentelemetry_otlp::WithExportConfig;
12use opentelemetry_sdk::export::trace::{ExportResult, SpanData, SpanExporter};
13use opentelemetry_sdk::trace::TracerProvider;
14use opentelemetry_sdk::{Resource, runtime};
15use serde::{Deserialize, Deserializer, Serialize, Serializer};
16use tracing_subscriber::filter::{Directive, EnvFilter};
17use tracing_subscriber::layer::SubscriberExt as _;
18use tracing_subscriber::util::SubscriberInitExt;
19use tracing_subscriber::{Registry, reload};
20
21use crate::{AddDaemonExt, CancellationToken, Config, ConfigSection, Daemon, DynamicConfig};
22
23pub struct Tracing {
24    default_level: tracing::Level,
25    directives: Vec<Directive>,
26    reload_handle: reload::Handle<EnvFilter, Registry>,
27    tracer_provider: TracerProvider,
28}
29
30impl Tracing {
31    pub fn build(app: &mut AppBuilder) -> Result<(), StdError> {
32        if app.has_component::<Tracing>() {
33            // Ensure that the tracing daemon is added.
34            if !app.has_daemon::<TracingDaemon>() {
35                app.add_daemon(TracingDaemon);
36            }
37            return Ok(());
38        }
39        let config = match app
40            .get_component_ref::<Config>()
41            .unwrap()
42            .get::<Option<TracingConfig>>("tracing")
43            .unwrap()
44        {
45            Some(v) => v,
46            None => return Ok(()),
47        };
48        let mut directives = Vec::new();
49        for directive in config.directives {
50            directives.push(directive.parse().map_err(Box::new)?);
51        }
52        // Setup dynamic config level filter.
53        let (env_filter, reload_handle) =
54            reload::Layer::new(new_env_filter(&directives, config.level));
55        // Setup OpenTelemetry tracer.
56        let tracer_provider = {
57            if let Some(otlp_exporter) = config.otlp_exporter {
58                let exporter_builder = opentelemetry_otlp::SpanExporter::builder()
59                    .with_tonic()
60                    .with_endpoint(
61                        otlp_exporter
62                            .endpoint
63                            .unwrap_or(DEFAULT_OTLP_EXPORTER_ENDPOINT.into()),
64                    )
65                    .with_timeout(
66                        otlp_exporter
67                            .timeout
68                            .unwrap_or(DEFAULT_OTLP_EXPORTER_TIMEOUT),
69                    );
70                let exporter = CustomSpanExporter::new(exporter_builder.build().unwrap());
71                TracerProvider::builder()
72                    .with_resource(Resource::new(vec![KeyValue::new(
73                        "service.name",
74                        otlp_exporter.service_name.unwrap_or("unknown".into()),
75                    )]))
76                    .with_batch_exporter(exporter, runtime::Tokio)
77                    .build()
78            } else {
79                TracerProvider::builder().build()
80            }
81        };
82        // Setup tracing registry.
83        tracing_subscriber::registry()
84            .with(env_filter)
85            .with(tracing_subscriber::fmt::Layer::default())
86            .with(tracing_opentelemetry::layer().with_tracer(tracer_provider.tracer("")))
87            .init();
88        // Add app components.
89        app.add_component(Self {
90            default_level: config.level,
91            directives,
92            reload_handle,
93            tracer_provider,
94        });
95        app.add_daemon(TracingDaemon);
96        Ok(())
97    }
98}
99
100impl Drop for Tracing {
101    fn drop(&mut self) {
102        if let Err(err) = self.tracer_provider.shutdown() {
103            tracing::error!("Cannot shutdown tracer provider: {err}");
104        }
105    }
106}
107
108struct TracingDaemon;
109
110const TRACING_LEVEL_CONFIG_KEY: &str = "tracing_level";
111
112impl Daemon for TracingDaemon {
113    async fn run(&self, app: &App, shutdown: CancellationToken) -> Result<(), StdError> {
114        let tracing = app.get_component_ref::<Tracing>().unwrap();
115        let default_level = tracing.default_level;
116        let directives = tracing.directives.clone();
117        let reload_handle = tracing.reload_handle.clone();
118        if let Some(dynamic_config) = app.get_component::<Arc<DynamicConfig>>() {
119            dynamic_config.subscribe(TRACING_LEVEL_CONFIG_KEY, move |level: Option<String>| {
120                let level = match level {
121                    Some(v) => match tracing::Level::from_str(&v) {
122                        Ok(v) => v,
123                        Err(err) => {
124                            tracing::error!("Cannot parse tracing level: {}", err);
125                            return;
126                        }
127                    },
128                    None => default_level,
129                };
130                reload_handle
131                    .reload(new_env_filter(&directives, level))
132                    .unwrap();
133            });
134        }
135        shutdown.cancelled_owned().await;
136        Ok(())
137    }
138}
139
140impl Default for TracingConfig {
141    fn default() -> Self {
142        Self {
143            level: default_level(),
144            directives: Default::default(),
145            otlp_exporter: None,
146        }
147    }
148}
149
150#[derive(Serialize, Deserialize)]
151pub struct OtlpExporterConfig {
152    #[serde(default)]
153    pub service_name: Option<String>,
154    #[serde(default)]
155    pub endpoint: Option<String>,
156    #[serde(default, deserialize_with = "deserialize_option_duration")]
157    pub timeout: Option<Duration>,
158}
159
160#[derive(Serialize, Deserialize)]
161pub struct TracingConfig {
162    #[serde(
163        serialize_with = "serialize_level",
164        deserialize_with = "deserialize_level",
165        default = "default_level"
166    )]
167    pub level: tracing::Level,
168    #[serde(default)]
169    pub directives: Vec<String>,
170    #[serde(default)]
171    pub otlp_exporter: Option<OtlpExporterConfig>,
172}
173
174impl ConfigSection for TracingConfig {
175    fn key() -> &'static str {
176        "tracing"
177    }
178}
179
180fn new_env_filter(directives: &Vec<Directive>, level: tracing::Level) -> EnvFilter {
181    let mut filter = EnvFilter::default();
182    for directive in directives {
183        filter = filter.add_directive(directive.clone());
184    }
185    filter.add_directive(level.into())
186}
187
188const DEFAULT_OTLP_EXPORTER_ENDPOINT: &str = "https://localhost:4317/v1/traces";
189const DEFAULT_OTLP_EXPORTER_TIMEOUT: Duration = Duration::from_secs(10);
190
191#[derive(Debug)]
192struct CustomSpanExporter<T> {
193    inner: T,
194}
195
196impl<T> CustomSpanExporter<T> {
197    pub fn new(inner: T) -> Self {
198        Self { inner }
199    }
200}
201
202impl<T> SpanExporter for CustomSpanExporter<T>
203where
204    T: SpanExporter,
205{
206    fn export(
207        &mut self,
208        mut batch: Vec<SpanData>,
209    ) -> Pin<Box<dyn Future<Output = ExportResult> + Send + 'static>> {
210        const OTEL_NAME_KEY: Key = Key::from_static_str("otel.name");
211        const OTEL_KIND_KEY: Key = Key::from_static_str("otel.kind");
212        const TRACE_ID_KEY: Key = Key::from_static_str("trace_id");
213        for span in batch.iter_mut() {
214            let mut otel_name = None;
215            let mut otel_kind = None;
216            span.attributes.retain(|v| {
217                if v.key == OTEL_NAME_KEY {
218                    otel_name = Some(v.value.clone());
219                    false
220                } else if v.key == OTEL_KIND_KEY {
221                    otel_kind = Some(v.value.clone());
222                    false
223                } else if v.key == TRACE_ID_KEY {
224                    false
225                } else {
226                    true
227                }
228            });
229            if let Some(v) = otel_name {
230                span.name = v.to_string().into();
231            }
232            if let Some(v) = otel_kind {
233                match v.as_str().as_ref() {
234                    "server" => span.span_kind = SpanKind::Server,
235                    "client" => span.span_kind = SpanKind::Client,
236                    "consumer" => span.span_kind = SpanKind::Consumer,
237                    "producer" => span.span_kind = SpanKind::Producer,
238                    _ => {}
239                }
240            }
241        }
242        self.inner.export(batch)
243    }
244
245    fn shutdown(&mut self) {
246        self.inner.shutdown();
247    }
248
249    fn force_flush(&mut self) -> Pin<Box<dyn Future<Output = ExportResult> + Send + 'static>> {
250        self.inner.force_flush()
251    }
252
253    fn set_resource(&mut self, resource: &opentelemetry_sdk::Resource) {
254        self.inner.set_resource(resource);
255    }
256}
257
258fn serialize_level<S>(v: &tracing::Level, serializer: S) -> Result<S::Ok, S::Error>
259where
260    S: Serializer,
261{
262    serializer.serialize_str(v.as_str())
263}
264
265fn deserialize_level<'de, D>(deserializer: D) -> Result<tracing::Level, D::Error>
266where
267    D: Deserializer<'de>,
268{
269    use serde::de::Error;
270    String::deserialize(deserializer)
271        .and_then(|v| tracing::Level::from_str(&v).map_err(|v| Error::custom(format!("{v}"))))
272}
273
274fn default_level() -> tracing::Level {
275    tracing::Level::DEBUG
276}
diode_base/tracing.rs

diode_base/
tracing.rs
