adding support for setting the status of the span through an extension

Author: webfinesse

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<Key>, value: impl Into<Value>);
+
+    /// 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::<WithContext>() {
+                get_context.with_context(subscriber, id, move |builder, _| {
+                    builder.builder.status = status.take().unwrap();
+                });
+            }
+        });
+    }
 }

tests/span_ext.rs

@@ -0,0 +1,75 @@
+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<Mutex<Vec<SpanData>>>);
+
+impl SpanExporter for TestExporter {
+    fn export(&mut self, mut batch: Vec<SpanData>) -> 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()
+}