diff --git a/src/span_ext.rs b/src/span_ext.rs index c89dc09..3dddd64 100644 --- a/src/span_ext.rs +++ b/src/span_ext.rs @@ -1,5 +1,5 @@ use crate::layer::WithContext; -use opentelemetry::{trace::SpanContext, Context, Key, KeyValue, Value}; +use opentelemetry::{trace::SpanContext, trace::Status, Context, Key, KeyValue, Value}; /// Utility functions to allow tracing [`Span`]s to accept and return /// [OpenTelemetry] [`Context`]s. @@ -133,6 +133,25 @@ pub trait OpenTelemetrySpanExt { /// app_root.set_attribute("http.request.header.x_forwarded_for", "example"); /// ``` fn set_attribute(&self, key: impl Into, value: impl Into); + + /// Sets an OpenTelemetry status for this span. + /// This is useful for setting the status of a span that was created by a library that does not declare + /// the otel.status_code field of the span in advance. + /// + /// # Examples + /// + /// ```rust + /// use opentelemetry::trace::Status; + /// use tracing_opentelemetry::OpenTelemetrySpanExt; + /// use tracing::Span; + /// + /// /// // Generate a tracing span as usual + /// let app_root = tracing::span!(tracing::Level::INFO, "app_start"); + /// + /// // Set the Status of the span to `Status::Ok`. + /// app_root.set_status(Status::Ok); + /// ``` + fn set_status(&self, status: Status); } impl OpenTelemetrySpanExt for tracing::Span { @@ -207,4 +226,15 @@ impl OpenTelemetrySpanExt for tracing::Span { } }); } + + fn set_status(&self, status: Status) { + self.with_subscriber(move |(id, subscriber)| { + let mut status = Some(status); + if let Some(get_context) = subscriber.downcast_ref::() { + get_context.with_context(subscriber, id, move |builder, _| { + builder.builder.status = status.take().unwrap(); + }); + } + }); + } } diff --git a/tests/span_ext.rs b/tests/span_ext.rs new file mode 100644 index 0000000..acfdcdf --- /dev/null +++ b/tests/span_ext.rs @@ -0,0 +1,76 @@ +use futures_util::future::BoxFuture; +use opentelemetry::trace::{Status, TracerProvider as _}; +use opentelemetry_sdk::{ + export::trace::{ExportResult, SpanData, SpanExporter}, + trace::{Tracer, TracerProvider}, +}; +use std::sync::{Arc, Mutex}; +use tracing::level_filters::LevelFilter; +use tracing::Subscriber; +use tracing_opentelemetry::{layer, OpenTelemetrySpanExt}; +use tracing_subscriber::prelude::*; + +#[derive(Clone, Default, Debug)] +struct TestExporter(Arc>>); + +impl SpanExporter for TestExporter { + fn export(&mut self, mut batch: Vec) -> BoxFuture<'static, ExportResult> { + let spans = self.0.clone(); + Box::pin(async move { + if let Ok(mut inner) = spans.lock() { + inner.append(&mut batch); + } + Ok(()) + }) + } +} + +fn test_tracer() -> (Tracer, TracerProvider, TestExporter, impl Subscriber) { + let exporter = TestExporter::default(); + let provider = TracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let tracer = provider.tracer("test"); + + let subscriber = tracing_subscriber::registry() + .with( + layer() + .with_tracer(tracer.clone()) + .with_filter(LevelFilter::DEBUG), + ) + .with(tracing_subscriber::fmt::layer().with_filter(LevelFilter::TRACE)); + + (tracer, provider, exporter, subscriber) +} + +#[test] +fn set_status_ok() { + let root_span = set_status_helper(Status::Ok); + assert_eq!(Status::Ok, root_span.status); +} + +#[test] +fn set_status_error() { + let expected_error = Status::Error { + description: std::borrow::Cow::Borrowed("Elon put in too much fuel in his rocket!"), + }; + let root_span = set_status_helper(expected_error.clone()); + assert_eq!(expected_error, root_span.status); +} + +fn set_status_helper(status: Status) -> SpanData { + let (_tracer, provider, exporter, subscriber) = test_tracer(); + + tracing::subscriber::with_default(subscriber, || { + let root = tracing::debug_span!("root").entered(); + + root.set_status(status); + }); + + drop(provider); // flush all spans + let spans = exporter.0.lock().unwrap(); + + assert_eq!(spans.len(), 1); + + spans.iter().find(|s| s.name == "root").unwrap().clone() +}