fix: Allow overlapping context scopes

opentelemetry/Cargo.toml

@@ -57,3 +57,8 @@ harness = false
 [[bench]]
 name = "anyvalue"
 harness = false
+
+[[bench]]
+name = "context_attach"
+harness = false
+required-features = ["tracing"]

opentelemetry/benches/context_attach.rs

@@ -0,0 +1,119 @@
+use criterion::{
+    black_box, criterion_group, criterion_main, measurement::WallTime, BenchmarkGroup, BenchmarkId,
+    Criterion, Throughput,
+};
+use opentelemetry::{
+    trace::{SpanContext, TraceContextExt},
+    Context,
+};
+
+// Run this benchmark with:
+// cargo bench --bench context_attach
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let span_context = Context::new().with_remote_span_context(SpanContext::empty_context());
+    let contexts = vec![
+        ("empty_cx", Context::new()),
+        ("single_value_cx", Context::new().with_value(Value(4711))),
+        ("span_cx", span_context),
+    ];
+    for (name, cx) in contexts {
+        single_cx_scope(&mut group(c), name, &cx);
+        nested_cx_scope(&mut group(c), name, &cx);
+        overlapping_cx_scope(&mut group(c), name, &cx);
+    }
+}
+
+fn single_cx_scope(
+    group: &mut BenchmarkGroup<'_, WallTime>,
+    context_type: &str,
+    context: &Context,
+) {
+    let _restore = Context::current().attach();
+    group.throughput(Throughput::Elements(1)).bench_function(
+        BenchmarkId::new("single_cx_scope", context_type),
+        |b| {
+            b.iter_batched(
+                || context.clone(),
+                |cx| {
+                    single_cx(cx);
+                },
+                criterion::BatchSize::SmallInput,
+            );
+        },
+    );
+}
+
+#[inline(never)]
+fn single_cx(cx: Context) {
+    let _ = black_box(cx.attach());
+    let _ = black_box(dummy_work());
+}
+
+fn nested_cx_scope(group: &mut BenchmarkGroup<'_, WallTime>, cx_type: &str, context: &Context) {
+    let _restore = Context::current().attach();
+    group.throughput(Throughput::Elements(1)).bench_function(
+        BenchmarkId::new("nested_cx_scope", cx_type),
+        |b| {
+            b.iter_batched(
+                || (context.clone(), context.clone()),
+                |(cx1, cx2)| {
+                    nested_cx(cx1, cx2);
+                },
+                criterion::BatchSize::SmallInput,
+            );
+        },
+    );
+}
+
+#[inline(never)]
+fn nested_cx(cx1: Context, cx2: Context) {
+    let _ = black_box(cx1.attach());
+    let _ = black_box(cx2.attach());
+    let _ = black_box(dummy_work());
+}
+
+fn overlapping_cx_scope(
+    group: &mut BenchmarkGroup<'_, WallTime>,
+    cx_type: &str,
+    context: &Context,
+) {
+    let _restore = Context::current().attach();
+    group.throughput(Throughput::Elements(1)).bench_function(
+        BenchmarkId::new("overlapping_cx_scope", cx_type),
+        |b| {
+            b.iter_batched(
+                || (context.clone(), context.clone()),
+                |(cx1, cx2)| {
+                    overlapping_cx(cx1, cx2);
+                },
+                criterion::BatchSize::SmallInput,
+            );
+        },
+    );
+}
+
+#[inline(never)]
+fn overlapping_cx(cx1: Context, cx2: Context) {
+    let outer = cx1.attach();
+    let inner = cx2.attach();
+    let _ = black_box(dummy_work());
+    drop(outer);
+    drop(inner);
+}
+
+#[inline(never)]
+fn dummy_work() -> i32 {
+    black_box(1 + 1)
+}
+
+fn group(c: &mut Criterion) -> BenchmarkGroup<WallTime> {
+    c.benchmark_group("context_attach")
+}
+
+#[derive(Debug, PartialEq)]
+struct Value(i32);
+
+criterion_group!(benches, criterion_benchmark);
+
+criterion_main!(benches);

opentelemetry/src/context.rs

@@ -9,7 +9,7 @@
 use std::sync::Arc;
 
 thread_local! {
-    static CURRENT_CONTEXT: RefCell<Context> = RefCell::new(Context::default());
+    static CURRENT_CONTEXT: RefCell<ContextStack> = RefCell::new(ContextStack::default());
 }
 
 /// An execution-scoped collection of values.
@@ -122,7 +122,7 @@
     /// Note: This function will panic if you attempt to attach another context
     /// while the current one is still borrowed.
     pub fn map_current<T>(f: impl FnOnce(&Context) -> T) -> T {
-        CURRENT_CONTEXT.with(|cx| f(&cx.borrow()))
+        CURRENT_CONTEXT.with(|cx| cx.borrow().map_current_cx(f))
     }
 
     /// Returns a clone of the current thread's context with the given value.
