telemetry: remove Box::pin from WithTelemetryContext

Cargo.toml

@@ -47,6 +47,7 @@ once_cell = "1.5"
 tonic = { version = "0.13", default-features = false }
 opentelemetry-proto = "0.30"
 parking_lot = "0.12"
+pin-project-lite = "0.2.16"
 proc-macro2 = { version = "1", default-features = false }
 prometheus = { version = "0.14", default-features = false }
 prometheus-client = "0.18"

foundations-macros/src/span_fn.rs

@@ -1,3 +1,6 @@
+// NOTE: required to allow `foundations_generic_telemetry_wrapper` cfg
+#![allow(unexpected_cfgs)]
+
 use crate::common::parse_optional_trailing_meta_list;
 use darling::FromMeta;
 use proc_macro::TokenStream;
@@ -44,6 +47,9 @@ struct Options {
 
     #[darling(default = "Options::default_async_local")]
     async_local: bool,
+
+    #[darling(default = "Options::default_generic")]
+    generic: bool,
 }
 
 impl Options {
@@ -54,6 +60,10 @@ impl Options {
     fn default_async_local() -> bool {
         false
     }
+
+    fn default_generic() -> bool {
+        cfg!(foundations_generic_telemetry_wrapper)
+    }
 }
 
 struct Args {
@@ -176,6 +186,8 @@ 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 if args.options.generic {
+        quote!(apply_generic)
     } else {
         quote!(apply)
     };
@@ -321,6 +333,38 @@ mod tests {
         assert_eq!(actual, expected);
     }
 
+    #[test]
+    fn expand_async_fn_generic() {
+        let args = parse_attr! {
+            #[span_fn("async_span", generic = 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_generic(async move {{
+                        do_something_else().await;
+
+                        Ok("foo".into())
+                    }})
+                    .await
+            }
+        };
+
+        assert_eq!(actual, expected);
+    }
+
     #[test]
     fn expand_async_trait_fn() {
         let args = parse_attr! {
@@ -453,6 +497,72 @@ mod tests {
         assert_eq!(actual, expected);
     }
 
+    #[test]
+    fn expand_async_trait_fn_generic() {
+        let args = parse_attr! {
+            #[span_fn("async_trait_span", generic = 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_generic(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_structure_with_crate_path() {
         let args = parse_attr! {

foundations/Cargo.toml

@@ -175,10 +175,10 @@ tracing-rs-compat = ["dep:tracing-slog"]
 
 [package.metadata.docs.rs]
 all-features = true
-rustdoc-args = ["--cfg", "docsrs", "--cfg", "tokio_unstable", "--cfg", "foundations_unstable"]
+rustdoc-args = ["--cfg", "docsrs", "--cfg", "tokio_unstable", "--cfg", "foundations_unstable", "--cfg", "foundations_generic_telemetry_wrapper"]
 # it's necessary to _also_ pass `--cfg tokio_unstable` and `--cfg foundations_unstable`
 # to rustc, or else dependencies will not be enabled, and the docs build will fail.
-rustc-args = ["--cfg", "tokio_unstable", "--cfg", "foundations_unstable"]
+rustc-args = ["--cfg", "tokio_unstable", "--cfg", "foundations_unstable", "--cfg", "foundations_generic_telemetry_wrapper"]
 
 [dependencies]
 anyhow = { workspace = true, features = ["backtrace", "std"] }
@@ -235,6 +235,7 @@ parking_lot_core = { workspace = true, optional = true }
 regex = { workspace = true, optional = true }
 thiserror = { workspace = true, optional = true }
 tower = { workspace = true, optional = true }
+pin-project-lite = { workspace = true }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 tikv-jemalloc-ctl = { workspace = true, optional = true, features = [

foundations/src/lib.rs

@@ -43,6 +43,13 @@
 //! - **cli**: Enables command line interface (CLI) functionality. Implicitly enabled **settings**
 //!   feature.
 //!
+//! # Generic telemetry
+//! Foundations currently box the future with TelemetryContext by default. A default generic
+//! wrapper is gated behind `--cfg foundations_generic_telemetry_wrapper`.
+//!
+//! To enable this, you must add `--cfg foundations_generic_telemetry_wrapper` to your RUSTFLAGS
+//! environment variable.
+//!
 //! # Unstable Features
 //! Foundations has unstable features which are gated behind `--cfg foundations_unstable`:
 //!

foundations/src/telemetry/telemetry_context.rs

@@ -1,5 +1,6 @@
 use super::TelemetryScope;
 use crate::utils::feature_use;
+use pin_project_lite::pin_project;
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
@@ -61,6 +62,26 @@ impl<T> Future for WithTelemetryContextLocal<'_, T> {
     }
 }
 
+pin_project! {
+    /// The same as [`WithTelemetryContext`], but for futures that are not boxed
+    pub struct WithTelemetryContextGeneric<T> {
+        #[pin]
+        inner: T,
+        ctx: TelemetryContext,
+    }
+}
+
+impl<T: Future> Future for WithTelemetryContextGeneric<T> {
+    type Output = T::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let _telemetry_scope = this.ctx.scope();
+
+        this.inner.poll(cx)
+    }
+}
+
 /// Implicit context for logging and tracing.
 ///
 /// Current context can be obtained with the [`TelemetryContext::current`] method.
@@ -181,7 +202,7 @@ impl TelemetryContext {
     /// #[tokio::main]
     /// async fn main() {
     ///     let ctx = TelemetryContext::test();
-    ///     
+    ///
     ///     {
     ///         let _scope = ctx.scope();
     ///         let _root = tracing::span("root");
@@ -206,7 +227,7 @@ impl TelemetryContext {
     ///             message: "Sync hello!".into(),
     ///             fields: vec![]
     ///         }
-    ///     ]);  
+    ///     ]);
     ///
     ///     assert_eq!(
     ///         ctx.traces(Default::default()),
@@ -291,6 +312,17 @@ impl TelemetryContext {
             ctx: self.clone(),
         }
     }
+
+    /// The same as [`TelemetryContext::apply`], but for futures that are not boxed.
+    pub fn apply_generic<F>(&self, fut: F) -> WithTelemetryContextGeneric<F>
+    where
+        F: Future,
+    {
+        WithTelemetryContextGeneric {
+            inner: fut,
+            ctx: self.clone(),
+        }
+    }
 }
 
 #[cfg(feature = "tracing")]