fix(datadog-lambda): track cold start per handler

Author: Dogbu-cyber

integrations/aws/datadog-lambda/src/invocation.rs

@@ -3,24 +3,17 @@
 
 use crate::attribute_keys as attr;
 use crate::span_inferrer::{InferredSpanScope, TriggerContext, TriggerExtractor};
-use crate::Config;
 
 use lambda_runtime::LambdaEvent;
 use opentelemetry::trace::{FutureExt, SpanKind, TraceContextExt, Tracer, TracerProvider as _};
 use opentelemetry::{Context, KeyValue};
 use opentelemetry_sdk::trace::{SdkTracer, SdkTracerProvider};
 use serde_json::Value;
 use std::future::Future;
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::SystemTime;
 
-static IS_COLD_START: AtomicBool = AtomicBool::new(true);
 static TRACER_NAME: &str = "datadog-lambda-rs";
 
-fn detect_cold_start() -> bool {
-    IS_COLD_START.swap(false, Ordering::Relaxed)
-}
-
 /// The Lambda invocation (aws.lambda) root span.
 pub(crate) struct LambdaSpan {
     cx: Context,
@@ -103,20 +96,22 @@ impl LambdaSpan {
 pub(crate) fn start_invocation(
     event: &LambdaEvent<Value>,
     provider: &SdkTracerProvider,
-    _config: &Config,
+    cold_start: bool,
 ) -> (LambdaSpan, InferredSpanScope) {
-    let is_cold = detect_cold_start();
     let tracer = provider.tracer(TRACER_NAME);
     let extraction = TriggerExtractor::extract(&event.payload);
-    let (parent_cx, inferred_spans) =
-        InferredSpanScope::start(&tracer, &extraction.upstream_cx, extraction.inferred_span.as_ref());
+    let (parent_cx, inferred_spans) = InferredSpanScope::start(
+        &tracer,
+        &extraction.upstream_cx,
+        extraction.inferred_span.as_ref(),
+    );
     let trigger = TriggerContext {
         parent_cx,
         is_async: extraction.is_async,
         event_source: extraction.event_source,
         event_source_arn: extraction.event_source_arn,
     };
-    let lambda_span = LambdaSpan::start(&tracer, &trigger, &event.context, is_cold);
+    let lambda_span = LambdaSpan::start(&tracer, &trigger, &event.context, cold_start);
 
     (lambda_span, inferred_spans)
 }
@@ -310,7 +305,7 @@ mod tests {
             lambda_runtime::Context::default(),
         );
 
-        let (span, _) = start_invocation(&raw_event, &provider, &Config::default());
+        let (span, _) = start_invocation(&raw_event, &provider, false);
         assert!(span.cx.span().span_context().is_valid());
     }
 
@@ -341,7 +336,7 @@ mod tests {
             test_lambda_cx(),
         );
 
-        let (span, inferred_spans) = start_invocation(&raw_event, &provider, &Config::default());
+        let (span, inferred_spans) = start_invocation(&raw_event, &provider, false);
         assert!(span.cx.span().span_context().is_valid());
         assert!(!inferred_spans.is_empty());
     }

integrations/aws/datadog-lambda/src/lib.rs

@@ -21,21 +21,23 @@ type BoxFuture<T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send
 pub struct Config {
     /// Service name. Overrides `DD_SERVICE`. Ignored when [`tracing`](Self::tracing) is `Some`.
     pub service: Option<String>,
-    /// Deployment environment. Overrides `DD_ENV`. Ignored when [`tracing`](Self::tracing) is `Some`.
+    /// Deployment environment. Overrides `DD_ENV`. Ignored when [`tracing`](Self::tracing) is
+    /// `Some`.
     pub env: Option<String>,
     /// Service version. Overrides `DD_VERSION`. Ignored when [`tracing`](Self::tracing) is `Some`.
     pub version: Option<String>,
     /// Full control over the OTel SDK and Datadog tracer config.
     ///
-    /// When `None` (default), Lambda-appropriate defaults are applied and `service`/`env`/`version`
-    /// above are forwarded. When `Some`, the builder is used as-is; `service`/`env`/`version` are
-    /// ignored and you are responsible for setting:
-    /// - `trace_stats_computation_enabled = false` (the Datadog agent handles stats for serverless environments)
-    /// - `trace_writer_synchronous_write = true` (so `force_flush()` blocks until spans reach agent)
+    /// When `None` (default), Lambda-appropriate defaults are applied and
+    /// `service`/`env`/`version` above are forwarded. When `Some`, the builder is used as-is;
+    /// `service`/`env`/`version` are ignored and you are responsible for setting:
+    /// - `trace_stats_computation_enabled = false` (the Datadog agent handles stats for serverless
+    ///   environments)
+    /// - `trace_writer_synchronous_write = true` (so `force_flush()` blocks until spans reach
+    ///   agent)
     pub tracing: Option<datadog_opentelemetry::DatadogTracingBuilder>,
 }
 
