chore: Move FutureExt into context (#2776)

Co-authored-by: Cijo Thomas <[email protected]>

Author: scottgerring

opentelemetry/src/context.rs

@@ -1,3 +1,14 @@
+//! Execution-scoped context propagation.
+//!
+//! The `context` module provides mechanisms for propagating values across API boundaries and between
+//! logically associated execution units. It enables cross-cutting concerns to access their data in-process
+//! using a shared context object.
+//!
+//! # Main Types
+//!
+//! - [`Context`]: An immutable, execution-scoped collection of values.
+//!
+
 use crate::otel_warn;
 #[cfg(feature = "trace")]
 use crate::trace::context::SynchronizedSpan;
@@ -9,6 +20,10 @@ use std::hash::{BuildHasherDefault, Hasher};
 use std::marker::PhantomData;
 use std::sync::Arc;
 
+mod future_ext;
+
+pub use future_ext::FutureExt;
+
 thread_local! {
     static CURRENT_CONTEXT: RefCell<ContextStack> = RefCell::new(ContextStack::default());
 }
@@ -78,7 +93,7 @@ thread_local! {
 #[derive(Clone, Default)]
 pub struct Context {
     #[cfg(feature = "trace")]
-    pub(super) span: Option<Arc<SynchronizedSpan>>,
+    pub(crate) span: Option<Arc<SynchronizedSpan>>,
     entries: Option<Arc<EntryMap>>,
 }
 
@@ -314,15 +329,15 @@ 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()),
         }
     }
 
     #[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(),

opentelemetry/src/context/future_ext.rs

@@ -0,0 +1,108 @@
+use crate::Context;
+use futures_core::Stream;
+use futures_sink::Sink;
+use pin_project_lite::pin_project;
+use std::pin::Pin;
+use std::task::Context as TaskContext;
+use std::task::Poll;
+impl<T: Sized> FutureExt for T {}
+
+impl<T: std::future::Future> std::future::Future for WithContext<T> {
+    type Output = T::Output;
+
+    fn poll(self: Pin<&mut Self>, task_cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let _guard = this.otel_cx.clone().attach();
+
+        this.inner.poll(task_cx)
+    }
+}
+
+impl<T: Stream> Stream for WithContext<T> {
+    type Item = T::Item;
+
+    fn poll_next(self: Pin<&mut Self>, task_cx: &mut TaskContext<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        let _guard = this.otel_cx.clone().attach();
+        T::poll_next(this.inner, task_cx)
+    }
+}
+
+pin_project! {
+    /// A future, stream, or sink that has an associated context.
+    #[derive(Clone, Debug)]
+    pub struct WithContext<T> {
+        #[pin]
+        inner: T,
+        otel_cx: Context,
+    }
+}
+
+impl<I, T: Sink<I>> Sink<I> for WithContext<T>
+where
+    T: Sink<I>,
+{
+    type Error = T::Error;
+
+    fn poll_ready(
+        self: Pin<&mut Self>,
+        task_cx: &mut TaskContext<'_>,
+    ) -> Poll<Result<(), Self::Error>> {
+        let this = self.project();
+        let _guard = this.otel_cx.clone().attach();
+        T::poll_ready(this.inner, task_cx)
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
+        let this = self.project();
+        let _guard = this.otel_cx.clone().attach();
+        T::start_send(this.inner, item)
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        task_cx: &mut TaskContext<'_>,
+    ) -> Poll<Result<(), Self::Error>> {
+        let this = self.project();
+        let _guard = this.otel_cx.clone().attach();
+        T::poll_flush(this.inner, task_cx)
+    }
+
+    fn poll_close(
+        self: Pin<&mut Self>,
+        task_cx: &mut TaskContext<'_>,
+    ) -> Poll<Result<(), Self::Error>> {
+        let this = self.project();
+        let _enter = this.otel_cx.clone().attach();
+        T::poll_close(this.inner, task_cx)
+    }
+}
+
+/// 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<Self> {
+        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<Self> {
+        let otel_cx = Context::current();
+        self.with_context(otel_cx)
+    }
+}

opentelemetry/src/lib.rs

