Repository: cloudflare/foundations (language: rust)

Issue:
Change LogScope to use an _inner field instead of allow(dead_code)
https://github.com/cloudflare/foundations/pull/38


Discussion / hints:

Fix patch:
diff --git a/foundations-macros/src/span_fn.rs b/foundations-macros/src/span_fn.rs
index 38068e3..5f8649e 100644
--- a/foundations-macros/src/span_fn.rs
+++ b/foundations-macros/src/span_fn.rs
@@ -41,12 +41,19 @@ impl Parse for SpanName {
 struct Options {
     #[darling(default = "Options::default_crate_path")]
     crate_path: Path,
+
+    #[darling(default = "Options::default_async_local")]
+    async_local: bool,
 }
 
 impl Options {
     fn default_crate_path() -> Path {
         parse_quote!(::foundations)
     }
+
+    fn default_async_local() -> bool {
+        false
+    }
 }
 
 struct Args {
@@ -100,17 +107,12 @@ fn expand_from_parsed(args: Args, item_fn: ItemFn) -> TokenStream2 {
         block,
     } = item_fn;
 
-    let span_name = args.span_name.as_tokens();
-    let crate_path = &args.options.crate_path;
-
     let body = match asyncness {
-        Some(_) => quote!(
-            #crate_path::telemetry::TelemetryContext::current().apply_with_tracing_span(
-                #span_name,
-                async move { #block }
-            ).await
-        ),
+        Some(_) => wrap_with_span(&args, quote!(async move { #block })),
         None => try_async_trait_fn_rewrite(&args, &block).unwrap_or_else(|| {
+            let span_name = args.span_name.as_tokens();
+            let crate_path = &args.options.crate_path;
+
             quote!(
                 let __span = #crate_path::telemetry::tracing::span(#span_name);
                 #block
@@ -158,18 +160,12 @@ fn try_async_trait_fn_rewrite(args: &Args, body: &Block) -> Option<TokenStream2>
         .map(|stmt| stmt.to_token_stream())
         .collect();
 
-    let span_name = args.span_name.as_tokens();
-    let crate_path = &args.options.crate_path;
+    let async_block = wrap_with_span(args, async_block.into_token_stream());
 
     // NOTE: OXY-1023 we do instrumentation inside additional future, so boxed
     // future can capture telemetry context on poll if it was instrumented.
     *body_stmts_token_streams.last_mut().unwrap() = quote!(
-        Box::pin(async move {
-            #crate_path::telemetry::TelemetryContext::current().apply_with_tracing_span(
-                #span_name,
-                #async_block
-            ).await
-        })
+        Box::pin(async move { #async_block })
     );
 
     Some(quote!(
@@ -177,6 +173,24 @@ fn try_async_trait_fn_rewrite(args: &Args, body: &Block) -> Option<TokenStream2>
     ))
 }
 
+fn wrap_with_span(args: &Args, block: TokenStream2) -> TokenStream2 {
+    let apply_fn = if args.options.async_local {
+        quote!(apply_local)
+    } else {
+        quote!(apply)
+    };
+
+    let span_name = args.span_name.as_tokens();
+    let crate_path = &args.options.crate_path;
+
+    quote!(
+        #crate_path::telemetry::tracing::span(#span_name)
+            .into_context()
+            .#apply_fn(#block)
+            .await
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -261,14 +275,46 @@ mod tests {
 
         let expected = code_str! {
             async fn do_async<>() -> io::Result<String> {
-                ::foundations::telemetry::TelemetryContext::current().apply_with_tracing_span(
-                    "async_span",
-                    async move {{
+                ::foundations::telemetry::tracing::span("async_span")
+                    .into_context()
+                    .apply(async move {{
+                        do_something_else().await;
+
+                        Ok("foo".into())
+                    }})
+                    .await
+            }
+        };
+
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn expand_async_fn_local() {
+        let args = parse_attr! {
+            #[span_fn("async_span", async_local = true)]
+        };
+
+        let item_fn = parse_quote! {
+            async fn do_async() -> io::Result<String> {
+                do_something_else().await;
+
+                Ok("foo".into())
+            }
+        };
+
+        let actual = expand_from_parsed(args, item_fn).to_string();
+
+        let expected = code_str! {
+            async fn do_async<>() -> io::Result<String> {
+                ::foundations::telemetry::tracing::span("async_span")
+                    .into_context()
+                    .apply_local(async move {{
                         do_something_else().await;
 
                         Ok("foo".into())
-                    }}
-                ).await
+                    }})
+                    .await
             }
         };
 
@@ -319,9 +365,75 @@ mod tests {
                 Self: 'async_trait,
             {
                 Box::pin(async move {
-                    ::foundations::telemetry::TelemetryContext::current().apply_with_tracing_span(
-                        "async_trait_span",
-                        async move {
+                    ::foundations::telemetry::tracing::span("async_trait_span")
+                        .into_context()
+                        .apply(async move {
+                            if let ::core::option::Option::Some(__ret) = ::core::option::Option::None::<String> {
+                                return __ret;
+                            }
+                            let __self = self;
+                            let __ret: String = {
+                                __self.do_something_else().await;
+                                "foo".into()
+                            };
+                            #[allow(unreachable_code)]
+                            __ret
+                        })
+                        .await
+                })
+            }
+        };
+
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn expand_async_trait_fn_local() {
+        let args = parse_attr! {
+            #[span_fn("async_trait_span", async_local = true)]
+        };
+
+        let item_fn = parse_quote! {
+            fn test<'life0, 'async_trait>(
+                &'life0 self,
+            ) -> ::core::pin::Pin<
+                Box<dyn ::core::future::Future<Output = String> + ::core::marker::Send + 'async_trait>
+            >
+            where
+                'life0: 'async_trait,
+                Self: 'async_trait,
+            {
+                Box::pin(async move {
+                    if let ::core::option::Option::Some(__ret) = ::core::option::Option::None::<String> {
+                        return __ret;
+                    }
+                    let __self = self;
+                    let __ret: String = {
+                        __self.do_something_else().await;
+                        "foo".into()
+                    };
+                    #[allow(unreachable_code)]
+                    __ret
+                })
+            }
+        };
+
+        let actual = expand_from_parsed(args, item_fn).to_string();
+
+        let expected = code_str! {
+            fn test<'life0, 'async_trait>(
+                &'life0 self,
+            ) -> ::core::pin::Pin<
+                Box<dyn ::core::future::Future<Output = String> + ::core::marker::Send + 'async_trait>
+            >
+            where
+                'life0: 'async_trait,
+                Self: 'async_trait,
+            {
+                Box::pin(async move {
+                    ::foundations::telemetry::tracing::span("async_trait_span")
+                        .into_context()
+                        .apply_local(async move {
                             if let ::core::option::Option::Some(__ret) = ::core::option::Option::None::<String> {
                                 return __ret;
                             }
@@ -332,7 +444,8 @@ mod tests {
                             };
                             #[allow(unreachable_code)]
                             __ret
-                        }).await
+                        })
+                        .await
                 })
             }
         };
diff --git a/foundations/src/telemetry/log/internal.rs b/foundations/src/telemetry/log/internal.rs
index 1a057b7..6739bd4 100644
--- a/foundations/src/telemetry/log/internal.rs
+++ b/foundations/src/telemetry/log/internal.rs
@@ -57,12 +57,16 @@ impl Deref for LoggerWithKvNestingTracking {
 }
 
 #[must_use]
-pub(crate) struct LogScope(#[allow(dead_code)] Scope<SharedLog>);
+pub(crate) struct LogScope {
+    _inner: Scope<SharedLog>,
+}
 
 impl LogScope {
     #[inline]
     pub(crate) fn new(log: SharedLog) -> Self {
-        Self(Scope::new(&LogHarness::get().log_scope_stack, log))
+        Self {
+            _inner: Scope::new(&LogHarness::get().log_scope_stack, log),
+        }
     }
 }
 
diff --git a/foundations/src/telemetry/mod.rs b/foundations/src/telemetry/mod.rs
index 9a637a6..c8c0857 100644
--- a/foundations/src/telemetry/mod.rs
+++ b/foundations/src/telemetry/mod.rs
@@ -50,6 +50,7 @@
 mod scope;
 
 mod driver;
+mod telemetry_context;
 
 #[cfg(all(feature = "tracing", feature = "telemetry-otlp-grpc"))]
 mod otlp_conversion;
@@ -93,26 +94,18 @@ use self::settings::TelemetrySettings;
 use crate::utils::feature_use;
 use crate::{BootstrapResult, ServiceInfo};
 use futures_util::stream::FuturesUnordered;
-use std::future::Future;
-use std::pin::Pin;
-use std::task::{Context, Poll};
-
-feature_use!(cfg(feature = "logging"), {
-    use self::log::internal::{current_log, fork_log, LogScope, SharedLog};
-    use std::sync::Arc;
-});
 
 feature_use!(cfg(feature = "tracing"), {
-    use self::tracing::internal::{create_span, current_span, fork_trace, SharedSpan};
     use self::tracing::SpanScope;
-    use std::borrow::Cow;
 
     feature_use!(cfg(feature = "testing"), {
-        use self::tracing::internal::Tracer;
-        use self::tracing::testing::{current_test_tracer, TestTracerScope};
+        use self::tracing::testing::TestTracerScope;
     });
 });
 
+#[cfg(feature = "logging")]
+use self::log::internal::LogScope;
+
 #[cfg(feature = "testing")]
 pub use self::testing::TestTelemetryContext;
 
@@ -123,6 +116,9 @@ pub use self::memory_profiler::MemoryProfiler;
 pub use self::server::{TelemetryRouteHandler, TelemetryRouteHandlerFuture, TelemetryServerRoute};
 
 pub use self::driver::TelemetryDriver;
+pub use self::telemetry_context::{
+    TelemetryContext, WithTelemetryContext, WithTelemetryContextLocal,
+};
 
 /// A macro that enables telemetry testing in `#[test]` and `#[tokio::test]`.
 ///
@@ -207,25 +203,6 @@ pub use self::driver::TelemetryDriver;
 #[cfg(feature = "testing")]
 pub use foundations_macros::with_test_telemetry;
 
-/// Wrapper for a future that provides it with [`TelemetryContext`].
-pub struct WithTelemetryContext<'f, T> {
-    // NOTE: we intentionally erase type here as we can get close to the type
-    // length limit, adding telemetry wrappers on top causes compiler to fail in some
-    // cases.
-    inner: Pin<Box<dyn Future<Output = T> + Send + 'f>>,
-    ctx: TelemetryContext,
-}
-
-impl<'f, T> Future for WithTelemetryContext<'f, T> {
-    type Output = T;
-
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let _telemetry_scope = self.ctx.scope();
-
-        self.inner.as_mut().poll(cx)
-    }
-}
-
 /// A handle for the scope in which certain [`TelemetryContext`] is active.
 ///
 /// Scope ends when the handle is dropped.
@@ -246,425 +223,6 @@ pub struct TelemetryScope {
     _test_tracer_scope: Option<TestTracerScope>,
 }
 
-/// Implicit context for logging and tracing.
-///
-/// Current context can be obtained with the [`TelemetryContext::current`] method.
-///
-/// The context can be forked with [`TelemetryContext::with_forked_log`] and
-/// [`TelemetryContext::with_forked_trace`] methods. And propagated in different scopes with
-/// [`TelemetryContext::scope`] and [`TelemetryContext::apply`] methods.
-#[derive(Debug, Clone)]
-pub struct TelemetryContext {
-    #[cfg(feature = "logging")]
-    log: SharedLog,
-
-    // NOTE: we might not have tracing root at this point
-    #[cfg(feature = "tracing")]
-    span: Option<SharedSpan>,
-
-    #[cfg(all(feature = "tracing", feature = "testing"))]
-    test_tracer: Option<Tracer>,
-}
-
-impl TelemetryContext {
-    /// Returns the telemetry context that is active in the current scope.
-    pub fn current() -> Self {
-        Self {
-            #[cfg(feature = "logging")]
-            log: current_log(),
-
-            #[cfg(feature = "tracing")]
-            span: current_span(),
-
-            #[cfg(all(feature = "tracing", feature = "testing"))]
-            test_tracer: current_test_tracer(),
-        }
-    }
-
-    /// Creates a scope handle for the telemetry context.
-    ///
-    /// The telemetry context is active in the scope unless the handle is dropped.
-    ///
-    /// Note that the scope can only be used in the sync contexts. To propagate telemetry context
-    /// to async contexts, [`TelemetryContext::apply`] should be used instead.
-    ///
-    /// The scope handle is useful to propagate the telemetry context to callbacks provided
-    /// by third party libraries or other threads.
-    ///
-    /// The other use case is situations where code control flow doesn't align with telemetry flow.
-    /// E.g. there is a centralized dispatcher that drives certain tasks and tasks are registered
-    /// with the dispatcher via callbacks. It's desirable for each task to have its own telemetry
-    /// context, so scope can be used to propagate task context to the dispatcher's callbacks.
-    ///
-    /// # Examples
-    /// ```
-    /// use foundations::telemetry::TelemetryContext;
-    /// use foundations::telemetry::tracing::{self, test_trace};
-    ///
-    /// // Test context is used for demonstration purposes to show the resulting traces.
-    /// let ctx = TelemetryContext::test();
-    ///
-    /// {
-    ///     let _scope = ctx.scope();
-    ///     let _root = tracing::span("root");
-    ///     let telemetry_ctx = TelemetryContext::current();
-    ///
-    ///     let handle = std::thread::spawn(move || {
-    ///         let _scope = telemetry_ctx.scope();
-    ///         let _child = tracing::span("child");
-    ///     });
-    ///
-    ///     handle.join();
-    /// }
-    ///
-    /// assert_eq!(
-    ///     ctx.traces(Default::default()),
-    ///     vec![
-    ///         test_trace! {
-    ///             "root" => {
-    ///                 "child"
-    ///             }
-    ///         },
-    ///     ]
-    /// );
-    /// ```
-    pub fn scope(&self) -> TelemetryScope {
-        TelemetryScope {
-            #[cfg(feature = "logging")]
-            _log_scope: LogScope::new(Arc::clone(&self.log)),
-
-            #[cfg(feature = "tracing")]
-            _span_scope: self.span.as_ref().cloned().map(SpanScope::new),
-
-            #[cfg(all(feature = "tracing", feature = "testing"))]
-            _test_tracer_scope: self.test_tracer.as_ref().cloned().map(TestTracerScope::new),
-        }
-    }
-
-    /// Creates a test telemetry context.
-    ///
-    /// Returned context has the same API as standard context, but also exposes API to obtain the
-    /// telemetry collected in it.
-    ///
-    /// # Examples
-    /// ```
-    /// use foundations::telemetry::TelemetryContext;
-    /// use foundations::telemetry::tracing::{self, test_trace};
-    /// use foundations::telemetry::log::{self, TestLogRecord};
-    /// use foundations::telemetry::settings::Level;
-    ///
-    /// #[tracing::span_fn("sync_fn")]
-    /// fn some_sync_production_fn_that_we_test() {
-    ///     log::warn!("Sync hello!");
-    /// }
-    ///
-    /// #[tracing::span_fn("async_fn")]
-    /// async fn some_async_production_fn_that_we_test() {
-    ///     log::warn!("Async hello!");
-    /// }
-    ///
-    /// #[tokio::main]
-    /// async fn main() {
-    ///     let ctx = TelemetryContext::test();
-    ///     
-    ///     {
-    ///         let _scope = ctx.scope();
-    ///         let _root = tracing::span("root");
-    ///
-    ///         let handle = tokio::spawn(TelemetryContext::current().apply(async {
-    ///             some_async_production_fn_that_we_test().await;
-    ///         }));
-    ///
-    ///         handle.await;
-    ///
-    ///         some_sync_production_fn_that_we_test();
-    ///     }
-    ///
-    ///     assert_eq!(*ctx.log_records(), &[
-    ///         TestLogRecord {
-    ///             level: Level::Warning,
-    ///             message: "Async hello!".into(),
-    ///             fields: vec![]
-    ///         },
-    ///         TestLogRecord {
-    ///             level: Level::Warning,
-    ///             message: "Sync hello!".into(),
-    ///             fields: vec![]
-    ///         }
-    ///     ]);  
-    ///
-    ///     assert_eq!(
-    ///         ctx.traces(Default::default()),
-    ///         vec![
-    ///             test_trace! {
-    ///                 "root" => {
-    ///                     "async_fn",
-    ///                     "sync_fn"
-    ///                 }
-    ///             }
-    ///         ]
-    ///     );
-    /// }
-    /// ```
-    #[cfg(feature = "testing")]
-    pub fn test() -> TestTelemetryContext {
-        TestTelemetryContext::new()
-    }
-
-    /// Wraps a future with the telemetry context.
-    ///
-    /// [`TelemetryScope`] can't be used across `await` points to propagate the telemetry context,
-    /// so to use telemetry context in async blocks, futures should be wrapped using this
-    /// method instead.
-    ///
-    /// Note that you don't need to use this method to wrap async function's bodies,
-    /// as [`tracing::span_fn`] macro takes care of that.
-    ///
-    /// # Examples
-    /// ```
-    /// use foundations::telemetry::TelemetryContext;
-    /// use foundations::telemetry::tracing::{self, test_trace};
-    ///
-    /// #[tokio::main]
-    /// async fn main() {
-    ///     // Test context is used for demonstration purposes to show the resulting traces.
-    ///     let ctx = TelemetryContext::test();
-    ///
-    ///     {
-    ///         let _scope = ctx.scope();
-    ///         let _root = tracing::span("root");
-    ///
-    ///         let handle = tokio::spawn(
-    ///             TelemetryContext::current().apply(async {
-    ///                 let _child = tracing::span("child");
-    ///             })
-    ///         );
-    ///
-    ///         handle.await;
-    ///     }
-    ///
-    ///     assert_eq!(
-    ///         ctx.traces(Default::default()),
-    ///         vec![
-    ///             test_trace! {
-    ///                 "root" => {
-    ///                     "child"
-    ///                 }
-    ///             }
-    ///         ]
-    ///     );
-    /// }
-    /// ```
-    pub fn apply<'f, F>(&self, fut: F) -> WithTelemetryContext<'f, F::Output>
-    where
-        F: Future + Send + 'f,
-    {
-        WithTelemetryContext {
-            inner: Box::pin(fut),
-            ctx: self.clone(),
-        }
-    }
-}
-
-#[cfg(feature = "tracing")]
-impl TelemetryContext {
-    /// Creates a new telemetry context, that includes a forked trace, creating a
-    /// linked child trace.
-    ///
-    /// If the current trace is sampled, the new child trace also will be sampled.
-    /// If the current trace isn't sampled, no new child trace is created.
-    ///
-    /// This method is useful to avoid a single trace from ballooning in size
-    /// while still keeping navigability from the source trace to the child
-    /// traces and vice-versa.
-    ///
-    /// # Examples
-    /// ```
-    /// use foundations::telemetry::TelemetryContext;
-    /// use foundations::telemetry::tracing::{self, test_trace};
-    ///
-    /// // Test context is used for demonstration purposes to show the resulting traces.
-    /// let ctx = TelemetryContext::test();
-    ///
-    /// {
-    ///     let _scope = ctx.scope();
-    ///     let _root = tracing::span("root");
-    ///
-    ///     {
-    ///         let _span1 = tracing::span("span1");
-    ///     }
-    ///
-    ///     let _scope = TelemetryContext::current()
-    ///         .with_forked_trace("new fork")
-    ///         .scope();
-    ///
-    ///     let _span2 = tracing::span("span2");
-    /// }
-    ///
-    /// assert_eq!(
-    ///     ctx.traces(Default::default()),
-    ///     vec![
-    ///         test_trace! {
-    ///             "root" => {
-    ///                 "span1",
-    ///                 "[new fork ref]"
-    ///             }
-    ///         },
-    ///         test_trace! {
-    ///             "new fork" => {
-    ///                 "span2"
-    ///             }
-    ///         }
-    ///     ]
-    /// );
-    pub fn with_forked_trace(&self, fork_name: impl Into<Cow<'static, str>>) -> Self {
-        Self {
-            #[cfg(feature = "logging")]
-            log: Arc::clone(&self.log),
-
-            span: Some(fork_trace(fork_name)),
-
-            #[cfg(feature = "testing")]
-            test_tracer: self.test_tracer.clone(),
-        }
-    }
-
-    /// Provides the same functionality as [`TelemetryContext::apply`], but also creates a tracing
-    /// span that is active during the future execution.
-    ///
-    /// # Examples
-    /// ```
-    /// use foundations::telemetry::TelemetryContext;
-    /// use foundations::telemetry::tracing::{self, test_trace};
-    ///
-    /// #[tokio::main]
-    /// async fn main() {
-    ///     // Test context is used for demonstration purposes to show the resulting traces.
-    ///     let ctx = TelemetryContext::test();
-    ///
-    ///     {
-    ///         let _scope = ctx.scope();
-    ///         let _root = tracing::span("root");
-    ///
-    ///         let handle = tokio::spawn(
-    ///             TelemetryContext::current().apply_with_tracing_span("future", async {
-    ///                 let _child = tracing::span("child");
-    ///             })
-    ///         );
-    ///
-    ///         handle.await;
-    ///     }
-    ///
-    ///     assert_eq!(
-    ///         ctx.traces(Default::default()),
-    ///         vec![
-    ///             test_trace! {
-    ///                 "root" => {
-    ///                     "future" => {
-    ///                         "child"
-    ///                     }
-    ///                 }
-    ///             }
-    ///         ]
-    ///     );
-    /// }
-    /// ```
-    pub fn apply_with_tracing_span<'f, F, N>(
-        &self,
-        span_name: N,
-        fut: F,
-    ) -> WithTelemetryContext<'f, F::Output>
-    where
-        F: Future + Send + 'f,
-        N: Into<Cow<'static, str>>,
-    {
-        let mut ctx = self.clone();
-        let _scope = ctx.span.as_ref().cloned().map(SpanScope::new);
-
-        ctx.span = Some(create_span(span_name));
-
-        WithTelemetryContext {
-            inner: Box::pin(fut),
-            ctx,
-        }
-    }
-}
-
-#[cfg(feature = "logging")]
-impl TelemetryContext {
-    /// Creates a telemetry context with log that is detached from the current context's log, but
-    /// inherits its log fields.
-    ///
-    /// For example, can be used in server software to produce separate logs for HTTP requests, each
-    /// of which has log fields added during the HTTP connection establishment.
-    ///
-    /// # Examples
-    /// ```
-    /// use foundations::telemetry::TelemetryContext;
-    /// use foundations::telemetry::log::{self, TestLogRecord};
-    /// use foundations::telemetry::settings::Level;
-    ///
-    /// // Test context is used for demonstration purposes to show the resulting log records.
-    /// let ctx = TelemetryContext::test();
-    /// let _scope = ctx.scope();
-    ///
-    /// log::add_fields!("conn_field" => 42);
-    ///
-    /// {
-    ///     let _scope = TelemetryContext::current().with_forked_log().scope();
-    ///
-    ///     log::add_fields!("req1_field" => "foo");
-    ///     log::warn!("Hello from request 1");
-    /// }
-    ///
-    /// {
-    ///     let _scope = TelemetryContext::current().with_forked_log().scope();
-    ///
-    ///     log::add_fields!("req2_field" => "bar");
-    ///     log::warn!("Hello from request 2");
-    /// }
-    ///
-    /// log::warn!("Hello from connection");
-    ///
-    /// assert_eq!(*ctx.log_records(), &[
-    ///     TestLogRecord {
-    ///         level: Level::Warning,
-    ///         message: "Hello from request 1".into(),
-    ///         fields: vec![
-    ///             ("req1_field".into(), "foo".into()),
-    ///             ("conn_field".into(), "42".into()),
-    ///         ]
-    ///     },
-    ///     TestLogRecord {
-    ///         level: Level::Warning,
-    ///         message: "Hello from request 2".into(),
-    ///         fields: vec![
-    ///             ("req2_field".into(), "bar".into()),
-    ///             ("conn_field".into(), "42".into()),
-    ///         ]
-    ///     },
-    ///     TestLogRecord {
-    ///         level: Level::Warning,
-    ///         message: "Hello from connection".into(),
-    ///         fields: vec![
-    ///             ("conn_field".into(), "42".into()),
-    ///         ]
-    ///     }
-    /// ]);
-    /// ```
-    pub fn with_forked_log(&self) -> Self {
-        Self {
-            log: fork_log(),
-
-            #[cfg(feature = "tracing")]
-            span: self.span.clone(),
-
-            #[cfg(all(feature = "tracing", feature = "testing"))]
-            test_tracer: self.test_tracer.clone(),
-        }
-    }
-}
-
 /// Telemetry configuration that is passed to [`init`].
 pub struct TelemetryConfig<'c> {
     /// Service information that is used in telemetry reporting.
diff --git a/foundations/src/telemetry/telemetry_context.rs b/foundations/src/telemetry/telemetry_context.rs
new file mode 100644
index 0000000..ba48dfc
--- /dev/null
+++ b/foundations/src/telemetry/telemetry_context.rs
@@ -0,0 +1,433 @@
+use super::TelemetryScope;
+use crate::utils::feature_use;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+feature_use!(cfg(feature = "logging"), {
+    use super::log::internal::{current_log, fork_log, LogScope, SharedLog};
+    use std::sync::Arc;
+});
+
+feature_use!(cfg(feature = "tracing"), {
+    use super::tracing::internal::{current_span, fork_trace, SharedSpan};
+    use super::tracing::SpanScope;
+    use std::borrow::Cow;
+
+    feature_use!(cfg(feature = "testing"), {
+        use super::tracing::internal::Tracer;
+        use super::tracing::testing::{current_test_tracer, TestTracerScope};
+    });
+});
+
+#[cfg(feature = "testing")]
+use super::testing::TestTelemetryContext;
+
+/// Wrapper for a future that provides it with [`TelemetryContext`].
+pub struct WithTelemetryContext<'f, T> {
+    // NOTE: we intentionally erase type here as we can get close to the type
+    // length limit, adding telemetry wrappers on top causes compiler to fail in some
+    // cases.
+    inner: Pin<Box<dyn Future<Output = T> + Send + 'f>>,
+    ctx: TelemetryContext,
+}
+
+impl<'f, T> Future for WithTelemetryContext<'f, T> {
+    type Output = T;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let _telemetry_scope = self.ctx.scope();
+
+        self.inner.as_mut().poll(cx)
+    }
+}
+
+/// The same as [`WithTelemetryContext`], but for futures that are `!Send`.
+pub struct WithTelemetryContextLocal<'f, T> {
+    // NOTE: we intentionally erase type here as we can get close to the type
+    // length limit, adding telemetry wrappers on top causes compiler to fail in some
+    // cases.
+    inner: Pin<Box<dyn Future<Output = T> + 'f>>,
+    ctx: TelemetryContext,
+}
+
+impl<'f, T> Future for WithTelemetryContextLocal<'f, T> {
+    type Output = T;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let _telemetry_scope = self.ctx.scope();
+
+        self.inner.as_mut().poll(cx)
+    }
+}
+
+/// Implicit context for logging and tracing.
+///
+/// Current context can be obtained with the [`TelemetryContext::current`] method.
+///
+/// The context can be forked with [`TelemetryContext::with_forked_log`] and
+/// [`TelemetryContext::with_forked_trace`] methods. And propagated in different scopes with
+/// [`TelemetryContext::scope`] and [`TelemetryContext::apply`] methods.
+#[derive(Debug, Clone)]
+pub struct TelemetryContext {
+    #[cfg(feature = "logging")]
+    pub(super) log: SharedLog,
+
+    // NOTE: we might not have tracing root at this point
+    #[cfg(feature = "tracing")]
+    pub(super) span: Option<SharedSpan>,
+
+    #[cfg(all(feature = "tracing", feature = "testing"))]
+    pub(super) test_tracer: Option<Tracer>,
+}
+
+impl TelemetryContext {
+    /// Returns the telemetry context that is active in the current scope.
+    pub fn current() -> Self {
+        Self {
+            #[cfg(feature = "logging")]
+            log: current_log(),
+
+            #[cfg(feature = "tracing")]
+            span: current_span(),
+
+            #[cfg(all(feature = "tracing", feature = "testing"))]
+            test_tracer: current_test_tracer(),
+        }
+    }
+
+    /// Creates a scope handle for the telemetry context.
+    ///
+    /// The telemetry context is active in the scope unless the handle is dropped.
+    ///
+    /// Note that the scope can only be used in the sync contexts. To propagate telemetry context
+    /// to async contexts, [`TelemetryContext::apply`] should be used instead.
+    ///
+    /// The scope handle is useful to propagate the telemetry context to callbacks provided
+    /// by third party libraries or other threads.
+    ///
+    /// The other use case is situations where code control flow doesn't align with telemetry flow.
+    /// E.g. there is a centralized dispatcher that drives certain tasks and tasks are registered
+    /// with the dispatcher via callbacks. It's desirable for each task to have its own telemetry
+    /// context, so scope can be used to propagate task context to the dispatcher's callbacks.
+    ///
+    /// # Examples
+    /// ```
+    /// use foundations::telemetry::TelemetryContext;
+    /// use foundations::telemetry::tracing::{self, test_trace};
+    ///
+    /// // Test context is used for demonstration purposes to show the resulting traces.
+    /// let ctx = TelemetryContext::test();
+    ///
+    /// {
+    ///     let _scope = ctx.scope();
+    ///     let _root = tracing::span("root");
+    ///     let telemetry_ctx = TelemetryContext::current();
+    ///
+    ///     let handle = std::thread::spawn(move || {
+    ///         let _scope = telemetry_ctx.scope();
+    ///         let _child = tracing::span("child");
+    ///     });
+    ///
+    ///     handle.join();
+    /// }
+    ///
+    /// assert_eq!(
+    ///     ctx.traces(Default::default()),
+    ///     vec![
+    ///         test_trace! {
+    ///             "root" => {
+    ///                 "child"
+    ///             }
+    ///         },
+    ///     ]
+    /// );
+    /// ```
+    pub fn scope(&self) -> TelemetryScope {
+        TelemetryScope {
+            #[cfg(feature = "logging")]
+            _log_scope: LogScope::new(Arc::clone(&self.log)),
+
+            #[cfg(feature = "tracing")]
+            _span_scope: self.span.as_ref().cloned().map(SpanScope::new),
+
+            #[cfg(all(feature = "tracing", feature = "testing"))]
+            _test_tracer_scope: self.test_tracer.as_ref().cloned().map(TestTracerScope::new),
+        }
+    }
+
+    /// Creates a test telemetry context.
+    ///
+    /// Returned context has the same API as standard context, but also exposes API to obtain the
+    /// telemetry collected in it.
+    ///
+    /// # Examples
+    /// ```
+    /// use foundations::telemetry::TelemetryContext;
+    /// use foundations::telemetry::tracing::{self, test_trace};
+    /// use foundations::telemetry::log::{self, TestLogRecord};
+    /// use foundations::telemetry::settings::Level;
+    ///
+    /// #[tracing::span_fn("sync_fn")]
+    /// fn some_sync_production_fn_that_we_test() {
+    ///     log::warn!("Sync hello!");
+    /// }
+    ///
+    /// #[tracing::span_fn("async_fn")]
+    /// async fn some_async_production_fn_that_we_test() {
+    ///     log::warn!("Async hello!");
+    /// }
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let ctx = TelemetryContext::test();
+    ///     
+    ///     {
+    ///         let _scope = ctx.scope();
+    ///         let _root = tracing::span("root");
+    ///
+    ///         let handle = tokio::spawn(TelemetryContext::current().apply(async {
+    ///             some_async_production_fn_that_we_test().await;
+    ///         }));
+    ///
+    ///         handle.await;
+    ///
+    ///         some_sync_production_fn_that_we_test();
+    ///     }
+    ///
+    ///     assert_eq!(*ctx.log_records(), &[
+    ///         TestLogRecord {
+    ///             level: Level::Warning,
+    ///             message: "Async hello!".into(),
+    ///             fields: vec![]
+    ///         },
+    ///         TestLogRecord {
+    ///             level: Level::Warning,
+    ///             message: "Sync hello!".into(),
+    ///             fields: vec![]
+    ///         }
+    ///     ]);  
+    ///
+    ///     assert_eq!(
+    ///         ctx.traces(Default::default()),
+    ///         vec![
+    ///             test_trace! {
+    ///                 "root" => {
+    ///                     "async_fn",
+    ///                     "sync_fn"
+    ///                 }
+    ///             }
+    ///         ]
+    ///     );
+    /// }
+    /// ```
+    #[cfg(feature = "testing")]
+    pub fn test() -> TestTelemetryContext {
+        TestTelemetryContext::new()
+    }
+
+    /// Wraps a future with the telemetry context.
+    ///
+    /// [`TelemetryScope`] can't be used across `await` points to propagate the telemetry context,
+    /// so to use telemetry context in async blocks, futures should be wrapped using this
+    /// method instead.
+    ///
+    /// Note that you don't need to use this method to wrap async function's bodies,
+    /// as [`tracing::span_fn`] macro takes care of that.
+    ///
+    /// # Examples
+    /// ```
+    /// use foundations::telemetry::TelemetryContext;
+    /// use foundations::telemetry::tracing::{self, test_trace};
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     // Test context is used for demonstration purposes to show the resulting traces.
+    ///     let ctx = TelemetryContext::test();
+    ///
+    ///     {
+    ///         let _scope = ctx.scope();
+    ///
+    ///         let handle = tokio::spawn(
+    ///             tracing::span("root").into_context().apply(async {
+    ///                 let _child = tracing::span("child");
+    ///             })
+    ///         );
+    ///
+    ///         handle.await;
+    ///     }
+    ///
+    ///     assert_eq!(
+    ///         ctx.traces(Default::default()),
+    ///         vec![
+    ///             test_trace! {
+    ///                 "root" => {
+    ///                     "child"
+    ///                 }
+    ///             }
+    ///         ]
+    ///     );
+    /// }
+    /// ```
+    ///
+    /// [`tracing::span_fn`]: crate::telemetry::tracing::span_fn
+    pub fn apply<'f, F>(&self, fut: F) -> WithTelemetryContext<'f, F::Output>
+    where
+        F: Future + Send + 'f,
+    {
+        WithTelemetryContext {
+            inner: Box::pin(fut),
+            ctx: self.clone(),
+        }
+    }
+
+    /// The same as [`TelemetryContext::apply`], but for futures that are `!Send`.
+    pub fn apply_local<'f, F>(&self, fut: F) -> WithTelemetryContextLocal<'f, F::Output>
+    where
+        F: Future + 'f,
+    {
+        WithTelemetryContextLocal {
+            inner: Box::pin(fut),
+            ctx: self.clone(),
+        }
+    }
+}
+
+#[cfg(feature = "tracing")]
+impl TelemetryContext {
+    /// Creates a new telemetry context, that includes a forked trace, creating a
+    /// linked child trace.
+    ///
+    /// If the current trace is sampled, the new child trace also will be sampled.
+    /// If the current trace isn't sampled, no new child trace is created.
+    ///
+    /// This method is useful to avoid a single trace from ballooning in size
+    /// while still keeping navigability from the source trace to the child
+    /// traces and vice-versa.
+    ///
+    /// # Examples
+    /// ```
+    /// use foundations::telemetry::TelemetryContext;
+    /// use foundations::telemetry::tracing::{self, test_trace};
+    ///
+    /// // Test context is used for demonstration purposes to show the resulting traces.
+    /// let ctx = TelemetryContext::test();
+    ///
+    /// {
+    ///     let _scope = ctx.scope();
+    ///     let _root = tracing::span("root");
+    ///
+    ///     {
+    ///         let _span1 = tracing::span("span1");
+    ///     }
+    ///
+    ///     let _scope = TelemetryContext::current()
+    ///         .with_forked_trace("new fork")
+    ///         .scope();
+    ///
+    ///     let _span2 = tracing::span("span2");
+    /// }
+    ///
+    /// assert_eq!(
+    ///     ctx.traces(Default::default()),
+    ///     vec![
+    ///         test_trace! {
+    ///             "root" => {
+    ///                 "span1",
+    ///                 "[new fork ref]"
+    ///             }
+    ///         },
+    ///         test_trace! {
+    ///             "new fork" => {
+    ///                 "span2"
+    ///             }
+    ///         }
+    ///     ]
+    /// );
+    pub fn with_forked_trace(&self, fork_name: impl Into<Cow<'static, str>>) -> Self {
+        Self {
+            #[cfg(feature = "logging")]
+            log: Arc::clone(&self.log),
+
+            span: Some(fork_trace(fork_name)),
+
+            #[cfg(feature = "testing")]
+            test_tracer: self.test_tracer.clone(),
+        }
+    }
+}
+
+#[cfg(feature = "logging")]
+impl TelemetryContext {
+    /// Creates a telemetry context with log that is detached from the current context's log, but
+    /// inherits its log fields.
+    ///
+    /// For example, can be used in server software to produce separate logs for HTTP requests, each
+    /// of which has log fields added during the HTTP connection establishment.
+    ///
+    /// # Examples
+    /// ```
+    /// use foundations::telemetry::TelemetryContext;
+    /// use foundations::telemetry::log::{self, TestLogRecord};
+    /// use foundations::telemetry::settings::Level;
+    ///
+    /// // Test context is used for demonstration purposes to show the resulting log records.
+    /// let ctx = TelemetryContext::test();
+    /// let _scope = ctx.scope();
+    ///
+    /// log::add_fields!("conn_field" => 42);
+    ///
+    /// {
+    ///     let _scope = TelemetryContext::current().with_forked_log().scope();
+    ///
+    ///     log::add_fields!("req1_field" => "foo");
+    ///     log::warn!("Hello from request 1");
+    /// }
+    ///
+    /// {
+    ///     let _scope = TelemetryContext::current().with_forked_log().scope();
+    ///
+    ///     log::add_fields!("req2_field" => "bar");
+    ///     log::warn!("Hello from request 2");
+    /// }
+    ///
+    /// log::warn!("Hello from connection");
+    ///
+    /// assert_eq!(*ctx.log_records(), &[
+    ///     TestLogRecord {
+    ///         level: Level::Warning,
+    ///         message: "Hello from request 1".into(),
+    ///         fields: vec![
+    ///             ("req1_field".into(), "foo".into()),
+    ///             ("conn_field".into(), "42".into()),
+    ///         ]
+    ///     },
+    ///     TestLogRecord {
+    ///         level: Level::Warning,
+    ///         message: "Hello from request 2".into(),
+    ///         fields: vec![
+    ///             ("req2_field".into(), "bar".into()),
+    ///             ("conn_field".into(), "42".into()),
+    ///         ]
+    ///     },
+    ///     TestLogRecord {
+    ///         level: Level::Warning,
+    ///         message: "Hello from connection".into(),
+    ///         fields: vec![
+    ///             ("conn_field".into(), "42".into()),
+    ///         ]
+    ///     }
+    /// ]);
+    /// ```
+    pub fn with_forked_log(&self) -> Self {
+        Self {
+            log: fork_log(),
+
+            #[cfg(feature = "tracing")]
+            span: self.span.clone(),
+
+            #[cfg(all(feature = "tracing", feature = "testing"))]
+            test_tracer: self.test_tracer.clone(),
+        }
+    }
+}
diff --git a/foundations/src/telemetry/tracing/mod.rs b/foundations/src/telemetry/tracing/mod.rs
index 133e104..b68d3c2 100644
--- a/foundations/src/telemetry/tracing/mod.rs
+++ b/foundations/src/telemetry/tracing/mod.rs
@@ -16,6 +16,7 @@ mod output_otlp_grpc;
 use self::init::TracingHarness;
 use self::internal::{create_span, current_span, span_trace_id, SharedSpan};
 use super::scope::Scope;
+use super::TelemetryContext;
 use cf_rustracing_jaeger::Span;
 use std::borrow::Cow;
 use std::sync::Arc;
@@ -86,6 +87,16 @@ pub use cf_rustracing_jaeger::span::SpanContextState as SerializableTraceState;
 /// );
 /// ```
 ///
+/// # Using with `async fn`'s that produce `!Send` futures.
+/// ```
+/// use foundations::telemetry::tracing;
+///
+/// #[tracing::span_fn("foo", async_local = true)]
+/// async fn foo() {
+///     // Does something that produces `!Send`` future...
+/// }
+/// ```
+///
 /// # Renamed or reexported crate
 ///
 /// The macro will fail to compile if `foundations` crate is reexported. However, the crate path
@@ -127,12 +138,68 @@ pub use foundations_macros::span_fn;
 ///
 /// Scope ends when the handle is dropped.
 #[must_use]
-pub struct SpanScope(#[allow(dead_code)] Scope<SharedSpan>);
+pub struct SpanScope {
+    span: SharedSpan,
+    _inner: Scope<SharedSpan>,
+}
 
 impl SpanScope {
     #[inline]
     pub(crate) fn new(span: SharedSpan) -> Self {
-        Self(Scope::new(&TracingHarness::get().span_scope_stack, span))
+        Self {
+            span: span.clone(),
+            _inner: Scope::new(&TracingHarness::get().span_scope_stack, span),
+        }
+    }
+
+    /// Converts the span scope to [`TelemetryContext`] that can be a applied to a future.
+    ///
+    /// This is effectively a shorthand for calling [`TelemetryContext::current`] with the span
+    /// being in scope.
+    ///
+    /// # Examples
+    /// ```
+    /// use foundations::telemetry::TelemetryContext;
+    /// use foundations::telemetry::tracing::{self, test_trace};
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     // Test context is used for demonstration purposes to show the resulting traces.
+    ///     let ctx = TelemetryContext::test();
+    ///
+    ///     {
+    ///         let _scope = ctx.scope();
+    ///         let _root = tracing::span("root");
+    ///
+    ///         let handle = tokio::spawn(
+    ///             tracing::span("future").into_context().apply(async {
+    ///                 let _child = tracing::span("child");
+    ///             })
+    ///         );
+    ///
+    ///         handle.await;
+    ///     }
+    ///
+    ///     assert_eq!(
+    ///         ctx.traces(Default::default()),
+    ///         vec![
+    ///             test_trace! {
+    ///                 "root" => {
+    ///                     "future" => {
+    ///                         "child"
+    ///                     }
+    ///                 }
+    ///             }
+    ///         ]
+    ///     );
+    /// }
+    /// ```
+    pub fn into_context(self) -> TelemetryContext {
+        let mut ctx = TelemetryContext::current();
+
+        ctx.span = Some(self.span);
+
+        ctx
     }
 }
 
@@ -220,10 +287,7 @@ pub fn state_for_trace_stitching() -> Option<SerializableTraceState> {
 /// If span covers whole function body it's preferable to use [`span_fn`] macro.
 ///
 /// Span ends when returned [`SpanScope`] is dropped. Note that [`SpanScope`] can't be used across
-/// `await` points. To span async scopes [`TelemetryContext::apply_with_tracing_span`] should be
-/// used.
-///
-/// [`TelemetryContext::apply_with_tracing_span`]: super::TelemetryContext::apply_with_tracing_span
+/// `await` points. To span async scopes [`SpanScope::into_context`] should be used.
 ///
 /// # Examples
 /// ```
diff --git a/foundations/src/telemetry/tracing/output_jaeger_thrift_udp.rs b/foundations/src/telemetry/tracing/output_jaeger_thrift_udp.rs
index 7dfd8c0..f432bb5 100644
--- a/foundations/src/telemetry/tracing/output_jaeger_thrift_udp.rs
+++ b/foundations/src/telemetry/tracing/output_jaeger_thrift_udp.rs
@@ -8,7 +8,6 @@ use cf_rustracing_jaeger::span::SpanReceiver;
 use futures_util::future::{BoxFuture, FutureExt as _};
 use std::net::SocketAddr;
 use std::net::{Ipv4Addr, Ipv6Addr};
-use std::sync::Arc;
 
 pub(super) fn start(
     service_info: ServiceInfo,
@@ -59,18 +58,10 @@ fn get_reporter_bind_addr(settings: &JaegerThriftUdpOutputSettings) -> Bootstrap
 }
 
 async fn do_export(reporter: JaegerCompactReporter, mut span_rx: SpanReceiver) {
-    let reporter = Arc::new(reporter);
-
     while let Some(span) = span_rx.recv().await {
         // NOTE: we are limited with a UDP dgram size here, so doing batching is risky.
-        tokio::spawn({
-            let reporter = Arc::clone(&reporter);
-
-            async move {
-                if let Err(err) = reporter.report(&[span][..]).await {
-                    reporter_error(err);
-                }
-            }
-        });
+        if let Err(err) = reporter.report(&[span][..]).await {
+            reporter_error(err);
+        }
     }
 }


Test patch:
diff --git a/foundations/src/telemetry/tracing/testing.rs b/foundations/src/telemetry/tracing/testing.rs
index 4e61bdb..32e6955 100644
--- a/foundations/src/telemetry/tracing/testing.rs
+++ b/foundations/src/telemetry/tracing/testing.rs
@@ -146,15 +146,16 @@ pub struct TestTraceOptions {
 }
 
 #[must_use]
-pub(crate) struct TestTracerScope(#[allow(dead_code)] Scope<Tracer>);
+pub(crate) struct TestTracerScope {
+    _inner: Scope<Tracer>,
+}
 
 impl TestTracerScope {
     #[inline]
     pub(crate) fn new(tracer: Tracer) -> Self {
-        Self(Scope::new(
-            &TracingHarness::get().test_tracer_scope_stack,
-            tracer,
-        ))
+        Self {
+            _inner: Scope::new(&TracingHarness::get().test_tracer_scope_stack, tracer),
+        }
     }
 }