diff --git a/opentelemetry/Cargo.toml b/opentelemetry/Cargo.toml index cf89a51bb0..0748eb810e 100644 --- a/opentelemetry/Cargo.toml +++ b/opentelemetry/Cargo.toml @@ -57,3 +57,8 @@ harness = false [[bench]] name = "anyvalue" harness = false + +[[bench]] +name = "context_attach" +harness = false +required-features = ["tracing"] diff --git a/opentelemetry/benches/context_attach.rs b/opentelemetry/benches/context_attach.rs new file mode 100644 index 0000000000..8bdcb06fc2 --- /dev/null +++ b/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 { + c.benchmark_group("context_attach") +} + +#[derive(Debug, PartialEq)] +struct Value(i32); + +criterion_group!(benches, criterion_benchmark); + +criterion_main!(benches); diff --git a/opentelemetry/src/context.rs b/opentelemetry/src/context.rs index 67eae958f6..b58011fbc9 100644 --- a/opentelemetry/src/context.rs +++ b/opentelemetry/src/context.rs @@ -9,7 +9,7 @@ use std::marker::PhantomData; use std::sync::Arc; thread_local! { - static CURRENT_CONTEXT: RefCell = RefCell::new(Context::default()); + static CURRENT_CONTEXT: RefCell = RefCell::new(ContextStack::default()); } /// An execution-scoped collection of values. @@ -122,7 +122,7 @@ impl Context { /// Note: This function will panic if you attempt to attach another context /// while the current one is still borrowed. pub fn map_current(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 @@ impl Context { /// assert_eq!(Context::current().get::(), 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 @@ impl fmt::Debug for Context { /// A guard that resets the current context to the prior context when dropped. #[allow(missing_debug_implementations)] pub struct ContextGuard { - previous_cx: Option, + 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 @@ impl Hasher for IdHasher { } } +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>, +} + +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(&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 @@ mod tests { 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::(), 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::(), None); + assert_eq!(current.get::(), None); + } }