return parent context in on_end instead of causing panic

Author: taisho6339

opentelemetry-sdk/src/trace/mod.rs

@@ -137,10 +137,24 @@ mod tests {
             }
         }
 
-        fn on_end(&self, _span: SpanData) {
-            // TODO: Accessing Context::current() will panic today and hence commented out.
+        fn on_end(&self, span: SpanData) {
+            // Fixed: Context::current() no longer panics from Drop
             // See https://github.com/open-telemetry/opentelemetry-rust/issues/2871
-            // let _c = Context::current();
+            let current_cx = Context::current();
+
+            // IMPORTANT NOTE: The context returned here is the PARENT context,
+            // NOT the span's context (which has already been popped from the stack).
+            // The span's baggage was extracted in on_start and stored as attributes.
+
+            // Verify: on_start stored the baggage as an attribute
+            assert!(
+                span.attributes.iter().any(|kv| kv.key.as_str() == "bag-key"),
+                "Baggage should have been stored as span attribute in on_start"
+            );
+
+            // The current context's baggage is the parent's baggage (if any)
+            println!("on_end: parent context baggage = {:?}",
+                current_cx.baggage().iter().map(|(k, _)| k).collect::<Vec<_>>());
         }
 
         fn force_flush(&self) -> crate::error::OTelSdkResult {

opentelemetry-sdk/src/trace/span_processor.rs

@@ -79,9 +79,43 @@ pub trait SpanProcessor: Send + Sync + std::fmt::Debug {
     /// synchronously on the thread that started the span, therefore it should
     /// not block or throw exceptions.
     fn on_start(&self, span: &mut Span, cx: &Context);
+
     /// `on_end` is called after a `Span` is ended (i.e., the end timestamp is
     /// already set). This method is called synchronously within the `Span::end`
     /// API, therefore it should not block or throw an exception.
+    ///
+    /// # Accessing Context
+    ///
+    /// If you need to access context information (such as baggage) in `on_end`,
+    /// note that calling [`Context::current()`] will return the **parent context**,
+    /// not the span's context (which has already been removed from the context stack).
+    ///
+    /// **Best Practice**: Extract any needed context information in [`on_start`]
+    /// and store it as span attributes. This ensures the information is available
+    /// in the [`SpanData`] passed to `on_end`.
+    ///
+    /// # Example
+    ///
+    /// ```rust,ignore
+    /// impl SpanProcessor for MyProcessor {
+    ///     fn on_start(&self, span: &mut Span, cx: &Context) {
+    ///         // Extract baggage and store as span attribute
+    ///         if let Some(value) = cx.baggage().get("my-key") {
+    ///             span.set_attribute(KeyValue::new("my-key", value.to_string()));
+    ///         }
+    ///     }
+    ///
+    ///     fn on_end(&self, span: SpanData) {
+    ///         // Access the attribute stored in on_start
+    ///         let my_value = span.attributes.iter()
+    ///             .find(|kv| kv.key.as_str() == "my-key");
+    ///     }
+    /// }
+    /// ```
+    ///
+    /// [`on_start`]: SpanProcessor::on_start
+    /// [`Context::current()`]: opentelemetry::Context::current
+    ///
     /// TODO - This method should take reference to `SpanData`
     fn on_end(&self, span: SpanData);
     /// Force the spans lying in the cache to be exported.

opentelemetry/src/context.rs

@@ -115,6 +115,18 @@ impl Context {
 
     /// Returns an immutable snapshot of the current thread's context.
     ///
+    /// # Behavior During Context Drop
+    ///
+    /// When called from within a [`Drop`] implementation that is triggered by
+    /// a [`ContextGuard`] being dropped (e.g., when a [`Span`] is dropped as part
+    /// of context cleanup), this function returns the **parent context**, not the
+    /// context being dropped.
+    ///
+    /// This behavior is by design and prevents panics that would otherwise occur
+    /// from attempting to borrow the context while it's being mutably borrowed
+    /// for cleanup. See [issue #2871](https://github.com/open-telemetry/opentelemetry-rust/issues/2871)
+    /// for details.
+    ///
     /// # Examples
     ///
     /// ```
@@ -130,6 +142,8 @@ impl Context {
     /// let _guard = Context::new().with_value(ValueA("a")).attach();
     /// do_work()
     /// ```
+    ///
+    /// [`Span`]: crate::trace::Span
     pub fn current() -> Self {
         Self::map_current(|cx| cx.clone())
     }
@@ -464,7 +478,12 @@ impl Drop for ContextGuard {
     fn drop(&mut self) {
         let id = self.cx_pos;
         if id > ContextStack::BASE_POS && id < ContextStack::MAX_POS {
-            CURRENT_CONTEXT.with(|context_stack| context_stack.borrow_mut().pop_id(id));
+            // Extract the context to drop outside of borrow_mut to avoid panic
+            // when the context's drop implementation (e.g., Span::drop) calls Context::current()
+            let _context_to_drop = CURRENT_CONTEXT.with(|context_stack| {
+                context_stack.borrow_mut().pop_id(id)
+            });
+            // Context is automatically dropped here, outside of borrow_mut scope
         }
     }
 }
@@ -543,7 +562,7 @@ impl ContextStack {
     }
 
     #[inline(always)]
-    fn pop_id(&mut self, pos: u16) {
+    fn pop_id(&mut self, pos: u16) -> Option<Context> {
         if pos == ContextStack::BASE_POS || pos == ContextStack::MAX_POS {
             // The empty context is always at the bottom of the [`ContextStack`]
             // and cannot be popped, and the overflow position is invalid, so do
@@ -557,7 +576,7 @@ impl ContextStack {
                     "Attempted to pop the overflow position which is not allowed"
                 }
             );
-            return;
+            return None;
         }
         let len: u16 = self.stack.len() as u16;
         // Are we at the top of the [`ContextStack`]?
@@ -570,7 +589,10 @@ impl ContextStack {
             // empty context is always at the bottom of the stack if the
             // [`ContextStack`] is not empty.
             if let Some(Some(next_cx)) = self.stack.pop() {
-                self.current_cx = next_cx;
+                // Return the old context to be dropped outside of borrow_mut
+                Some(std::mem::replace(&mut self.current_cx, next_cx))
+            } else {
+                None
             }
         } else {
             // This is an out of order pop.
@@ -582,10 +604,10 @@ impl ContextStack {
                     stack_length = len,
                     message = "Attempted to pop beyond the end of the context stack"
                 );
-                return;
+                return None;
             }
-            // Clear out the entry at the given id.
-            _ = self.stack[pos as usize].take();
+            // Clear out the entry at the given id and return it
+            self.stack[pos as usize].take()
         }
     }