Skip to main content

nanocodex_observability/
lib.rs

1#![doc = include_str!("../README.md")]
2#![deny(missing_docs, rustdoc::broken_intra_doc_links)]
3
4extern crate self as nanocodex_observability;
5
6use std::{fs::OpenOptions, io, path::PathBuf};
7
8use opentelemetry::trace::TracerProvider as _;
9use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig, WithHttpConfig};
10use opentelemetry_sdk::{
11    Resource, runtime,
12    trace::{
13        SdkTracerProvider,
14        span_processor_with_async_runtime::BatchSpanProcessor as TokioBatchSpanProcessor,
15    },
16};
17use opentelemetry_semantic_conventions::{SCHEMA_URL, attribute::SERVICE_VERSION};
18use tracing_appender::non_blocking::WorkerGuard;
19use tracing_subscriber::{
20    EnvFilter, Layer, fmt::format::FmtSpan, layer::SubscriberExt, util::SubscriberInitExt,
21};
22
23const DEPLOYMENT_ENVIRONMENT_NAME: &str = "deployment.environment.name";
24
25/// Human-readable or structured local tracing output.
26#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
27pub enum LogFormat {
28    /// Multiline human-readable records.
29    Pretty,
30    /// Dense human-readable records.
31    #[default]
32    Compact,
33    /// Structured JSON records.
34    Json,
35}
36
37/// Destination for the local formatting layer.
38#[derive(Clone, Debug, Default, Eq, PartialEq)]
39pub enum LogOutput {
40    /// Write local records to standard error.
41    #[default]
42    Stderr,
43    /// Append local records to a file, creating its parent directories.
44    File(PathBuf),
45}
46
47/// Builder for local formatting and optional OTLP/HTTP trace export.
48#[derive(Clone, Debug)]
49pub struct ObservabilityBuilder {
50    filter: String,
51    otel_filter: String,
52    format: LogFormat,
53    output: LogOutput,
54    service_name: String,
55    service_version: String,
56    environment: Option<String>,
57    otlp_endpoint: Option<String>,
58}
59
60/// Keeps asynchronous formatting and OTLP providers alive and flushes them on drop.
61pub struct ObservabilityGuard {
62    tracer_provider: Option<SdkTracerProvider>,
63    _writer: WorkerGuard,
64}
65
66/// Failure to configure, install, or flush observability.
67#[derive(Debug, thiserror::Error)]
68pub enum ObservabilityError {
69    /// A local or OpenTelemetry filter is invalid.
70    #[error("invalid tracing filter: {0}")]
71    Filter(#[from] tracing_subscriber::filter::ParseError),
72    /// The configured local output could not be opened.
73    #[error("failed to open tracing output: {0}")]
74    Output(#[from] io::Error),
75    /// The OTLP exporter could not be constructed.
76    #[error("failed to configure OTLP exporter: {0}")]
77    Otlp(#[from] opentelemetry_otlp::ExporterBuildError),
78    /// The OpenTelemetry SDK could not flush or shut down.
79    #[error("failed to flush or shut down the OpenTelemetry exporter: {0}")]
80    OTelSdk(#[from] opentelemetry_sdk::error::OTelSdkError),
81    /// Another process-global tracing subscriber is already installed.
82    #[error("a global tracing subscriber is already installed")]
83    Subscriber,
84}
85
86impl ObservabilityBuilder {
87    /// Starts a builder with the service identity attached to exported spans.
88    #[must_use]
89    pub fn new(service_name: impl Into<String>, service_version: impl Into<String>) -> Self {
90        let filter = "warn,nanocodex=info,nanocodex_oai_api=info,nanocodex_tools=info".to_owned();
91        Self {
92            otel_filter: filter.clone(),
93            filter,
94            format: LogFormat::Compact,
95            output: LogOutput::Stderr,
96            service_name: service_name.into(),
97            service_version: service_version.into(),
98            environment: None,
99            otlp_endpoint: None,
100        }
101    }
102
103    /// Sets the filter applied to local formatted records.
104    #[must_use]
105    pub fn filter(mut self, filter: impl Into<String>) -> Self {
106        self.filter = filter.into();
107        self
108    }
109
110    /// Sets the independent filter applied only to exported OpenTelemetry spans.
111    #[must_use]
112    pub fn otel_filter(mut self, filter: impl Into<String>) -> Self {
113        self.otel_filter = filter.into();
114        self
115    }
116
117    /// Selects the local record format.
118    #[must_use]
119    pub const fn format(mut self, format: LogFormat) -> Self {
120        self.format = format;
121        self
122    }
123
124    /// Selects the local record destination.
125    #[must_use]
126    pub fn output(mut self, output: LogOutput) -> Self {
127        self.output = output;
128        self
129    }
130
131    /// Sets the deployment environment attached to exported spans.
132    #[must_use]
133    pub fn environment(mut self, environment: impl Into<String>) -> Self {
134        self.environment = Some(environment.into());
135        self
136    }
137
138    /// Sets the OTLP/HTTP collector base endpoint.
139    ///
140    /// The standard `/v1/traces` signal path is appended unless it is already
141    /// present, matching `OTEL_EXPORTER_OTLP_ENDPOINT` semantics.
142    #[must_use]
143    pub fn otlp_endpoint(mut self, endpoint: impl Into<String>) -> Self {
144        let endpoint = endpoint.into();
145        self.otlp_endpoint = Some(trace_endpoint(&endpoint));
146        self
147    }
148
149    /// Installs the process-global subscriber.
150    ///
151    /// # Errors
152    ///
153    /// Returns an error for an invalid filter, unusable output file, invalid
154    /// exporter configuration, or an already-installed global subscriber.
155    pub fn install(self) -> Result<ObservabilityGuard, ObservabilityError> {
156        let (writer, writer_guard) = tracing_appender::non_blocking(self.writer()?);
157        let filter = EnvFilter::try_new(self.filter.as_str())?;
158        let otel_filter = EnvFilter::try_new(self.otel_filter.as_str())?;
159        let fmt_layer = match self.format {
160            LogFormat::Pretty => tracing_subscriber::fmt::layer()
161                .pretty()
162                .with_writer(writer)
163                .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
164                .with_filter(filter)
165                .boxed(),
166            LogFormat::Compact => tracing_subscriber::fmt::layer()
167                .compact()
168                .with_writer(writer)
169                .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
170                .with_filter(filter)
171                .boxed(),
172            LogFormat::Json => tracing_subscriber::fmt::layer()
173                .json()
174                .with_span_list(true)
175                .with_current_span(false)
176                .with_writer(writer)
177                .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
178                .with_filter(filter)
179                .boxed(),
180        };
181        let tracer_provider = self.tracer_provider()?;
182        let tracing_layer = tracer_provider.as_ref().map(|provider| {
183            tracing_opentelemetry::layer()
184                .with_tracer(provider.tracer(self.service_name))
185                .with_filter(otel_filter)
186        });
187        tracing_subscriber::registry()
188            .with(fmt_layer)
189            .with(tracing_layer)
190            .try_init()
191            .map_err(|_| ObservabilityError::Subscriber)?;
192        tracing::callsite::rebuild_interest_cache();
193        Ok(ObservabilityGuard {
194            tracer_provider,
195            _writer: writer_guard,
196        })
197    }
198
199    fn writer(&self) -> Result<Box<dyn io::Write + Send>, io::Error> {
200        match &self.output {
201            LogOutput::Stderr => Ok(Box::new(io::stderr())),
202            LogOutput::File(path) => {
203                if let Some(parent) = path.parent() {
204                    std::fs::create_dir_all(parent)?;
205                }
206                Ok(Box::new(
207                    OpenOptions::new().create(true).append(true).open(path)?,
208                ))
209            }
210        }
211    }
212
213    fn tracer_provider(&self) -> Result<Option<SdkTracerProvider>, ObservabilityError> {
214        let Some(endpoint) = self.otlp_endpoint.as_deref() else {
215            return Ok(None);
216        };
217        let resource = self.resource();
218        if current_tokio_runtime_is_multi_thread() {
219            let exporter = SpanExporter::builder()
220                .with_http()
221                .with_endpoint(endpoint)
222                .with_protocol(Protocol::HttpBinary)
223                .with_http_client(reqwest::Client::new())
224                .build()?;
225            let processor = TokioBatchSpanProcessor::builder(exporter, runtime::Tokio).build();
226            return Ok(Some(
227                SdkTracerProvider::builder()
228                    .with_resource(resource)
229                    .with_span_processor(processor)
230                    .build(),
231            ));
232        }
233
234        // The async processor's synchronous flush waits for work scheduled on
235        // Tokio. Keep the SDK's dedicated-thread processor when no runtime is
236        // active or when a current-thread runtime could deadlock that flush.
237        let exporter = SpanExporter::builder()
238            .with_http()
239            .with_endpoint(endpoint)
240            .with_protocol(Protocol::HttpBinary)
241            .with_http_client(reqwest::blocking::Client::new())
242            .build()?;
243        let processor = opentelemetry_sdk::trace::BatchSpanProcessor::builder(exporter).build();
244        Ok(Some(
245            SdkTracerProvider::builder()
246                .with_resource(resource)
247                .with_span_processor(processor)
248                .build(),
249        ))
250    }
251
252    fn resource(&self) -> Resource {
253        Resource::builder()
254            .with_service_name(self.service_name.clone())
255            .with_schema_url(
256                [
257                    opentelemetry::KeyValue::new(SERVICE_VERSION, self.service_version.clone()),
258                    opentelemetry::KeyValue::new(
259                        DEPLOYMENT_ENVIRONMENT_NAME,
260                        self.environment
261                            .clone()
262                            .unwrap_or_else(|| "development".to_owned()),
263                    ),
264                ],
265                SCHEMA_URL,
266            )
267            .build()
268    }
269}
270
271fn current_tokio_runtime_is_multi_thread() -> bool {
272    tokio::runtime::Handle::try_current()
273        .is_ok_and(|handle| handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread)
274}
275
276fn trace_endpoint(endpoint: &str) -> String {
277    let endpoint = endpoint.trim_end_matches('/');
278    if endpoint.ends_with("/v1/traces") {
279        endpoint.to_owned()
280    } else {
281        format!("{endpoint}/v1/traces")
282    }
283}
284
285impl ObservabilityGuard {
286    /// Flushes pending spans and shuts down the exporter. Later calls are no-ops.
287    ///
288    /// # Errors
289    ///
290    /// Returns an error when the exporter cannot flush or shut down cleanly.
291    pub fn shutdown(&mut self) -> Result<(), ObservabilityError> {
292        if let Some(provider) = self.tracer_provider.take() {
293            provider.force_flush()?;
294            provider.shutdown()?;
295        }
296        Ok(())
297    }
298}
299
300impl Drop for ObservabilityGuard {
301    fn drop(&mut self) {
302        drop(self.shutdown());
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use std::{
309        io::{Read, Write},
310        net::TcpListener,
311        sync::mpsc,
312        thread,
313        time::{Duration, Instant},
314    };
315
316    use super::*;
317
318    const OTLP_TEST_TIMEOUT: Duration = Duration::from_secs(30);
319
320    #[test]
321    fn resource_uses_configured_service_identity_and_semantic_schema() {
322        let resource = ObservabilityBuilder::new("nanocodex-test", "1.2.3")
323            .environment("test")
324            .resource();
325
326        assert_eq!(
327            resource.get(&opentelemetry::Key::new("service.name")),
328            Some(opentelemetry::Value::from("nanocodex-test"))
329        );
330        assert_eq!(
331            resource.get(&opentelemetry::Key::new(SERVICE_VERSION)),
332            Some(opentelemetry::Value::from("1.2.3"))
333        );
334        assert_eq!(
335            resource.get(&opentelemetry::Key::new(DEPLOYMENT_ENVIRONMENT_NAME)),
336            Some(opentelemetry::Value::from("test"))
337        );
338        assert_eq!(resource.schema_url(), Some(SCHEMA_URL));
339    }
340
341    #[test]
342    fn async_export_requires_a_multithreaded_tokio_runtime() {
343        assert!(!current_tokio_runtime_is_multi_thread());
344
345        let current_thread = tokio::runtime::Builder::new_current_thread()
346            .build()
347            .unwrap();
348        assert!(!current_thread.block_on(async { current_tokio_runtime_is_multi_thread() }));
349
350        let multi_thread = tokio::runtime::Builder::new_multi_thread()
351            .worker_threads(2)
352            .build()
353            .unwrap();
354        assert!(multi_thread.block_on(async { current_tokio_runtime_is_multi_thread() }));
355    }
356
357    #[test]
358    fn formatting_and_otlp_export_share_the_installed_span_stream() {
359        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
360        listener.set_nonblocking(true).unwrap();
361        let address = listener.local_addr().unwrap();
362        let (request_seen, request_received) = mpsc::channel();
363        let server = thread::spawn(move || {
364            let deadline = Instant::now() + OTLP_TEST_TIMEOUT;
365            loop {
366                assert!(Instant::now() < deadline, "OTLP exporter did not connect");
367                match listener.accept() {
368                    Ok((mut stream, _)) => {
369                        stream
370                            .set_read_timeout(Some(Duration::from_secs(2)))
371                            .unwrap();
372                        let mut request = Vec::with_capacity(4 * 1024);
373                        let mut chunk = [0_u8; 1024];
374                        while request.len() < 16 * 1024 {
375                            let read = stream.read(&mut chunk).unwrap_or_default();
376                            if read == 0 {
377                                break;
378                            }
379                            request.extend_from_slice(&chunk[..read]);
380                            if request.windows(4).any(|window| window == b"\r\n\r\n") {
381                                break;
382                            }
383                        }
384                        // The HTTP client may open and close an empty warm-up
385                        // connection before it sends the export request.
386                        if request.is_empty() {
387                            continue;
388                        }
389                        stream
390                            .write_all(
391                                b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\nconnection: close\r\n\r\n",
392                            )
393                            .unwrap();
394                        request_seen.send(request).unwrap();
395                        return;
396                    }
397                    Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
398                        thread::sleep(Duration::from_millis(10));
399                    }
400                    Err(error) => panic!("loopback OTLP listener failed: {error}"),
401                }
402            }
403        });
404        let log_path = std::env::temp_dir().join(format!(
405            "nanocodex-observability-{}.jsonl",
406            std::process::id()
407        ));
408        let mut guard = ObservabilityBuilder::new("nanocodex-test", "0.0.0")
409            .filter("trace")
410            .format(LogFormat::Json)
411            .output(LogOutput::File(log_path.clone()))
412            .otlp_endpoint(format!("http://{address}"))
413            .install()
414            .unwrap();
415        {
416            let span = tracing::info_span!("test.operation", work_units = 3);
417            let _entered = span.enter();
418            tracing::info!("test event");
419        }
420        guard.shutdown().unwrap();
421        guard.shutdown().unwrap();
422        let request = request_received.recv_timeout(OTLP_TEST_TIMEOUT).unwrap();
423        assert!(
424            request.starts_with(b"POST ")
425                && request
426                    .windows(b"/v1/traces".len())
427                    .any(|window| window == b"/v1/traces"),
428            "unexpected OTLP request headers: {}",
429            String::from_utf8_lossy(&request)
430        );
431        server.join().unwrap();
432        drop(guard);
433        let log = std::fs::read_to_string(&log_path).unwrap();
434        assert!(log.lines().any(|line| line.contains("test.operation")));
435        std::fs::remove_file(log_path).unwrap();
436    }
437}