@@ -298,12 +298,10 @@
     /// assert_eq!(Context::current().get::<ValueA>(), None);
     /// ```
     pub fn attach(self) -> ContextGuard {
-        let previous_cx = CURRENT_CONTEXT
-            .try_with(|current| current.replace(self))
-            .ok();
+        let cx_id = CURRENT_CONTEXT.with(|cx| cx.borrow_mut().push(self));
 
         ContextGuard {
-            previous_cx,
+            cx_id,
             _marker: PhantomData,
         }
     }
@@ -336,15 +334,16 @@
 /// A guard that resets the current context to the prior context when dropped.
 #[allow(missing_debug_implementations)]
 pub struct ContextGuard {
-    previous_cx: Option<Context>,
+    cx_id: usize,
     // ensure this type is !Send as it relies on thread locals
     _marker: PhantomData<*const ()>,
 }
 
 impl Drop for ContextGuard {
     fn drop(&mut self) {
-        if let Some(previous_cx) = self.previous_cx.take() {
-            let _ = CURRENT_CONTEXT.try_with(|current| current.replace(previous_cx));
+        let id = self.cx_id;
+        if id > 0 {
+            CURRENT_CONTEXT.with(|context_stack| context_stack.borrow_mut().pop_id(id));
         }
     }
 }
@@ -371,6 +370,75 @@
     }
 }
 
+struct ContextStack {
+    current_cx: Context,
+    current_id: usize,
+    // TODO:ban wrap the whole id thing in its own type
+    id_count: usize,
+    // TODO:ban wrap the the tuple in its own type
+    stack: Vec<Option<(usize, Context)>>,
+}
+
+impl ContextStack {
+    #[inline(always)]
+    fn push(&mut self, cx: Context) -> usize {
+        self.id_count += 512; // TODO:ban clean up this
+        let next_id = self.stack.len() + 1 + self.id_count;
+        let current_cx = std::mem::replace(&mut self.current_cx, cx);
+        self.stack.push(Some((self.current_id, current_cx)));
+        self.current_id = next_id;
+        next_id
+    }
+
+    #[inline(always)]
+    fn pop_id(&mut self, id: usize) {
+        if id == 0 {
+            return;
+        }
+        // Are we at the top of the stack?
+        if id == self.current_id {
+            // Shrink the stack if possible
+            while let Some(None) = self.stack.last() {
+                self.stack.pop();
+            }
+            // There is always the initial context at the bottom of the stack
+            if let Some(Some((next_id, next_cx))) = self.stack.pop() {
+                self.current_cx = next_cx;
+                self.current_id = next_id;
+            }
+        } else {
+            let pos = id & 511; // TODO:ban clean up this
+            if pos >= self.stack.len() {
+                // This is an invalid id, ignore it
+                return;
+            }
+            if let Some((pos_id, _)) = self.stack[pos] {
+                // Is the correct id at this position?
+                if pos_id == id {
+                    // Clear out this entry
+                    self.stack[pos] = None;
+                }
+            }
+        }
+    }
+
+    #[inline(always)]
+    fn map_current_cx<T>(&self, f: impl FnOnce(&Context) -> T) -> T {
+        f(&self.current_cx)
+    }
+}
+
+impl Default for ContextStack {
+    fn default() -> Self {
+        ContextStack {
+            current_id: 0,
+            current_cx: Context::default(),
+            id_count: 0,
+            stack: Vec::with_capacity(64),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -413,4 +481,45 @@
             true
         }));
     }
+
+    #[test]
+    fn overlapping_contexts() {
+        #[derive(Debug, PartialEq)]
+        struct ValueA(&'static str);
+        #[derive(Debug, PartialEq)]
+        struct ValueB(u64);
+
+        let outer_guard = Context::new().with_value(ValueA("a")).attach();
+
+        // Only value `a` is set
+        let current = Context::current();
+        assert_eq!(current.get(), Some(&ValueA("a")));
+        assert_eq!(current.get::<ValueB>(), None);
+
+        let inner_guard = Context::current_with_value(ValueB(42)).attach();
+        // Both values are set in inner context
+        let current = Context::current();
+        assert_eq!(current.get(), Some(&ValueA("a")));
+        assert_eq!(current.get(), Some(&ValueB(42)));
+
+        assert!(Context::map_current(|cx| {
+            assert_eq!(cx.get(), Some(&ValueA("a")));
+            assert_eq!(cx.get(), Some(&ValueB(42)));
+            true
+        }));
+
+        drop(outer_guard);
+
+        // `inner_guard` is still alive so both `ValueA` and `ValueB` should still be accessible
+        let current = Context::current();
+        assert_eq!(current.get(), Some(&ValueA("a")));
+        assert_eq!(current.get(), Some(&ValueB(42)));
+
+        drop(inner_guard);
+
+        // Both guards are dropped and neither value should be accessible.
+        let current = Context::current();
+        assert_eq!(current.get::<ValueA>(), None);
+        assert_eq!(current.get::<ValueB>(), None);
+    }
 }