diff --git a/opentelemetry/src/context.rs b/opentelemetry/src/context/context_store.rs similarity index 99% rename from opentelemetry/src/context.rs rename to opentelemetry/src/context/context_store.rs index 3707bd3758..809eec43d8 100644 --- a/opentelemetry/src/context.rs +++ b/opentelemetry/src/context/context_store.rs @@ -1,3 +1,4 @@ + use crate::otel_warn; #[cfg(feature = "trace")] use crate::trace::context::SynchronizedSpan; @@ -78,7 +79,7 @@ thread_local! { #[derive(Clone, Default)] pub struct Context { #[cfg(feature = "trace")] - pub(super) span: Option>, + pub(crate) span: Option>, entries: Option>, } @@ -314,7 +315,7 @@ impl Context { } #[cfg(feature = "trace")] - pub(super) fn current_with_synchronized_span(value: SynchronizedSpan) -> Self { + pub(crate) fn current_with_synchronized_span(value: SynchronizedSpan) -> Self { Context { span: Some(Arc::new(value)), entries: Context::map_current(|cx| cx.entries.clone()), @@ -322,7 +323,7 @@ impl Context { } #[cfg(feature = "trace")] - pub(super) fn with_synchronized_span(&self, value: SynchronizedSpan) -> Self { + pub(crate) fn with_synchronized_span(&self, value: SynchronizedSpan) -> Self { Context { span: Some(Arc::new(value)), entries: self.entries.clone(), diff --git a/opentelemetry/src/context/future_ext.rs b/opentelemetry/src/context/future_ext.rs new file mode 100644 index 0000000000..7c0831af34 --- /dev/null +++ b/opentelemetry/src/context/future_ext.rs @@ -0,0 +1,31 @@ +use crate::Context; +use crate::trace::WithContext; + +/// Extension trait allowing futures, streams, and sinks to be traced with a span. +pub trait FutureExt: Sized { + /// Attaches the provided [`Context`] to this type, returning a `WithContext` + /// wrapper. + /// + /// When the wrapped type is a future, stream, or sink, the attached context + /// will be set as current while it is being polled. + /// + /// [`Context`]: Context + fn with_context(self, otel_cx: Context) -> WithContext { + WithContext { + inner: self, + otel_cx, + } + } + + /// Attaches the current [`Context`] to this type, returning a `WithContext` + /// wrapper. + /// + /// When the wrapped type is a future, stream, or sink, the attached context + /// will be set as the default while it is being polled. + /// + /// [`Context`]: Context + fn with_current_context(self) -> WithContext { + let otel_cx = Context::current(); + self.with_context(otel_cx) + } +} diff --git a/opentelemetry/src/context/mod.rs b/opentelemetry/src/context/mod.rs new file mode 100644 index 0000000000..042c34d0d8 --- /dev/null +++ b/opentelemetry/src/context/mod.rs @@ -0,0 +1,6 @@ +mod context_store; +mod future_ext; + +pub use context_store::Context; +pub use context_store::ContextGuard; +pub use future_ext::FutureExt; \ No newline at end of file diff --git a/opentelemetry/src/trace/context.rs b/opentelemetry/src/trace/context.rs index 125c96f5f5..7db6404676 100644 --- a/opentelemetry/src/trace/context.rs +++ b/opentelemetry/src/trace/context.rs @@ -15,6 +15,9 @@ use std::{ task::{Context as TaskContext, Poll}, }; +// Re-export for compatability. This used to be contained here. +pub use crate::context::FutureExt; + const NOOP_SPAN: SynchronizedSpan = SynchronizedSpan { span_context: SpanContext::NONE, inner: None, @@ -377,8 +380,8 @@ pin_project! { #[derive(Clone, Debug)] pub struct WithContext { #[pin] - inner: T, - otel_cx: Context, + pub(crate) inner: T, + pub(crate) otel_cx: Context, } } @@ -445,31 +448,3 @@ where } } -/// Extension trait allowing futures, streams, and sinks to be traced with a span. -pub trait FutureExt: Sized { - /// Attaches the provided [`Context`] to this type, returning a `WithContext` - /// wrapper. - /// - /// When the wrapped type is a future, stream, or sink, the attached context - /// will be set as current while it is being polled. - /// - /// [`Context`]: crate::Context - fn with_context(self, otel_cx: Context) -> WithContext { - WithContext { - inner: self, - otel_cx, - } - } - - /// Attaches the current [`Context`] to this type, returning a `WithContext` - /// wrapper. - /// - /// When the wrapped type is a future, stream, or sink, the attached context - /// will be set as the default while it is being polled. - /// - /// [`Context`]: crate::Context - fn with_current_context(self) -> WithContext { - let otel_cx = Context::current(); - self.with_context(otel_cx) - } -}