opentelemetry/context.rs
1//! Execution-scoped context propagation.
2//!
3//! The `context` module provides mechanisms for propagating values across API boundaries and between
4//! logically associated execution units. It enables cross-cutting concerns to access their data in-process
5//! using a shared context object.
6//!
7//! # Main Types
8//!
9//! - [`Context`]: An immutable, execution-scoped collection of values.
10//!
11
12use crate::otel_warn;
13#[cfg(feature = "trace")]
14use crate::trace::context::SynchronizedSpan;
15use std::any::{Any, TypeId};
16use std::cell::RefCell;
17use std::collections::HashMap;
18use std::fmt;
19use std::hash::{BuildHasherDefault, Hasher};
20use std::marker::PhantomData;
21use std::sync::Arc;
22
23#[cfg(feature = "futures")]
24mod future_ext;
25
26#[cfg(feature = "futures")]
27pub use future_ext::{FutureExt, WithContext};
28
29thread_local! {
30 static CURRENT_CONTEXT: RefCell<ContextStack> = RefCell::new(ContextStack::default());
31}
32
33/// An execution-scoped collection of values.
34///
35/// A [`Context`] is a propagation mechanism which carries execution-scoped
36/// values across API boundaries and between logically associated execution
37/// units. Cross-cutting concerns access their data in-process using the same
38/// shared context object.
39///
40/// [`Context`]s are immutable, and their write operations result in the creation
41/// of a new context containing the original values and the new specified values.
42///
43/// ## Context state
44///
45/// Concerns can create and retrieve their local state in the current execution
46/// state represented by a context through the [`get`] and [`with_value`]
47/// methods. It is recommended to use application-specific types when storing new
48/// context values to avoid unintentionally overwriting existing state.
49///
50/// ## Managing the current context
51///
52/// Contexts can be associated with the caller's current execution unit on a
53/// given thread via the [`attach`] method, and previous contexts can be restored
54/// by dropping the returned [`ContextGuard`]. Context can be nested, and will
55/// restore their parent outer context when detached on drop. To access the
56/// values of the context, a snapshot can be created via the [`Context::current`]
57/// method.
58///
59/// [`Context::current`]: Context::current()
60/// [`get`]: Context::get()
61/// [`with_value`]: Context::with_value()
62/// [`attach`]: Context::attach()
63///
64/// # Examples
65///
66/// ```
67/// use opentelemetry::Context;
68///
69/// // Application-specific `a` and `b` values
70/// #[derive(Debug, PartialEq)]
71/// struct ValueA(&'static str);
72/// #[derive(Debug, PartialEq)]
73/// struct ValueB(u64);
74///
75/// let _outer_guard = Context::new().with_value(ValueA("a")).attach();
76///
77/// // Only value a has been set
78/// let current = Context::current();
79/// assert_eq!(current.get::<ValueA>(), Some(&ValueA("a")));
80/// assert_eq!(current.get::<ValueB>(), None);
81///
82/// {
83/// let _inner_guard = Context::current_with_value(ValueB(42)).attach();
84/// // Both values are set in inner context
85/// let current = Context::current();
86/// assert_eq!(current.get::<ValueA>(), Some(&ValueA("a")));
87/// assert_eq!(current.get::<ValueB>(), Some(&ValueB(42)));
88/// }
89///
90/// // Resets to only the `a` value when inner guard is dropped
91/// let current = Context::current();
92/// assert_eq!(current.get::<ValueA>(), Some(&ValueA("a")));
93/// assert_eq!(current.get::<ValueB>(), None);
94/// ```
95#[derive(Clone, Default)]
96pub struct Context {
97 #[cfg(feature = "trace")]
98 pub(crate) span: Option<Arc<SynchronizedSpan>>,
99 entries: Option<Arc<EntryMap>>,
100 suppress_telemetry: bool,
101}
102
103type EntryMap = HashMap<TypeId, Arc<dyn Any + Sync + Send>, BuildHasherDefault<IdHasher>>;
104
105impl Context {
106 /// Creates an empty `Context`.
107 ///
108 /// The context is initially created with a capacity of 0, so it will not
109 /// allocate. Use [`with_value`] to create a new context that has entries.
110 ///
111 /// [`with_value`]: Context::with_value()
112 pub fn new() -> Self {
113 Context::default()
114 }
115
116 /// Returns an immutable snapshot of the current thread's context.
117 ///
118 /// # Behavior During Context Drop
119 ///
120 /// When called from within a [`Drop`] implementation that is triggered by
121 /// a [`ContextGuard`] being dropped (e.g., when a [`Span`] is dropped as part
122 /// of context cleanup), this function returns **whatever context happens to be
123 /// current after the guard is popped**, not the context being dropped.
124 ///
125 /// **Important**: The returned context may be completely unrelated to the
126 /// context being dropped, as contexts can be activated in any order and are
127 /// not necessarily hierarchical. Do not rely on any relationship between them.
128 ///
129 /// This behavior is by design and prevents panics that would otherwise occur
130 /// from attempting to borrow the context while it's being mutably borrowed
131 /// for cleanup. See [issue #2871](https://github.com/open-telemetry/opentelemetry-rust/issues/2871)
132 /// for details.
133 ///
134 /// # Examples
135 ///
136 /// ```
137 /// use opentelemetry::Context;
138 ///
139 /// #[derive(Debug, PartialEq)]
140 /// struct ValueA(&'static str);
141 ///
142 /// fn do_work() {
143 /// assert_eq!(Context::current().get(), Some(&ValueA("a")));
144 /// }
145 ///
146 /// let _guard = Context::new().with_value(ValueA("a")).attach();
147 /// do_work()
148 /// ```
149 ///
150 /// [`Span`]: crate::trace::Span
151 pub fn current() -> Self {
152 Self::map_current(|cx| cx.clone())
153 }
154
155 /// Applies a function to the current context returning its value.
156 ///
157 /// This can be used to build higher performing algebraic expressions for
158 /// optionally creating a new context without the overhead of cloning the
159 /// current one and dropping it.
160 ///
161 /// Note: This function will panic if you attempt to attach another context
162 /// while the current one is still borrowed.
163 pub fn map_current<T>(f: impl FnOnce(&Context) -> T) -> T {
164 CURRENT_CONTEXT.with(|cx| cx.borrow().map_current_cx(f))
165 }
166
167 /// Returns a clone of the current thread's context with the given value.
168 ///
169 /// This is a more efficient form of `Context::current().with_value(value)`
170 /// as it avoids the intermediate context clone.
171 ///
172 /// # Examples
173 ///
174 /// ```
175 /// use opentelemetry::Context;
176 ///
177 /// // Given some value types defined in your application
178 /// #[derive(Debug, PartialEq)]
179 /// struct ValueA(&'static str);
180 /// #[derive(Debug, PartialEq)]
181 /// struct ValueB(u64);
182 ///
183 /// // You can create and attach context with the first value set to "a"
184 /// let _guard = Context::new().with_value(ValueA("a")).attach();
185 ///
186 /// // And create another context based on the fist with a new value
187 /// let all_current_and_b = Context::current_with_value(ValueB(42));
188 ///
189 /// // The second context now contains all the current values and the addition
190 /// assert_eq!(all_current_and_b.get::<ValueA>(), Some(&ValueA("a")));
191 /// assert_eq!(all_current_and_b.get::<ValueB>(), Some(&ValueB(42)));
192 /// ```
193 pub fn current_with_value<T: 'static + Send + Sync>(value: T) -> Self {
194 Self::map_current(|cx| cx.with_value(value))
195 }
196
197 /// Returns a reference to the entry for the corresponding value type.
198 ///
199 /// # Examples
200 ///
201 /// ```
202 /// use opentelemetry::Context;
203 ///
204 /// // Given some value types defined in your application
205 /// #[derive(Debug, PartialEq)]
206 /// struct ValueA(&'static str);
207 /// #[derive(Debug, PartialEq)]
208 /// struct MyUser();
209 ///
210 /// let cx = Context::new().with_value(ValueA("a"));
211 ///
212 /// // Values can be queried by type
213 /// assert_eq!(cx.get::<ValueA>(), Some(&ValueA("a")));
214 ///
215 /// // And return none if not yet set
216 /// assert_eq!(cx.get::<MyUser>(), None);
217 /// ```
218 pub fn get<T: 'static>(&self) -> Option<&T> {
219 self.entries
220 .as_ref()?
221 .get(&TypeId::of::<T>())?
222 .downcast_ref()
223 }
224
225 /// Returns a copy of the context with the new value included.
226 ///
227 /// # Examples
228 ///
229 /// ```
230 /// use opentelemetry::Context;
231 ///
232 /// // Given some value types defined in your application
233 /// #[derive(Debug, PartialEq)]
234 /// struct ValueA(&'static str);
235 /// #[derive(Debug, PartialEq)]
236 /// struct ValueB(u64);
237 ///
238 /// // You can create a context with the first value set to "a"
239 /// let cx_with_a = Context::new().with_value(ValueA("a"));
240 ///
241 /// // And create another context based on the fist with a new value
242 /// let cx_with_a_and_b = cx_with_a.with_value(ValueB(42));
243 ///
244 /// // The first context is still available and unmodified
245 /// assert_eq!(cx_with_a.get::<ValueA>(), Some(&ValueA("a")));
246 /// assert_eq!(cx_with_a.get::<ValueB>(), None);
247 ///
248 /// // The second context now contains both values
249 /// assert_eq!(cx_with_a_and_b.get::<ValueA>(), Some(&ValueA("a")));
250 /// assert_eq!(cx_with_a_and_b.get::<ValueB>(), Some(&ValueB(42)));
251 /// ```
252 pub fn with_value<T: 'static + Send + Sync>(&self, value: T) -> Self {
253 let entries = if let Some(current_entries) = &self.entries {
254 let mut inner_entries = (**current_entries).clone();
255 inner_entries.insert(TypeId::of::<T>(), Arc::new(value));
256 Some(Arc::new(inner_entries))
257 } else {
258 let mut entries = EntryMap::default();
259 entries.insert(TypeId::of::<T>(), Arc::new(value));
260 Some(Arc::new(entries))
261 };
262 Context {
263 entries,
264 #[cfg(feature = "trace")]
265 span: self.span.clone(),
266 suppress_telemetry: self.suppress_telemetry,
267 }
268 }
269
270 /// Replaces the current context on this thread with this context.
271 ///
272 /// Dropping the returned [`ContextGuard`] will reset the current context to the
273 /// previous value.
274 ///
275 ///
276 /// # Examples
277 ///
278 /// ```
279 /// use opentelemetry::Context;
280 ///
281 /// #[derive(Debug, PartialEq)]
282 /// struct ValueA(&'static str);
283 ///
284 /// let my_cx = Context::new().with_value(ValueA("a"));
285 ///
286 /// // Set the current thread context
287 /// let cx_guard = my_cx.attach();
288 /// assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA("a")));
289 ///
290 /// // Drop the guard to restore the previous context
291 /// drop(cx_guard);
292 /// assert_eq!(Context::current().get::<ValueA>(), None);
293 /// ```
294 ///
295 /// Guards do not need to be explicitly dropped:
296 ///
297 /// ```
298 /// use opentelemetry::Context;
299 ///
300 /// #[derive(Debug, PartialEq)]
301 /// struct ValueA(&'static str);
302 ///
303 /// fn my_function() -> String {
304 /// // attach a context the duration of this function.
305 /// let my_cx = Context::new().with_value(ValueA("a"));
306 /// // NOTE: a variable name after the underscore is **required** or rust
307 /// // will drop the guard, restoring the previous context _immediately_.
308 /// let _guard = my_cx.attach();
309 ///
310 /// // anything happening in functions we call can still access my_cx...
311 /// my_other_function();
312 ///
313 /// // returning from the function drops the guard, exiting the span.
314 /// return "Hello world".to_owned();
315 /// }
316 ///
317 /// fn my_other_function() {
318 /// // ...
319 /// }
320 /// ```
321 /// Sub-scopes may be created to limit the duration for which the span is
322 /// entered:
323 ///
324 /// ```
325 /// use opentelemetry::Context;
326 ///
327 /// #[derive(Debug, PartialEq)]
328 /// struct ValueA(&'static str);
329 ///
330 /// let my_cx = Context::new().with_value(ValueA("a"));
331 ///
332 /// {
333 /// let _guard = my_cx.attach();
334 ///
335 /// // the current context can access variables in
336 /// assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA("a")));
337 ///
338 /// // exiting the scope drops the guard, detaching the context.
339 /// }
340 ///
341 /// // this is back in the default empty context
342 /// assert_eq!(Context::current().get::<ValueA>(), None);
343 /// ```
344 pub fn attach(self) -> ContextGuard {
345 let cx_id = CURRENT_CONTEXT.with(|cx| cx.borrow_mut().push(self));
346
347 ContextGuard {
348 cx_pos: cx_id,
349 _marker: PhantomData,
350 }
351 }
352
353 /// Returns whether telemetry is suppressed in this context.
354 #[inline]
355 pub fn is_telemetry_suppressed(&self) -> bool {
356 self.suppress_telemetry
357 }
358
359 /// Returns a new context with telemetry suppression enabled.
360 pub fn with_telemetry_suppressed(&self) -> Self {
361 Context {
362 entries: self.entries.clone(),
363 #[cfg(feature = "trace")]
364 span: self.span.clone(),
365 suppress_telemetry: true,
366 }
367 }
368
369 /// Enters a scope where telemetry is suppressed.
370 ///
371 /// This method is specifically designed for OpenTelemetry components (like Exporters,
372 /// Processors etc.) to prevent generating recursive or self-referential
373 /// telemetry data when performing their own operations.
374 ///
375 /// Without suppression, we have a telemetry-induced-telemetry situation
376 /// where, operations like exporting telemetry could generate new telemetry
377 /// about the export process itself, potentially causing:
378 /// - Infinite telemetry feedback loops
379 /// - Excessive resource consumption
380 ///
381 /// This method:
382 /// 1. Takes the current context
383 /// 2. Creates a new context from current, with `suppress_telemetry` set to `true`
384 /// 3. Attaches it to the current thread
385 /// 4. Returns a guard that restores the previous context when dropped
386 ///
387 /// OTel SDK components would check `is_current_telemetry_suppressed()` before
388 /// generating new telemetry, but not end users.
389 ///
390 /// # Examples
391 ///
392 /// ```
393 /// use opentelemetry::Context;
394 ///
395 /// // Example: Inside an exporter's implementation
396 /// fn example_export_function() {
397 /// // Prevent telemetry-generating operations from creating more telemetry
398 /// let _guard = Context::enter_telemetry_suppressed_scope();
399 ///
400 /// // Verify suppression is active
401 /// assert_eq!(Context::is_current_telemetry_suppressed(), true);
402 ///
403 /// // Here you would normally perform operations that might generate telemetry
404 /// // but now they won't because the context has suppression enabled
405 /// }
406 ///
407 /// // Demonstrate the function
408 /// example_export_function();
409 /// ```
410 pub fn enter_telemetry_suppressed_scope() -> ContextGuard {
411 Self::map_current(|cx| cx.with_telemetry_suppressed()).attach()
412 }
413
414 /// Returns whether telemetry is suppressed in the current context.
415 ///
416 /// This method is used by OpenTelemetry components to determine whether they should
417 /// generate new telemetry in the current execution context. It provides a performant
418 /// way to check the suppression state.
419 ///
420 /// End-users generally should not use this method directly, as it is primarily intended for
421 /// OpenTelemetry SDK components.
422 ///
423 ///
424 #[inline]
425 pub fn is_current_telemetry_suppressed() -> bool {
426 Self::map_current(|cx| cx.is_telemetry_suppressed())
427 }
428
429 #[cfg(feature = "trace")]
430 pub(crate) fn current_with_synchronized_span(value: SynchronizedSpan) -> Self {
431 Self::map_current(|cx| Context {
432 span: Some(Arc::new(value)),
433 entries: cx.entries.clone(),
434 suppress_telemetry: cx.suppress_telemetry,
435 })
436 }
437
438 #[cfg(feature = "trace")]
439 pub(crate) fn with_synchronized_span(&self, value: SynchronizedSpan) -> Self {
440 Context {
441 span: Some(Arc::new(value)),
442 entries: self.entries.clone(),
443 suppress_telemetry: self.suppress_telemetry,
444 }
445 }
446}
447
448impl fmt::Debug for Context {
449 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
450 let mut dbg = f.debug_struct("Context");
451
452 #[cfg(feature = "trace")]
453 let mut entries = self.entries.as_ref().map_or(0, |e| e.len());
454 #[cfg(feature = "trace")]
455 {
456 if let Some(span) = &self.span {
457 dbg.field("span", &span.span_context());
458 entries += 1;
459 } else {
460 dbg.field("span", &"None");
461 }
462 }
463 #[cfg(not(feature = "trace"))]
464 let entries = self.entries.as_ref().map_or(0, |e| e.len());
465
466 dbg.field("entries count", &entries)
467 .field("suppress_telemetry", &self.suppress_telemetry)
468 .finish()
469 }
470}
471
472/// A guard that resets the current context to the prior context when dropped.
473#[derive(Debug)]
474pub struct ContextGuard {
475 // The position of the context in the stack. This is used to pop the context.
476 cx_pos: u16,
477 // Ensure this type is !Send as it relies on thread locals
478 _marker: PhantomData<*const ()>,
479}
480
481impl Drop for ContextGuard {
482 fn drop(&mut self) {
483 let id = self.cx_pos;
484 if id > ContextStack::BASE_POS && id < ContextStack::MAX_POS {
485 // Extract only the span to drop outside of borrow_mut to avoid panic
486 // when the span's drop implementation calls Context::current()
487 #[cfg(feature = "trace")]
488 let _to_drop =
489 CURRENT_CONTEXT.with(|context_stack| context_stack.borrow_mut().pop_id(id));
490 #[cfg(not(feature = "trace"))]
491 CURRENT_CONTEXT.with(|context_stack| context_stack.borrow_mut().pop_id(id));
492 // Span (if any) is automatically dropped here, outside of borrow_mut scope
493 }
494 }
495}
496
497/// With TypeIds as keys, there's no need to hash them. They are already hashes
498/// themselves, coming from the compiler. The IdHasher holds the u64 of
499/// the TypeId, and then returns it, instead of doing any bit fiddling.
500#[derive(Clone, Default, Debug)]
501struct IdHasher(u64);
502
503impl Hasher for IdHasher {
504 fn write(&mut self, _: &[u8]) {
505 unreachable!("TypeId calls write_u64");
506 }
507
508 #[inline]
509 fn write_u64(&mut self, id: u64) {
510 self.0 = id;
511 }
512
513 #[inline]
514 fn finish(&self) -> u64 {
515 self.0
516 }
517}
518
519/// A stack for keeping track of the [`Context`] instances that have been attached
520/// to a thread.
521///
522/// The stack allows for popping of contexts by position, which is used to do out
523/// of order dropping of [`ContextGuard`] instances. Only when the top of the
524/// stack is popped, the topmost [`Context`] is actually restored.
525///
526/// The stack relies on the fact that it is thread local and that the
527/// [`ContextGuard`] instances that are constructed using ids from it can't be
528/// moved to other threads. That means that the ids are always valid and that
529/// they are always within the bounds of the stack.
530struct ContextStack {
531 /// This is the current [`Context`] that is active on this thread, and the top
532 /// of the [`ContextStack`]. It is always present, and if the `stack` is empty
533 /// it's an empty [`Context`].
534 ///
535 /// Having this here allows for fast access to the current [`Context`].
536 current_cx: Context,
537 /// A `stack` of the other contexts that have been attached to the thread.
538 stack: Vec<Option<Context>>,
539 /// Ensure this type is !Send as it relies on thread locals
540 _marker: PhantomData<*const ()>,
541}
542
543// Type alias for what pop_id returns - only return the span when trace feature is enabled
544#[cfg(feature = "trace")]
545type PopIdReturn = Option<Arc<SynchronizedSpan>>;
546#[cfg(not(feature = "trace"))]
547type PopIdReturn = ();
548
549impl ContextStack {
550 const BASE_POS: u16 = 0;
551 const MAX_POS: u16 = u16::MAX;
552 const INITIAL_CAPACITY: usize = 8;
553
554 #[inline(always)]
555 fn push(&mut self, cx: Context) -> u16 {
556 // The next id is the length of the `stack`, plus one since we have the
557 // top of the [`ContextStack`] as the `current_cx`.
558 let next_id = self.stack.len() + 1;
559 if next_id < ContextStack::MAX_POS.into() {
560 let current_cx = std::mem::replace(&mut self.current_cx, cx);
561 self.stack.push(Some(current_cx));
562 next_id as u16
563 } else {
564 // This is an overflow, log it and ignore it.
565 otel_warn!(
566 name: "Context.AttachFailed",
567 message = format!("Too many contexts. Max limit is {}. \
568 Context::current() remains unchanged as this attach failed. \
569 Dropping the returned ContextGuard will have no impact on Context::current().",
570 ContextStack::MAX_POS)
571 );
572 ContextStack::MAX_POS
573 }
574 }
575
576 #[inline(always)]
577 fn pop_id(&mut self, pos: u16) -> PopIdReturn {
578 if pos == ContextStack::BASE_POS || pos == ContextStack::MAX_POS {
579 // The empty context is always at the bottom of the [`ContextStack`]
580 // and cannot be popped, and the overflow position is invalid, so do
581 // nothing.
582 otel_warn!(
583 name: "Context.OutOfOrderDrop",
584 position = pos,
585 message = if pos == ContextStack::BASE_POS {
586 "Attempted to pop the base context which is not allowed"
587 } else {
588 "Attempted to pop the overflow position which is not allowed"
589 }
590 );
591 #[cfg(feature = "trace")]
592 return None;
593 }
594 let len: u16 = self.stack.len() as u16;
595 // Are we at the top of the [`ContextStack`]?
596 if pos == len {
597 // Shrink the stack if possible to clear out any out of order pops.
598 while let Some(None) = self.stack.last() {
599 _ = self.stack.pop();
600 }
601 // Restore the previous context. This will always happen since the
602 // empty context is always at the bottom of the stack if the
603 // [`ContextStack`] is not empty.
604 if let Some(Some(next_cx)) = self.stack.pop() {
605 // Extract and return only the span to avoid cloning the entire Context
606 #[cfg(feature = "trace")]
607 {
608 let old_cx = std::mem::replace(&mut self.current_cx, next_cx);
609 return old_cx.span;
610 }
611 #[cfg(not(feature = "trace"))]
612 {
613 self.current_cx = next_cx;
614 }
615 }
616 #[cfg(feature = "trace")]
617 return None;
618 } else {
619 // This is an out of order pop.
620 if pos >= len {
621 // This is an invalid id, ignore it.
622 otel_warn!(
623 name: "Context.PopOutOfBounds",
624 position = pos,
625 stack_length = len,
626 message = "Attempted to pop beyond the end of the context stack"
627 );
628 #[cfg(feature = "trace")]
629 return None;
630 }
631 // Clear out the entry at the given id and extract its span
632 #[cfg(feature = "trace")]
633 return self.stack[pos as usize].take().and_then(|cx| cx.span);
634 #[cfg(not(feature = "trace"))]
635 {
636 self.stack[pos as usize] = None;
637 }
638 }
639 }
640
641 #[inline(always)]
642 fn map_current_cx<T>(&self, f: impl FnOnce(&Context) -> T) -> T {
643 f(&self.current_cx)
644 }
645}
646
647impl Default for ContextStack {
648 fn default() -> Self {
649 ContextStack {
650 current_cx: Context::default(),
651 stack: Vec::with_capacity(ContextStack::INITIAL_CAPACITY),
652 _marker: PhantomData,
653 }
654 }
655}
656
657#[cfg(test)]
658mod tests {
659 use super::*;
660 use std::time::Duration;
661 use tokio::time::sleep;
662
663 #[derive(Debug, PartialEq)]
664 struct ValueA(u64);
665 #[derive(Debug, PartialEq)]
666 struct ValueB(u64);
667
668 #[test]
669 fn context_immutable() {
670 // start with Current, which should be an empty context
671 let cx = Context::current();
672 assert_eq!(cx.get::<ValueA>(), None);
673 assert_eq!(cx.get::<ValueB>(), None);
674
675 // with_value should return a new context,
676 // leaving the original context unchanged
677 let cx_new = cx.with_value(ValueA(1));
678
679 // cx should be unchanged
680 assert_eq!(cx.get::<ValueA>(), None);
681 assert_eq!(cx.get::<ValueB>(), None);
682
683 // cx_new should contain the new value
684 assert_eq!(cx_new.get::<ValueA>(), Some(&ValueA(1)));
685
686 // cx_new should be unchanged
687 let cx_newer = cx_new.with_value(ValueB(1));
688
689 // Cx and cx_new are unchanged
690 assert_eq!(cx.get::<ValueA>(), None);
691 assert_eq!(cx.get::<ValueB>(), None);
692 assert_eq!(cx_new.get::<ValueA>(), Some(&ValueA(1)));
693 assert_eq!(cx_new.get::<ValueB>(), None);
694
695 // cx_newer should contain both values
696 assert_eq!(cx_newer.get::<ValueA>(), Some(&ValueA(1)));
697 assert_eq!(cx_newer.get::<ValueB>(), Some(&ValueB(1)));
698 }
699
700 #[test]
701 fn nested_contexts() {
702 let _outer_guard = Context::new().with_value(ValueA(1)).attach();
703
704 // Only value `a` is set
705 let current = Context::current();
706 assert_eq!(current.get(), Some(&ValueA(1)));
707 assert_eq!(current.get::<ValueB>(), None);
708
709 {
710 let _inner_guard = Context::current_with_value(ValueB(42)).attach();
711 // Both values are set in inner context
712 let current = Context::current();
713 assert_eq!(current.get(), Some(&ValueA(1)));
714 assert_eq!(current.get(), Some(&ValueB(42)));
715
716 assert!(Context::map_current(|cx| {
717 assert_eq!(cx.get(), Some(&ValueA(1)));
718 assert_eq!(cx.get(), Some(&ValueB(42)));
719 true
720 }));
721 }
722
723 // Resets to only value `a` when inner guard is dropped
724 let current = Context::current();
725 assert_eq!(current.get(), Some(&ValueA(1)));
726 assert_eq!(current.get::<ValueB>(), None);
727
728 assert!(Context::map_current(|cx| {
729 assert_eq!(cx.get(), Some(&ValueA(1)));
730 assert_eq!(cx.get::<ValueB>(), None);
731 true
732 }));
733 }
734
735 #[test]
736 fn overlapping_contexts() {
737 let outer_guard = Context::new().with_value(ValueA(1)).attach();
738
739 // Only value `a` is set
740 let current = Context::current();
741 assert_eq!(current.get(), Some(&ValueA(1)));
742 assert_eq!(current.get::<ValueB>(), None);
743
744 let inner_guard = Context::current_with_value(ValueB(42)).attach();
745 // Both values are set in inner context
746 let current = Context::current();
747 assert_eq!(current.get(), Some(&ValueA(1)));
748 assert_eq!(current.get(), Some(&ValueB(42)));
749
750 assert!(Context::map_current(|cx| {
751 assert_eq!(cx.get(), Some(&ValueA(1)));
752 assert_eq!(cx.get(), Some(&ValueB(42)));
753 true
754 }));
755
756 drop(outer_guard);
757
758 // `inner_guard` is still alive so both `ValueA` and `ValueB` should still be accessible
759 let current = Context::current();
760 assert_eq!(current.get(), Some(&ValueA(1)));
761 assert_eq!(current.get(), Some(&ValueB(42)));
762
763 drop(inner_guard);
764
765 // Both guards are dropped and neither value should be accessible.
766 let current = Context::current();
767 assert_eq!(current.get::<ValueA>(), None);
768 assert_eq!(current.get::<ValueB>(), None);
769 }
770
771 #[test]
772 fn too_many_contexts() {
773 let mut guards: Vec<ContextGuard> = Vec::with_capacity(ContextStack::MAX_POS as usize);
774 let stack_max_pos = ContextStack::MAX_POS as u64;
775 // Fill the stack up until the last position
776 for i in 1..stack_max_pos {
777 let cx_guard = Context::current().with_value(ValueB(i)).attach();
778 assert_eq!(Context::current().get(), Some(&ValueB(i)));
779 assert_eq!(cx_guard.cx_pos, i as u16);
780 guards.push(cx_guard);
781 }
782 // Let's overflow the stack a couple of times
783 for _ in 0..16 {
784 let cx_guard = Context::current().with_value(ValueA(1)).attach();
785 assert_eq!(cx_guard.cx_pos, ContextStack::MAX_POS);
786 assert_eq!(Context::current().get::<ValueA>(), None);
787 assert_eq!(Context::current().get(), Some(&ValueB(stack_max_pos - 1)));
788 guards.push(cx_guard);
789 }
790 // Drop the overflow contexts
791 for _ in 0..16 {
792 guards.pop();
793 assert_eq!(Context::current().get::<ValueA>(), None);
794 assert_eq!(Context::current().get(), Some(&ValueB(stack_max_pos - 1)));
795 }
796 // Drop one more so we can add a new one
797 guards.pop();
798 assert_eq!(Context::current().get::<ValueA>(), None);
799 assert_eq!(Context::current().get(), Some(&ValueB(stack_max_pos - 2)));
800 // Push a new context and see that it works
801 let cx_guard = Context::current().with_value(ValueA(2)).attach();
802 assert_eq!(cx_guard.cx_pos, ContextStack::MAX_POS - 1);
803 assert_eq!(Context::current().get(), Some(&ValueA(2)));
804 assert_eq!(Context::current().get(), Some(&ValueB(stack_max_pos - 2)));
805 guards.push(cx_guard);
806 // Let's overflow the stack a couple of times again
807 for _ in 0..16 {
808 let cx_guard = Context::current().with_value(ValueA(1)).attach();
809 assert_eq!(cx_guard.cx_pos, ContextStack::MAX_POS);
810 assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA(2)));
811 assert_eq!(Context::current().get(), Some(&ValueB(stack_max_pos - 2)));
812 guards.push(cx_guard);
813 }
814 }
815
816 /// Tests that a new ContextStack is created with the correct initial capacity.
817 #[test]
818 fn test_initial_capacity() {
819 let stack = ContextStack::default();
820 assert_eq!(stack.stack.capacity(), ContextStack::INITIAL_CAPACITY);
821 }
822
823 /// Tests that map_current_cx correctly accesses the current context.
824 #[test]
825 fn test_map_current_cx() {
826 let mut stack = ContextStack::default();
827 let test_value = ValueA(42);
828 stack.current_cx = Context::new().with_value(test_value);
829
830 let result = stack.map_current_cx(|cx| {
831 assert_eq!(cx.get::<ValueA>(), Some(&ValueA(42)));
832 true
833 });
834 assert!(result);
835 }
836
837 /// Tests popping contexts in non-sequential order.
838 #[test]
839 fn test_pop_id_out_of_order() {
840 let mut stack = ContextStack::default();
841
842 // Push three contexts
843 let cx1 = Context::new().with_value(ValueA(1));
844 let cx2 = Context::new().with_value(ValueA(2));
845 let cx3 = Context::new().with_value(ValueA(3));
846
847 let id1 = stack.push(cx1);
848 let id2 = stack.push(cx2);
849 let id3 = stack.push(cx3);
850
851 // Pop middle context first - should not affect current context
852 stack.pop_id(id2);
853 assert_eq!(stack.current_cx.get::<ValueA>(), Some(&ValueA(3)));
854 assert_eq!(stack.stack.len(), 3); // Length unchanged for middle pops
855
856 // Pop last context - should restore previous valid context
857 stack.pop_id(id3);
858 assert_eq!(stack.current_cx.get::<ValueA>(), Some(&ValueA(1)));
859 assert_eq!(stack.stack.len(), 1);
860
861 // Pop first context - should restore to empty state
862 stack.pop_id(id1);
863 assert_eq!(stack.current_cx.get::<ValueA>(), None);
864 assert_eq!(stack.stack.len(), 0);
865 }
866
867 /// Tests edge cases in context stack operations. IRL these should log
868 /// warnings, and definitely not panic.
869 #[test]
870 fn test_pop_id_edge_cases() {
871 let mut stack = ContextStack::default();
872
873 // Test popping BASE_POS - should be no-op
874 stack.pop_id(ContextStack::BASE_POS);
875 assert_eq!(stack.stack.len(), 0);
876
877 // Test popping MAX_POS - should be no-op
878 stack.pop_id(ContextStack::MAX_POS);
879 assert_eq!(stack.stack.len(), 0);
880
881 // Test popping invalid position - should be no-op
882 stack.pop_id(1000);
883 assert_eq!(stack.stack.len(), 0);
884
885 // Test popping from empty stack - should be safe
886 stack.pop_id(1);
887 assert_eq!(stack.stack.len(), 0);
888 }
889
890 /// Tests stack behavior when reaching maximum capacity.
891 /// Once we push beyond this point, we should end up with a context
892 /// that points _somewhere_, but mutating it should not affect the current
893 /// active context.
894 #[test]
895 fn test_push_overflow() {
896 let mut stack = ContextStack::default();
897 let max_pos = ContextStack::MAX_POS as usize;
898
899 // Fill stack up to max position
900 for i in 0..max_pos {
901 let cx = Context::new().with_value(ValueA(i as u64));
902 let id = stack.push(cx);
903 assert_eq!(id, (i + 1) as u16);
904 }
905
906 // Try to push beyond capacity
907 let cx = Context::new().with_value(ValueA(max_pos as u64));
908 let id = stack.push(cx);
909 assert_eq!(id, ContextStack::MAX_POS);
910
911 // Verify current context remains unchanged after overflow
912 assert_eq!(
913 stack.current_cx.get::<ValueA>(),
914 Some(&ValueA((max_pos - 2) as u64))
915 );
916 }
917
918 /// Tests that:
919 /// 1. Parent context values are properly propagated to async operations
920 /// 2. Values added during async operations do not affect parent context
921 #[tokio::test]
922 async fn test_async_context_propagation() {
923 // A nested async operation we'll use to test propagation
924 async fn nested_operation() {
925 // Verify we can see the parent context's value
926 assert_eq!(
927 Context::current().get::<ValueA>(),
928 Some(&ValueA(42)),
929 "Parent context value should be available in async operation"
930 );
931
932 // Create new context
933 let cx_with_both = Context::current()
934 .with_value(ValueA(43)) // override ValueA
935 .with_value(ValueB(24)); // Add new ValueB
936
937 // Run nested async operation with both values
938 async {
939 // Verify both values are available
940 assert_eq!(
941 Context::current().get::<ValueA>(),
942 Some(&ValueA(43)),
943 "Parent value should still be available after adding new value"
944 );
945 assert_eq!(
946 Context::current().get::<ValueB>(),
947 Some(&ValueB(24)),
948 "New value should be available in async operation"
949 );
950
951 // Do some async work to simulate real-world scenario
952 sleep(Duration::from_millis(10)).await;
953
954 // Values should still be available after async work
955 assert_eq!(
956 Context::current().get::<ValueA>(),
957 Some(&ValueA(43)),
958 "Parent value should persist across await points"
959 );
960 assert_eq!(
961 Context::current().get::<ValueB>(),
962 Some(&ValueB(24)),
963 "New value should persist across await points"
964 );
965 }
966 .with_context(cx_with_both)
967 .await;
968 }
969
970 // Set up initial context with ValueA
971 let parent_cx = Context::new().with_value(ValueA(42));
972
973 // Create and run async operation with the parent context explicitly propagated
974 nested_operation().with_context(parent_cx.clone()).await;
975
976 // After async operation completes:
977 // 1. Parent context should be unchanged
978 assert_eq!(
979 parent_cx.get::<ValueA>(),
980 Some(&ValueA(42)),
981 "Parent context should be unchanged"
982 );
983 assert_eq!(
984 parent_cx.get::<ValueB>(),
985 None,
986 "Parent context should not see values added in async operation"
987 );
988
989 // 2. Current context should be back to default
990 assert_eq!(
991 Context::current().get::<ValueA>(),
992 None,
993 "Current context should be back to default"
994 );
995 assert_eq!(
996 Context::current().get::<ValueB>(),
997 None,
998 "Current context should not have async operation's values"
999 );
1000 }
1001
1002 ///
1003 /// Tests that unnatural parent->child relationships in nested async
1004 /// operations behave properly.
1005 ///
1006 #[tokio::test]
1007 async fn test_out_of_order_context_detachment_futures() {
1008 // This function returns a future, but doesn't await it
1009 // It will complete before the future that it creates.
1010 async fn create_a_future() -> impl std::future::Future<Output = ()> {
1011 // Create a future that will do some work, referencing our current
1012 // context, but don't await it.
1013 async {
1014 assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA(42)));
1015
1016 // Longer work
1017 sleep(Duration::from_millis(50)).await;
1018 }
1019 .with_context(Context::current())
1020 }
1021
1022 // Create our base context
1023 let parent_cx = Context::new().with_value(ValueA(42));
1024
1025 // await our nested function, which will create and detach a context
1026 let future = create_a_future().with_context(parent_cx).await;
1027
1028 // Execute the future. The future that created it is long gone, but this shouldn't
1029 // cause issues.
1030 let _a = future.await;
1031
1032 // Nothing terrible (e.g., panics!) should happen, and we should definitely not have any
1033 // values attached to our current context that were set in the nested operations.
1034 assert_eq!(Context::current().get::<ValueA>(), None);
1035 assert_eq!(Context::current().get::<ValueB>(), None);
1036 }
1037
1038 #[test]
1039 fn test_is_telemetry_suppressed() {
1040 // Default context has suppression disabled
1041 let cx = Context::new();
1042 assert!(!cx.is_telemetry_suppressed());
1043
1044 // With suppression enabled
1045 let suppressed = cx.with_telemetry_suppressed();
1046 assert!(suppressed.is_telemetry_suppressed());
1047 }
1048
1049 #[test]
1050 fn test_with_telemetry_suppressed() {
1051 // Start with a normal context
1052 let cx = Context::new();
1053 assert!(!cx.is_telemetry_suppressed());
1054
1055 // Create a suppressed context
1056 let suppressed = cx.with_telemetry_suppressed();
1057
1058 // Original should remain unchanged
1059 assert!(!cx.is_telemetry_suppressed());
1060
1061 // New context should be suppressed
1062 assert!(suppressed.is_telemetry_suppressed());
1063
1064 // Test with values to ensure they're preserved
1065 let cx_with_value = cx.with_value(ValueA(42));
1066 let suppressed_with_value = cx_with_value.with_telemetry_suppressed();
1067
1068 assert!(!cx_with_value.is_telemetry_suppressed());
1069 assert!(suppressed_with_value.is_telemetry_suppressed());
1070 assert_eq!(suppressed_with_value.get::<ValueA>(), Some(&ValueA(42)));
1071 }
1072
1073 #[test]
1074 fn test_enter_telemetry_suppressed_scope() {
1075 // Ensure we start with a clean context
1076 let _reset_guard = Context::new().attach();
1077
1078 // Default context should not be suppressed
1079 assert!(!Context::is_current_telemetry_suppressed());
1080
1081 // Add an entry to the current context
1082 let cx_with_value = Context::current().with_value(ValueA(42));
1083 let _guard_with_value = cx_with_value.attach();
1084
1085 // Verify the entry is present and context is not suppressed
1086 assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA(42)));
1087 assert!(!Context::is_current_telemetry_suppressed());
1088
1089 // Enter a suppressed scope
1090 {
1091 let _guard = Context::enter_telemetry_suppressed_scope();
1092
1093 // Verify suppression is active and the entry is still present
1094 assert!(Context::is_current_telemetry_suppressed());
1095 assert!(Context::current().is_telemetry_suppressed());
1096 assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA(42)));
1097 }
1098
1099 // After guard is dropped, should be back to unsuppressed and entry should still be present
1100 assert!(!Context::is_current_telemetry_suppressed());
1101 assert!(!Context::current().is_telemetry_suppressed());
1102 assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA(42)));
1103 }
1104
1105 #[test]
1106 fn test_nested_suppression_scopes() {
1107 // Ensure we start with a clean context
1108 let _reset_guard = Context::new().attach();
1109
1110 // Default context should not be suppressed
1111 assert!(!Context::is_current_telemetry_suppressed());
1112
1113 // First level suppression
1114 {
1115 let _outer = Context::enter_telemetry_suppressed_scope();
1116 assert!(Context::is_current_telemetry_suppressed());
1117
1118 // Second level. This component is unaware of Suppression,
1119 // and just attaches a new context. Since it is from current,
1120 // it'll already have suppression enabled.
1121 {
1122 let _inner = Context::current().with_value(ValueA(1)).attach();
1123 assert!(Context::is_current_telemetry_suppressed());
1124 assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA(1)));
1125 }
1126
1127 // Another scenario. This component is unaware of Suppression,
1128 // and just attaches a new context, not from Current. Since it is
1129 // not from current it will not have suppression enabled.
1130 {
1131 let _inner = Context::new().with_value(ValueA(1)).attach();
1132 assert!(!Context::is_current_telemetry_suppressed());
1133 assert_eq!(Context::current().get::<ValueA>(), Some(&ValueA(1)));
1134 }
1135
1136 // Still suppressed after inner scope
1137 assert!(Context::is_current_telemetry_suppressed());
1138 }
1139
1140 // Back to unsuppressed
1141 assert!(!Context::is_current_telemetry_suppressed());
1142 }
1143
1144 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1145 async fn test_async_suppression() {
1146 async fn nested_operation() {
1147 assert!(Context::is_current_telemetry_suppressed());
1148
1149 let cx_with_additional_value = Context::current().with_value(ValueB(24));
1150
1151 async {
1152 assert_eq!(
1153 Context::current().get::<ValueB>(),
1154 Some(&ValueB(24)),
1155 "Parent value should still be available after adding new value"
1156 );
1157 assert!(Context::is_current_telemetry_suppressed());
1158
1159 // Do some async work to simulate real-world scenario
1160 sleep(Duration::from_millis(10)).await;
1161
1162 // Values should still be available after async work
1163 assert_eq!(
1164 Context::current().get::<ValueB>(),
1165 Some(&ValueB(24)),
1166 "Parent value should still be available after adding new value"
1167 );
1168 assert!(Context::is_current_telemetry_suppressed());
1169 }
1170 .with_context(cx_with_additional_value)
1171 .await;
1172 }
1173
1174 // Set up suppressed context, but don't attach it to current
1175 let suppressed_parent = Context::new().with_telemetry_suppressed();
1176 // Current should not be suppressed as we haven't attached it
1177 assert!(!Context::is_current_telemetry_suppressed());
1178
1179 // Create and run async operation with the suppressed context explicitly propagated
1180 nested_operation()
1181 .with_context(suppressed_parent.clone())
1182 .await;
1183
1184 // After async operation completes:
1185 // Suppression should be active
1186 assert!(suppressed_parent.is_telemetry_suppressed());
1187
1188 // Current should still be not suppressed
1189 assert!(!Context::is_current_telemetry_suppressed());
1190 }
1191}