@@ -246,7 +246,7 @@ pub mod global;
 
 pub mod baggage;
 
-mod context;
+pub mod context;
 
 pub use context::{Context, ContextGuard};
 

opentelemetry/src/trace/context.rs

@@ -4,16 +4,10 @@ use crate::{
     trace::{Span, SpanContext, Status},
     Context, ContextGuard, KeyValue,
 };
-use futures_core::stream::Stream;
-use futures_sink::Sink;
-use pin_project_lite::pin_project;
-use std::{
-    borrow::Cow,
-    error::Error,
-    pin::Pin,
-    sync::Mutex,
-    task::{Context as TaskContext, Poll},
-};
+use std::{borrow::Cow, error::Error, sync::Mutex};
+
+// Re-export for compatability. This used to be contained here.
+pub use crate::context::FutureExt;
 
 const NOOP_SPAN: SynchronizedSpan = SynchronizedSpan {
     span_context: SpanContext::NONE,
@@ -371,105 +365,3 @@ where
 {
     Context::map_current(|cx| f(cx.span()))
 }
-
-pin_project! {
-    /// A future, stream, or sink that has an associated context.
-    #[derive(Clone, Debug)]
-    pub struct WithContext<T> {
-        #[pin]
-        inner: T,
-        otel_cx: Context,
-    }
-}
-
-impl<T: Sized> FutureExt for T {}
-
-impl<T: std::future::Future> std::future::Future for WithContext<T> {
-    type Output = T::Output;
-
-    fn poll(self: Pin<&mut Self>, task_cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
-        let this = self.project();
-        let _guard = this.otel_cx.clone().attach();
-
-        this.inner.poll(task_cx)
-    }
-}
-
-impl<T: Stream> Stream for WithContext<T> {
-    type Item = T::Item;
-
-    fn poll_next(self: Pin<&mut Self>, task_cx: &mut TaskContext<'_>) -> Poll<Option<Self::Item>> {
-        let this = self.project();
-        let _guard = this.otel_cx.clone().attach();
-        T::poll_next(this.inner, task_cx)
-    }
-}
-
-impl<I, T: Sink<I>> Sink<I> for WithContext<T>
-where
-    T: Sink<I>,
-{
-    type Error = T::Error;
-
-    fn poll_ready(
-        self: Pin<&mut Self>,
-        task_cx: &mut TaskContext<'_>,
-    ) -> Poll<Result<(), Self::Error>> {
-        let this = self.project();
-        let _guard = this.otel_cx.clone().attach();
-        T::poll_ready(this.inner, task_cx)
-    }
-
-    fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
-        let this = self.project();
-        let _guard = this.otel_cx.clone().attach();
-        T::start_send(this.inner, item)
-    }
-
-    fn poll_flush(
-        self: Pin<&mut Self>,
-        task_cx: &mut TaskContext<'_>,
-    ) -> Poll<Result<(), Self::Error>> {
-        let this = self.project();
-        let _guard = this.otel_cx.clone().attach();
-        T::poll_flush(this.inner, task_cx)
-    }
-
-    fn poll_close(
-        self: Pin<&mut Self>,
-        task_cx: &mut TaskContext<'_>,
-    ) -> Poll<Result<(), Self::Error>> {
-        let this = self.project();
-        let _enter = this.otel_cx.clone().attach();
-        T::poll_close(this.inner, task_cx)
-    }
-}
-
-/// 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<Self> {
-        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<Self> {
-        let otel_cx = Context::current();
-        self.with_context(otel_cx)
-    }
-}

opentelemetry/src/trace/mod.rs

@@ -177,9 +177,7 @@ mod tracer;
 mod tracer_provider;
 
 pub use self::{
-    context::{
-        get_active_span, mark_span_as_active, FutureExt, SpanRef, TraceContextExt, WithContext,
-    },
+    context::{get_active_span, mark_span_as_active, FutureExt, SpanRef, TraceContextExt},
     span::{Span, SpanKind, Status},
     span_context::{SpanContext, TraceState},
     tracer::{SamplingDecision, SamplingResult, SpanBuilder, Tracer},