From 9c78908dbf5de77484e91d78ce974cba343c402b Mon Sep 17 00:00:00 2001 From: Brett Knapik Date: Thu, 17 Oct 2024 00:45:44 -0400 Subject: [PATCH 1/2] adding support for setting the status of the span through an extension --- src/span_ext.rs | 32 +++++++++++++++++++- tests/span_ext.rs | 75 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 1 deletion(-) create mode 100644 tests/span_ext.rs diff --git a/src/span_ext.rs b/src/span_ext.rs index c89dc09..c99f374 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..0f1376f --- /dev/null +++ b/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>>); + +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() +} \ No newline at end of file From f5d6c56fea29ec82d722e1adbdc77675ef02f353 Mon Sep 17 00:00:00 2001 From: Brett Knapik Date: Thu, 17 Oct 2024 18:28:07 -0400 Subject: [PATCH 2/2] format fixes --- src/span_ext.rs | 6 +++--- tests/span_ext.rs | 9 +++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/span_ext.rs b/src/span_ext.rs index c99f374..3dddd64 100644 --- a/src/span_ext.rs +++ b/src/span_ext.rs @@ -137,19 +137,19 @@ pub trait OpenTelemetrySpanExt { /// 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); } diff --git a/tests/span_ext.rs b/tests/span_ext.rs index 0f1376f..acfdcdf 100644 --- a/tests/span_ext.rs +++ b/tests/span_ext.rs @@ -51,13 +51,14 @@ fn set_status_ok() { #[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 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 -{ +fn set_status_helper(status: Status) -> SpanData { let (_tracer, provider, exporter, subscriber) = test_tracer(); tracing::subscriber::with_default(subscriber, || { @@ -72,4 +73,4 @@ fn set_status_helper(status: Status) -> SpanData assert_eq!(spans.len(), 1); spans.iter().find(|s| s.name == "root").unwrap().clone() -} \ No newline at end of file +}