refactor: only return span instead of entire context

Author: taisho6339

opentelemetry/src/context.rs

@@ -478,11 +478,11 @@ impl Drop for ContextGuard {
     fn drop(&mut self) {
         let id = self.cx_pos;
         if id > ContextStack::BASE_POS && id < ContextStack::MAX_POS {
-            // 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 =
+            // Extract only the span to drop outside of borrow_mut to avoid panic
+            // when the span's drop implementation calls Context::current()
+            let _to_drop =
                 CURRENT_CONTEXT.with(|context_stack| context_stack.borrow_mut().pop_id(id));
-            // Context is automatically dropped here, outside of borrow_mut scope
+            // Span (if any) is automatically dropped here, outside of borrow_mut scope
         }
     }
 }
@@ -533,6 +533,12 @@ struct ContextStack {
     _marker: PhantomData<*const ()>,
 }
 
+// Type alias for what pop_id returns - only return the span when trace feature is enabled
+#[cfg(feature = "trace")]
+type PopIdReturn = Option<Arc<SynchronizedSpan>>;
+#[cfg(not(feature = "trace"))]
+type PopIdReturn = ();
+
 impl ContextStack {
     const BASE_POS: u16 = 0;
     const MAX_POS: u16 = u16::MAX;
@@ -561,7 +567,7 @@ impl ContextStack {
     }
 
     #[inline(always)]
-    fn pop_id(&mut self, pos: u16) -> Option<Context> {
+    fn pop_id(&mut self, pos: u16) -> PopIdReturn {
         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
@@ -575,7 +581,10 @@ impl ContextStack {
                     "Attempted to pop the overflow position which is not allowed"
                 }
             );
+            #[cfg(feature = "trace")]
             return None;
+            #[cfg(not(feature = "trace"))]
+            return;
         }
         let len: u16 = self.stack.len() as u16;
         // Are we at the top of the [`ContextStack`]?
@@ -588,10 +597,17 @@ 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() {
-                // Return the old context to be dropped outside of borrow_mut
-                Some(std::mem::replace(&mut self.current_cx, next_cx))
+                let old_cx = std::mem::replace(&mut self.current_cx, next_cx);
+                // Extract and return only the span to avoid cloning the entire Context
+                #[cfg(feature = "trace")]
+                return old_cx.span;
+                #[cfg(not(feature = "trace"))]
+                return;
             } else {
-                None
+                #[cfg(feature = "trace")]
+                return None;
+                #[cfg(not(feature = "trace"))]
+                return;
             }
         } else {
             // This is an out of order pop.
@@ -603,10 +619,19 @@ impl ContextStack {
                     stack_length = len,
                     message = "Attempted to pop beyond the end of the context stack"
                 );
+                #[cfg(feature = "trace")]
                 return None;
+                #[cfg(not(feature = "trace"))]
+                return;
+            }
+            // Clear out the entry at the given id and extract its span
+            #[cfg(feature = "trace")]
+            return self.stack[pos as usize].take().and_then(|cx| cx.span);
+            #[cfg(not(feature = "trace"))]
+            {
+                self.stack[pos as usize] = None;
+                return;
             }
-            // Clear out the entry at the given id and return it
-            self.stack[pos as usize].take()
         }
     }