From ab629c304ea581d241e0d8fc76164eccf6573fe2 Mon Sep 17 00:00:00 2001 From: webfinesse Date: Fri, 18 Oct 2024 04:05:00 -0400 Subject: [PATCH] adding support for setting the status of the span through SpanExt (#176) ## Motivation I have found that when working with `tracing` spans created by 3rd party libraries, it is incredibly difficult to set the OpenTelemetry Status of the span, which results in always having an unset span in my visualization tool. image In this case, I am working with the span created by the `aws_lambda_runtime` [crate](https://github.com/awslabs/aws-lambda-rust-runtime/blob/main/lambda-runtime/src/layers/otel.rs#L75). Unfortunately, since the span in the `aws_lambda_runtime` crate does not include the `otel.status_code` attribute upfront the `OpenTelemetryLayer::SpanAttributeVisitor::record_str` and `OpenTelemetryLayer::SpanAttributeVisitor::record_debug` never gets triggered when calling `tracing::Span::current().record("otel.status_code", "ok");` This behavior of the `record` function ignoring new fields is documented [here](https://docs.rs/tracing/latest/tracing/struct.Span.html#method.record) ## Solution My solution to this is add a new member `OpenTelemetrySpanExt::set_status` that takes an `opentelemetry::trace::Status` enum that gets written to the underlying `SpanBuilder`. --- src/span_ext.rs | 32 +++++++++++++++++++- tests/span_ext.rs | 76 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) create mode 100644 tests/span_ext.rs 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() +}