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, AppContext, 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(ctx: &AppContext) -> Result<(), StdError> {
32        if ctx.has_component::<Self>() {
33            if !ctx.has_daemon::<TracingDaemon>() {
34                ctx.add_daemon(TracingDaemon);
35            }
36            return Ok(());
37        }
38        let config = match ctx
39            .get_component_ref::<Config>()
40            .unwrap()
41            .get::<Option<TracingConfig>>("tracing")?
42        {
43            Some(v) => v,
44            None => return Ok(()),
45        };
46        let mut directives = Vec::new();
47        for directive in config.directives {
48            directives.push(directive.parse().map_err(Box::new)?);
49        }
50        // Setup dynamic config level filter.
51        let (env_filter, reload_handle) =
52            reload::Layer::new(new_env_filter(&directives, config.level));
53        // Setup OpenTelemetry tracer.
54        let tracer_provider = {
55            if let Some(otlp_exporter) = config.otlp_exporter {
56                let exporter_builder = opentelemetry_otlp::SpanExporter::builder()
57                    .with_tonic()
58                    .with_endpoint(
59                        otlp_exporter
60                            .endpoint
61                            .unwrap_or(DEFAULT_OTLP_EXPORTER_ENDPOINT.into()),
62                    )
63                    .with_timeout(
64                        otlp_exporter
65                            .timeout
66                            .unwrap_or(DEFAULT_OTLP_EXPORTER_TIMEOUT),
67                    );
68                let exporter = CustomSpanExporter::new(exporter_builder.build().unwrap());
69                TracerProvider::builder()
70                    .with_resource(Resource::new(vec![KeyValue::new(
71                        "service.name",
72                        otlp_exporter.service_name.unwrap_or("unknown".into()),
73                    )]))
74                    .with_batch_exporter(exporter, runtime::Tokio)
75                    .build()
76            } else {
77                TracerProvider::builder().build()
78            }
79        };
80        // Setup tracing registry.
81        tracing_subscriber::registry()
82            .with(env_filter)
83            .with(tracing_subscriber::fmt::Layer::default())
84            .with(tracing_opentelemetry::layer().with_tracer(tracer_provider.tracer("")))
85            .init();
86        // Add app components.
87        ctx.add_component(Self {
88            default_level: config.level,
89            directives,
90            reload_handle,
91            tracer_provider,
92        });
93        if !ctx.has_daemon::<TracingDaemon>() {
94            ctx.add_daemon(TracingDaemon);
95        }
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 TracingOtlpExporterConfig {
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<TracingOtlpExporterConfig>,
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            #[allow(clippy::needless_bool)]
217            span.attributes.retain(|v| {
218                if v.key == OTEL_NAME_KEY {
219                    otel_name = Some(v.value.clone());
220                    false
221                } else if v.key == OTEL_KIND_KEY {
222                    otel_kind = Some(v.value.clone());
223                    false
224                } else if v.key == TRACE_ID_KEY {
225                    false
226                } else {
227                    true
228                }
229            });
230            if let Some(v) = otel_name {
231                span.name = v.to_string().into();
232            }
233            if let Some(v) = otel_kind {
234                match v.as_str().as_ref() {
235                    "server" => span.span_kind = SpanKind::Server,
236                    "client" => span.span_kind = SpanKind::Client,
237                    "consumer" => span.span_kind = SpanKind::Consumer,
238                    "producer" => span.span_kind = SpanKind::Producer,
239                    _ => {}
240                }
241            }
242        }
243        self.inner.export(batch)
244    }
245
246    fn shutdown(&mut self) {
247        self.inner.shutdown();
248    }
249
250    fn force_flush(&mut self) -> Pin<Box<dyn Future<Output = ExportResult> + Send + 'static>> {
251        self.inner.force_flush()
252    }
253
254    fn set_resource(&mut self, resource: &opentelemetry_sdk::Resource) {
255        self.inner.set_resource(resource);
256    }
257}
258
259fn serialize_level<S>(v: &tracing::Level, serializer: S) -> Result<S::Ok, S::Error>
260where
261    S: Serializer,
262{
263    serializer.serialize_str(v.as_str())
264}
265
266fn deserialize_level<'de, D>(deserializer: D) -> Result<tracing::Level, D::Error>
267where
268    D: Deserializer<'de>,
269{
270    use serde::de::Error;
271    String::deserialize(deserializer)
272        .and_then(|v| tracing::Level::from_str(&v).map_err(|v| Error::custom(format!("{v}"))))
273}
274
275fn default_level() -> tracing::Level {
276    tracing::Level::DEBUG
277}
diode_base/tracing.rs

diode_base/
tracing.rs
