layer: Fix race condition in on_event

Author: mladedav

src/layer.rs

@@ -1,7 +1,7 @@
 use crate::{OtelData, PreSampledTracer};
 use once_cell::unsync;
 use opentelemetry::{
-    trace::{self as otel, noop, TraceContextExt},
+    trace::{self as otel, noop, Status, TraceContextExt},
     Context as OtelContext, Key, KeyValue, StringValue, Value,
 };
 use std::any::TypeId;
@@ -117,9 +117,15 @@ fn str_to_status(s: &str) -> otel::Status {
     }
 }
 
+#[derive(Default)]
+struct ErrorState {
+    status: Option<Status>,
+    attributes: Option<Vec<KeyValue>>,
+}
+
 struct SpanEventVisitor<'a, 'b> {
     event_builder: &'a mut otel::Event,
-    span_builder: Option<&'b mut otel::SpanBuilder>,
+    error_state: &'b mut Option<ErrorState>,
     sem_conv_config: SemConvConfig,
 }
 
@@ -186,9 +192,10 @@ impl<'a, 'b> field::Visit for SpanEventVisitor<'a, 'b> {
             // In both cases, an event with an empty name and with an error attribute is created.
             "error" if self.event_builder.name.is_empty() => {
                 if self.sem_conv_config.error_events_to_status {
-                    if let Some(span) = &mut self.span_builder {
-                        span.status = otel::Status::error(format!("{:?}", value));
-                    }
+                    self.error_state
+                        .get_or_insert_with(ErrorState::default)
+                        .status
+                        .replace(otel::Status::error(format!("{:?}", value)));
                 }
                 if self.sem_conv_config.error_events_to_exceptions {
                     self.event_builder.name = EVENT_EXCEPTION_NAME.into();
@@ -225,9 +232,10 @@ impl<'a, 'b> field::Visit for SpanEventVisitor<'a, 'b> {
             // In both cases, an event with an empty name and with an error attribute is created.
             "error" if self.event_builder.name.is_empty() => {
                 if self.sem_conv_config.error_events_to_status {
-                    if let Some(span) = &mut self.span_builder {
-                        span.status = otel::Status::error(format!("{:?}", value));
-                    }
+                    self.error_state
+                        .get_or_insert_with(ErrorState::default)
+                        .status
+                        .replace(otel::Status::error(format!("{:?}", value)));
                 }
                 if self.sem_conv_config.error_events_to_exceptions {
                     self.event_builder.name = EVENT_EXCEPTION_NAME.into();
@@ -288,25 +296,27 @@ impl<'a, 'b> field::Visit for SpanEventVisitor<'a, 'b> {
         }
 
         if self.sem_conv_config.error_records_to_exceptions {
-            if let Some(span) = &mut self.span_builder {
-                if let Some(attrs) = span.attributes.as_mut() {
-                    attrs.push(KeyValue::new(
-                        FIELD_EXCEPTION_MESSAGE,
-                        Value::String(error_msg.clone().into()),
-                    ));
+            let attributes = self
+                .error_state
+                .get_or_insert_with(ErrorState::default)
+                .attributes
+                .get_or_insert_with(Vec::new);
 
-                    // NOTE: This is actually not the stacktrace of the exception. This is
-                    // the "source chain". It represents the heirarchy of errors from the
-                    // app level to the lowest level such as IO. It does not represent all
-                    // of the callsites in the code that led to the error happening.
-                    // `std::error::Error::backtrace` is a nightly-only API and cannot be
-                    // used here until the feature is stabilized.
-                    attrs.push(KeyValue::new(
-                        FIELD_EXCEPTION_STACKTRACE,
-                        Value::Array(chain.clone().into()),
-                    ));
-                }
-            }
+            attributes.push(KeyValue::new(
+                FIELD_EXCEPTION_MESSAGE,
+                Value::String(error_msg.clone().into()),
+            ));
+
+            // NOTE: This is actually not the stacktrace of the exception. This is
+            // the "source chain". It represents the heirarchy of errors from the
+            // app level to the lowest level such as IO. It does not represent all
+            // of the callsites in the code that led to the error happening.
+            // `std::error::Error::backtrace` is a nightly-only API and cannot be
+            // used here until the feature is stabilized.
+            attributes.push(KeyValue::new(
+                FIELD_EXCEPTION_STACKTRACE,
+                Value::Array(chain.clone().into()),
+            ));
         }
 
         self.event_builder
@@ -1023,24 +1033,24 @@ where
             #[cfg(not(feature = "tracing-log"))]
             let target = target.string(meta.target());
 
-            // Move out extension data to not hold the extensions lock across the event.record() call, which could result in a deadlock
-            let mut otel_data = span.extensions_mut().remove::<OtelData>();
-            let span_builder = otel_data.as_mut().map(|data| &mut data.builder);
-
             let mut otel_event = otel::Event::new(
                 String::new(),
                 crate::time::now(),
                 vec![Key::new("level").string(meta.level().as_str()), target],
                 0,
             );
 
+            let mut error_state = None;
             event.record(&mut SpanEventVisitor {
                 event_builder: &mut otel_event,
-                span_builder,
+                error_state: &mut error_state,
                 sem_conv_config: self.sem_conv_config,
             });
 
-            if let Some(mut otel_data) = otel_data {
+            let mut extensions = span.extensions_mut();
+            let otel_data = extensions.get_mut::<OtelData>();
+
+            if let Some(otel_data) = otel_data {
                 let builder = &mut otel_data.builder;
 
                 if builder.status == otel::Status::Unset
@@ -1049,6 +1059,18 @@ where
                     builder.status = otel::Status::error("")
                 }
 
+                if let Some(error_state) = error_state {
+                    if let Some(status) = error_state.status {
+                        builder.status = status;
+                    }
+                    if let Some(attributes) = error_state.attributes {
+                        builder
+                            .attributes
+                            .get_or_insert_with(Vec::new)
+                            .extend(attributes);
+                    }
+                }
+
                 if self.location {
                     #[cfg(not(feature = "tracing-log"))]
                     let normalized_meta: Option<tracing_core::Metadata<'_>> = None;
@@ -1085,8 +1107,6 @@ where
                 } else {
                     builder.events = Some(vec![otel_event]);
                 }
-
-                span.extensions_mut().replace(otel_data);
             }
         };
     }