chore: Move FutureExt into context

opentelemetry/src/context.rs → 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<Arc<SynchronizedSpan>>,
+    pub(crate) span: Option<Arc<SynchronizedSpan>>,
     entries: Option<Arc<EntryMap>>,
 }
 
@@ -314,15 +315,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,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<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/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;

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<T> {
         #[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<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)
-    }
-}