Add API for automatically recording trace spans duration to a histogram

This adds an API for automatically tracking span durations with a
`TimeHistogram`. The durations are stored regardless of if the
underlying span is sampled.

Author: evanrittenhouse

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_time_histogram")]
+    with_time_histogram: Option<Expr>,
 }
 
 impl Options {
@@ -54,6 +57,10 @@ impl Options {
     fn default_async_local() -> bool {
         false
     }
+
+    fn default_time_histogram() -> Option<Expr> {
+        None
+    }
 }
 
 struct Args {
@@ -113,10 +120,19 @@ fn expand_from_parsed(args: Args, item_fn: ItemFn) -> TokenStream2 {
             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
-            )
+            if let Some(timer) = args.options.with_time_histogram {
+                quote!(
+                    let __span = #crate_path::telemetry::tracing::span(#span_name)
+                        .into_context()
+                        .with_time_histogram(#timer);
+                    #block
+                )
+            } else {
+                quote!(
+                    let __span = #crate_path::telemetry::tracing::span(#span_name);
+                    #block
+                )
+            }
         }),
     };
 
@@ -182,13 +198,22 @@ fn wrap_with_span(args: &Args, block: TokenStream2) -> TokenStream2 {
 
     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
-    )
+    if let Some(timer) = &args.options.with_time_histogram {
+        quote!(
+            #crate_path::telemetry::tracing::span(#span_name)
+                .into_context()
+                .with_time_histogram(#timer)
+                .#apply_fn(#block)
+                .await
+        )
+    } else {
+        quote!(
+            #crate_path::telemetry::tracing::span(#span_name)
+                .into_context()
+                .#apply_fn(#block)
+                .await
+        )
+    }
 }
 
 #[cfg(test)]
@@ -482,4 +507,69 @@ mod tests {
 
         assert_eq!(actual, expected);
     }
+
+    #[test]
+    fn span_sync_fn_with_metric_tracker() {
+        let args = parse_attr! {
+            #[span_fn("sync_span", with_time_histogram = ::foo::bar::metrics::time_histogram(label = false))]
+        };
+
+        let item_fn = parse_quote! {
+            fn do_sync() -> io::Result<String> {
+                do_something_else();
+
+                Ok("foo".into())
+            }
+        };
+
+        let actual = expand_from_parsed(args, item_fn).to_string();
+
+        let expected = code_str! {
+            fn do_sync<>() -> io::Result<String> {
+                let __span = ::foundations::telemetry::tracing::span("sync_span")
+                    .into_context()
+                    .with_time_histogram(::foo::bar::metrics::time_histogram(label = false));
+                {
+                    do_something_else();
+
+                    Ok("foo".into())
+                }
+            }
+        };
+
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn span_async_fn_with_metric_tracker() {
+        let args = parse_attr! {
+            #[span_fn("async_span", with_time_histogram = ::foo::bar::metrics::time_histogram(label = false))]
+        };
+
+        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()
+                    .with_time_histogram(::foo::bar::metrics::time_histogram(label = false))
+                    .apply(async move {{
+                        do_something_else().await;
+
+                        Ok("foo".into())
+                    }})
+                .await
+            }
+        };
+
+        assert_eq!(actual, expected);
+    }
 }

foundations/src/telemetry/telemetry_context.rs

@@ -18,6 +18,10 @@ feature_use!(cfg(feature = "tracing"), {
         use super::tracing::internal::Tracer;
         use super::tracing::testing::{current_test_tracer, TestTracerScope};
     });
+
+    feature_use!(cfg(feature = "testing"), {
+        use prometools::histogram::TimeHistogram;
+    });
 });
 
 #[cfg(feature = "testing")]
@@ -226,6 +230,54 @@ impl TelemetryContext {
         TestTelemetryContext::new()
     }
 