-
 fn build_tracing(
     service: Option<String>,
     env: Option<String>,
@@ -98,23 +100,32 @@ fn build_tracing(
 pub struct WrappedHandler<F, E, R> {
     inner: Arc<F>,
     provider: opentelemetry_sdk::trace::SdkTracerProvider,
-    config: Arc<Config>,
+    cold_start: bool,
     _phantom: PhantomData<fn(E) -> R>,
 }
 
 impl<F, E, R> WrappedHandler<F, E, R> {
     pub fn new(handler: F, config: Config) -> Self {
-        let Config { tracing, service, env, version } = config;
+        let Config {
+            tracing,
+            service,
+            env,
+            version,
+        } = config;
         let provider = tracing
             .unwrap_or_else(|| build_tracing(service, env, version))
             .init();
         Self {
             inner: Arc::new(handler),
             provider,
-            config: Arc::new(Config::default()),
+            cold_start: true,
             _phantom: PhantomData,
         }
     }
+
+    fn take_cold_start(&mut self) -> bool {
+        std::mem::replace(&mut self.cold_start, false)
+    }
 }
 
 impl<F, Fut, E, R> Service<LambdaEvent<Value>> for WrappedHandler<F, E, R>
@@ -133,18 +144,58 @@ where
     }
 
     fn call(&mut self, event: LambdaEvent<Value>) -> Self::Future {
+        let cold_start = self.take_cold_start();
         let inner_handler = Arc::clone(&self.inner);
         let provider = self.provider.clone();
-        let lambda_config = Arc::clone(&self.config);
-        let (lambda_span, inferred_spans) = start_invocation(&event, &provider, &lambda_config);
+        let (lambda_span, inferred_spans) = start_invocation(&event, &provider, cold_start);
         let typed_payload = match serde_json::from_value::<E>(event.payload) {
             Ok(payload) => payload,
             Err(err) => {
-                return Box::pin(run_handler(lambda_span, inferred_spans, provider, std::future::ready(Err(err.into()))));
+                return Box::pin(run_handler(
+                    lambda_span,
+                    inferred_spans,
+                    provider,
+                    std::future::ready(Err(err.into())),
+                ));
             }
         };
         let typed_event = LambdaEvent::new(typed_payload, event.context);
         let fut = inner_handler(typed_event);
         Box::pin(async move { run_handler(lambda_span, inferred_spans, provider, fut).await })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn noop_handler(
+        _: LambdaEvent<Value>,
+    ) -> std::future::Ready<Result<(), lambda_runtime::Error>> {
+        std::future::ready(Ok(()))
+    }
+
+    fn test_handler() -> WrappedHandler<
+        fn(LambdaEvent<Value>) -> std::future::Ready<Result<(), lambda_runtime::Error>>,
+        Value,
+        (),
+    > {
+        WrappedHandler {
+            inner: Arc::new(noop_handler),
+            provider: opentelemetry_sdk::trace::SdkTracerProvider::builder().build(),
+            cold_start: true,
+            _phantom: PhantomData,
+        }
+    }
+
+    #[test]
+    fn cold_start_is_tracked_per_handler() {
+        let mut first = test_handler();
+        assert!(first.take_cold_start());
+        assert!(!first.take_cold_start());
+
+        let mut second = test_handler();
+        assert!(second.take_cold_start());
+        assert!(!second.take_cold_start());
+    }
+}