telemetry: add apply_generic to telemetry context

This allows user to skip the boxing step and save a memory allocation.
This is configure with `generic=true` in attribute. For example

```
 #[span_fn("async_trait_span", generic = true)]
```

Also add a feature `generic-telemetry-wrapper` so user can enable this
behavior by default.

Author: dqminh

Cargo.toml

@@ -47,6 +47,7 @@ once_cell = "1.5"
 tonic = { version = "0.12", default-features = false }
 opentelemetry-proto = "0.7"
 parking_lot = "0.12.1"
+pin-project-lite = "0.2.16"
 proc-macro2 = { version = "1", default-features = false }
 prometheus = { version = "0.13.3", default-features = false }
 prometheus-client = "0.18.1"

foundations-macros/Cargo.toml

@@ -11,6 +11,8 @@ readme = "../README.md"
 [features]
 default = ["settings_deny_unknown_fields_by_default"]
 settings_deny_unknown_fields_by_default = []
+# Enable generic telemetry wrapper for futures
+generic-telemetry-wrapper = []
 
 [lib]
 proc-macro = true

foundations-macros/src/span_fn.rs

@@ -44,6 +44,9 @@ struct Options {
 
     #[darling(default = "Options::default_async_local")]
     async_local: bool,
+
+    #[darling(default = "Options::default_generic")]
+    generic: bool,
 }
 
 impl Options {
@@ -54,6 +57,16 @@ impl Options {
     fn default_async_local() -> bool {
         false
     }
+
+    #[cfg(feature = "generic-telemetry-wrapper")]
+    fn default_generic() -> bool {
+        true
+    }
+
+    #[cfg(not(feature = "generic-telemetry-wrapper"))]
+    fn default_generic() -> bool {
+        false
+    }
 }
 
 struct Args {
@@ -176,6 +189,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 +336,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 +500,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

@@ -60,6 +60,9 @@ settings = [
 # Whether settings structs annotated with `#[settings]` will, by default, error on unknown fields.
 settings_deny_unknown_fields_by_default = ["foundations-macros?/settings_deny_unknown_fields_by_default"]
 
+# Enable generic telemetry wrapper for futures
+generic-telemetry-wrapper = ["foundations-macros?/generic-telemetry-wrapper"]
+
 # Enables all the telemetry-related features ("logging", "metrics", "tracing", "telemetry-server").
 telemetry = [
     "logging",
@@ -233,6 +236,7 @@ log = { workspace = true, optional = true }
 parking_lot_core = { workspace = true, optional = true }
 regex = { workspace = true, optional = true }
 thiserror = { 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/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")]