+    /// Track the duration of the current `span` in this `TelemetryContext`. The span's duration
+    /// will be inserted into `histogram`, regardless of whether `span` was sampled or not.
+    ///
+    /// # Examples
+    /// ```
+    /// use foundations::telemetry::TelemetryContext;
+    /// use foundations::telemetry::metrics::{metrics, HistogramBuilder, TimeHistogram};
+    /// use foundations::telemetry::tracing::{self, test_trace};
+    ///
+    /// #[metrics]
+    /// mod test_metrics {
+    ///     #[ctor = HistogramBuilder { buckets: &[1E-4, 2E-4] }]
+    ///    pub fn histogram(label: bool) -> TimeHistogram;
+    /// }
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     // Test context is used for demonstration purposes to show the resulting traces.
+    ///     let ctx = TelemetryContext::test();
+    ///     let hist = test_metrics::histogram(true);
+    ///     let clone = hist.clone();
+    ///
+    ///     {
+    ///         let _scope = ctx.scope();
+    ///         let _root = tracing::span("root");
+    ///         let telemetry_ctx = TelemetryContext::current().with_time_histogram(clone);
+    ///
+    ///         let handle = std::thread::spawn(move || {
+    ///             let _scope = telemetry_ctx.scope();
+    ///             let _child = tracing::span("child");
+    ///         });
+    ///
+    ///         handle.join();
+    ///     }
+    ///
+    ///     let snapshot = hist.snapshot();
+    ///     assert_eq!(snapshot.count(), 1);
+    /// }
+    /// ```
+    #[cfg(all(feature = "metrics", feature = "tracing"))]
+    pub fn with_time_histogram(mut self, histogram: TimeHistogram) -> Self {
+        if let Some(inner) = self.span.as_mut() {
+            inner.track_with_time_histogram(histogram);
+        }
+
+        self
+    }
+
     /// Wraps a future with the telemetry context.
     ///
     /// [`TelemetryScope`] can't be used across `await` points to propagate the telemetry context,

foundations/src/telemetry/tracing/internal.rs

@@ -6,6 +6,7 @@ use cf_rustracing::sampler::BoxSampler;
 use cf_rustracing::tag::Tag;
 use cf_rustracing_jaeger::span::{Span, SpanContext, SpanContextState};
 use parking_lot::RwLock;
+use prometools::histogram::{HistogramTimer, TimeHistogram};
 use rand::{self, Rng};
 use std::borrow::Cow;
 use std::error::Error;
@@ -54,14 +55,31 @@ impl From<SharedSpanHandle> for Arc<RwLock<Span>> {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub(crate) struct SharedSpan {
     // NOTE: we intentionally use a lock without poisoning here to not
     // panic the threads if they just share telemetry with failed thread.
     pub(crate) inner: SharedSpanHandle,
     // NOTE: store sampling flag separately, so we don't need to acquire lock
     // every time we need to check the flag.
     is_sampled: bool,
+    timer: Option<Arc<HistogramTimer>>,
+}
+
+impl SharedSpan {
+    #[cfg(all(feature = "metrics", feature = "tracing"))]
+    pub(crate) fn track_with_time_histogram(&mut self, timer: TimeHistogram) {
+        let _ = self.timer.insert(Arc::new(timer.start_timer()));
+    }
+}
+
+impl std::fmt::Debug for SharedSpan {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SharedSpan")
+            .field("inner", &self.inner)
+            .field("is_sampled", &self.is_sampled)
+            .finish()
+    }
 }
 
 impl From<Span> for SharedSpan {
@@ -71,6 +89,7 @@ impl From<Span> for SharedSpan {
         Self {
             inner: SharedSpanHandle::new(inner),
             is_sampled,
+            timer: None,
         }
     }
 }
@@ -83,12 +102,22 @@ pub fn write_current_span(write_fn: impl FnOnce(&mut Span)) {
     }
 }
 
-pub(crate) fn create_span(name: impl Into<Cow<'static, str>>) -> SharedSpan {
-    match current_span() {
+pub(crate) fn create_span(
+    name: impl Into<Cow<'static, str>>,
+    histogram: Option<TimeHistogram>,
+) -> SharedSpan {
+    let span = match current_span() {
         Some(parent) => parent.inner.read().child(name, |o| o.start()),
         None => start_trace(name, Default::default()),
+    };
+
+    let timer = histogram.map(|t| Arc::new(t.start_timer()));
+
+    SharedSpan {
+        is_sampled: span.is_sampled(),
+        inner: SharedSpanHandle::new(span),
+        timer,
     }
-    .into()
 }
 
 pub(crate) fn current_span() -> Option<SharedSpan> {