diff --git a/Cargo.toml b/Cargo.toml index f72207da..b69a29c8 100644 --- a/Cargo.toml +++ b/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" diff --git a/foundations-macros/src/span_fn.rs b/foundations-macros/src/span_fn.rs index 8e5dd5ac..4a88d758 100644 --- a/foundations-macros/src/span_fn.rs +++ b/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 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 { + 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 { + ::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 + ::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:: { + 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 + ::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:: { + 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! { diff --git a/foundations/Cargo.toml b/foundations/Cargo.toml index c20ec217..eaad9b6e 100644 --- a/foundations/Cargo.toml +++ b/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 = [ diff --git a/foundations/src/lib.rs b/foundations/src/lib.rs index 8c71bfca..dfd478f5 100644 --- a/foundations/src/lib.rs +++ b/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`: //! diff --git a/foundations/src/telemetry/telemetry_context.rs b/foundations/src/telemetry/telemetry_context.rs index 195f9880..8ec720b2 100644 --- a/foundations/src/telemetry/telemetry_context.rs +++ b/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 Future for WithTelemetryContextLocal<'_, T> { } } +pin_project! { + /// The same as [`WithTelemetryContext`], but for futures that are not boxed + pub struct WithTelemetryContextGeneric { + #[pin] + inner: T, + ctx: TelemetryContext, + } +} + +impl Future for WithTelemetryContextGeneric { + type Output = T::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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(&self, fut: F) -> WithTelemetryContextGeneric + where + F: Future, + { + WithTelemetryContextGeneric { + inner: fut, + ctx: self.clone(), + } + } } #[cfg(feature = "tracing")]