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)]
```

Author: dqminh

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,10 @@ impl Options {
     fn default_async_local() -> bool {
         false
     }
+
+    fn default_generic() -> bool {
+        false
+    }
 }
 
 struct Args {
@@ -176,6 +183,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 +330,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 +494,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

@@ -230,6 +230,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 = "1.1.10"
 
 [target.'cfg(target_os = "linux")'.dependencies]
 tikv-jemalloc-ctl = { workspace = true, optional = true, features = [

foundations/src/telemetry/telemetry_context.rs

@@ -3,6 +3,7 @@ use crate::utils::feature_use;
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use pin_project::pin_project;
 
 feature_use!(cfg(feature = "logging"), {
     use super::log::internal::{current_log, fork_log, LogScope, SharedLog};
@@ -23,6 +24,24 @@ feature_use!(cfg(feature = "tracing"), {
 #[cfg(feature = "testing")]
 use super::testing::TestTelemetryContext;
 
+#[pin_project]
+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)
+    }
+}
+
 /// 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
@@ -181,7 +200,7 @@ impl TelemetryContext {
     /// #[tokio::main]
     /// async fn main() {
     ///     let ctx = TelemetryContext::test();
-    ///     
+    ///
     ///     {
     ///         let _scope = ctx.scope();
     ///         let _root = tracing::span("root");
@@ -206,7 +225,7 @@ impl TelemetryContext {
     ///             message: "Sync hello!".into(),
     ///             fields: vec![]
     ///         }
-    ///     ]);  
+    ///     ]);
     ///
     ///     assert_eq!(
     ///         ctx.traces(Default::default()),
@@ -291,6 +310,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")]