Merge remote-tracking branch 'origin/main' into zarir/base_service

Author: zarirhamza

.gitlab/templates/pipeline.yaml.tpl

@@ -52,9 +52,9 @@ bottlecap ({{ $flavor.name }}):
   image: registry.ddbuild.io/images/docker:20.10
   tags: ["arch:{{ $flavor.arch }}"]
   needs: []
-  # This job sometimes times out on GitLab for unclear reason.
-  # Set a short timeout with retries to work around this.
-  timeout: 10m
+  # This job sometimes times out on GitLab for unclear reasons.
+  # Set a timeout with retries to work around this.
+  timeout: 20m
   retry:
     max: 2
     when:
@@ -464,6 +464,7 @@ integration-suite:
   stage: integration-tests
   tags: ["arch:amd64"]
   image: ${CI_DOCKER_TARGET_IMAGE}:${CI_DOCKER_TARGET_VERSION}
+  retry: 2
   parallel:
     matrix:
       - TEST_SUITE: {{ range (ds "test_suites").test_suites }}

bottlecap/src/config/aws.rs

@@ -1,6 +1,8 @@
 use std::env;
 use tokio::time::Instant;
 
+use crate::tags::lambda::tags::SNAP_START_VALUE;
+
 const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION";
 const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID";
 const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
@@ -46,6 +48,11 @@ impl AwsConfig {
         self.initialization_type
             .eq(LAMBDA_MANAGED_INSTANCES_INIT_TYPE)
     }
+
+    #[must_use]
+    pub fn is_snapstart(&self) -> bool {
+        self.initialization_type.eq(SNAP_START_VALUE)
+    }
 }
 
 #[allow(clippy::module_name_repetitions)]

bottlecap/src/lifecycle/invocation/context.rs

@@ -43,16 +43,6 @@ pub struct Context {
     /// tracing.
     ///
     pub extracted_span_context: Option<SpanContext>,
-    /// Whether the tracer has already generated stats for the dummy `aws.lambda` span
-    /// for this invocation (which only exists for universal instrumentation
-    /// languages).
-    ///
-    /// Set from the `Datadog-Client-Computed-Stats` header when the tracer's
-    /// dummy `aws.lambda` span (with `resource_name="dd-tracer-serverless-span"`) is received.
-    /// The value will be propagated to the same field of the real `aws.lambda` span generated by
-    /// the extension, which will be used to decide whether stats should be generated on extension
-    /// or backend.
-    pub client_computed_stats: bool,
 }
 
 /// Struct containing the information needed to reparent a span.
@@ -104,7 +94,6 @@ impl Default for Context {
             snapstart_restore_span: None,
             tracer_span: None,
             extracted_span_context: None,
-            client_computed_stats: false,
         }
     }
 }
@@ -519,12 +508,7 @@ impl ContextBuffer {
 
     /// Adds the tracer span to a `Context` in the buffer.
     ///
-    pub fn add_tracer_span(
-        &mut self,
-        request_id: &String,
-        tracer_span: &Span,
-        client_computed_stats: bool,
-    ) {
+    pub fn add_tracer_span(&mut self, request_id: &String, tracer_span: &Span) {
         if let Some(context) = self
             .buffer
             .iter_mut()
@@ -544,7 +528,6 @@ impl ContextBuffer {
                 .extend(tracer_span.metrics.clone());
 
             context.tracer_span = Some(tracer_span.clone());
-            context.client_computed_stats = client_computed_stats;
         } else {
             debug!("Could not add tracer span - context not found");
         }

bottlecap/src/lifecycle/invocation/processor.rs

@@ -49,6 +49,11 @@ use crate::{
 pub const MS_TO_NS: f64 = 1_000_000.0;
 pub const S_TO_MS: u64 = 1_000;
 pub const S_TO_NS: f64 = 1_000_000_000.0;
+/// Threshold for classifying a Lambda cold start as proactive initialization.
+///
+/// Proactive initialization is a Lambda optimization where the runtime pre-initializes
+/// a sandbox before any invocation is scheduled, to reduce cold start latency for future
+/// requests.
 pub const PROACTIVE_INITIALIZATION_THRESHOLD_MS: u64 = 10_000;
 
 pub const DATADOG_INVOCATION_ERROR_MESSAGE_KEY: &str = "x-datadog-invocation-error-msg";
@@ -96,6 +101,8 @@ pub struct Processor {
     /// logs agent. Decouples the trace agent from the logs agent: the trace agent sends spans
     /// to the lifecycle processor, which extracts durable context and relays it here.
     durable_context_tx: mpsc::Sender<DurableContextUpdate>,
+    /// Time of the `SnapStart` restore event, set when `PlatformRestoreStart` is received.
+    restore_time: Option<Instant>,
 }
 
 impl Processor {
@@ -137,6 +144,7 @@ impl Processor {
             active_invocations: 0,
             awaiting_first_invocation: false,
             durable_context_tx,
+            restore_time: None,
         }
     }
 
@@ -252,12 +260,35 @@ impl Processor {
 
         // If it's empty, then we are in a cold start
         if self.context_buffer.is_empty() {
-            let now = Instant::now();
-            let time_since_sandbox_init = now.duration_since(self.aws_config.sandbox_init_time);
-            if time_since_sandbox_init.as_millis() > PROACTIVE_INITIALIZATION_THRESHOLD_MS.into() {
-                proactive_initialization = true;
+            if self.aws_config.is_snapstart() {
+                match self.restore_time {
+                    None => {
+                        // PlatformRestoreStart hasn't arrived yet — restore and invoke
+                        // happened close together, so this is a cold start (not proactive).
+                        cold_start = true;
+                    }
+                    Some(restore_time) => {
+                        let now = Instant::now();
+                        let time_since_restore = now.duration_since(restore_time);
+                        if time_since_restore.as_millis()
+                            > PROACTIVE_INITIALIZATION_THRESHOLD_MS.into()
+                        {
+                            proactive_initialization = true;
+                        } else {
+                            cold_start = true;
+                        }
+                    }
+                }
             } else {
-                cold_start = true;
+                let now = Instant::now();
+                let time_since_sandbox_init = now.duration_since(self.aws_config.sandbox_init_time);
+                if time_since_sandbox_init.as_millis()
+                    > PROACTIVE_INITIALIZATION_THRESHOLD_MS.into()
+                {
+                    proactive_initialization = true;
+                } else {
+                    cold_start = true;
+                }
             }
 
             // Resolve runtime only once
@@ -383,6 +414,8 @@ impl Processor {
     /// This is used to create a `snapstart_restore` span, since this telemetry event does not
     /// provide a `request_id`, we try to guess which invocation is the restore similar to init.
     pub fn on_platform_restore_start(&mut self, time: DateTime<Utc>) {
+        self.restore_time = Some(Instant::now());
+
         let start_time: i64 = SystemTime::from(time)
             .duration_since(UNIX_EPOCH)
             .expect("time went backwards")
@@ -620,16 +653,9 @@ impl Processor {
         trace_sender: &Arc<SendingTraceProcessor>,
         context: Context,
     ) {
-        let client_computed_stats = context.client_computed_stats;
         let (traces, body_size) = self.get_ctx_spans(context);
-        self.send_spans(
-            traces,
-            body_size,
-            client_computed_stats,
-            tags_provider,
-            trace_sender,
-        )
-        .await;
+        self.send_spans(traces, body_size, tags_provider, trace_sender)
+            .await;
     }
 
     fn get_ctx_spans(&mut self, context: Context) -> (Vec<Span>, usize) {
@@ -694,7 +720,7 @@ impl Processor {
             let traces = vec![cold_start_span.clone()];
             let body_size = size_of_val(cold_start_span);
 
-            self.send_spans(traces, body_size, false, tags_provider, trace_sender)
+            self.send_spans(traces, body_size, tags_provider, trace_sender)
                 .await;
         }
     }
@@ -706,7 +732,6 @@ impl Processor {
         &mut self,
         traces: Vec<Span>,
         body_size: usize,
-        client_computed_stats: bool,
         tags_provider: &Arc<provider::Provider>,
         trace_sender: &Arc<SendingTraceProcessor>,
     ) {
@@ -719,7 +744,7 @@ impl Processor {
             tracer_version: "",
             container_id: "",
             client_computed_top_level: false,
-            client_computed_stats,
+            client_computed_stats: false,
             dropped_p0_traces: 0,
             dropped_p0_spans: 0,
         };
@@ -1366,10 +1391,9 @@ impl Processor {
     ///
     /// This is used to enrich the invocation span with additional metadata from the tracers
     /// top level span, since we discard the tracer span when we create the invocation span.
-    pub fn add_tracer_span(&mut self, span: &Span, client_computed_stats: bool) {
+    pub fn add_tracer_span(&mut self, span: &Span) {
         if let Some(request_id) = span.meta.get("request_id") {
-            self.context_buffer
-                .add_tracer_span(request_id, span, client_computed_stats);
+            self.context_buffer.add_tracer_span(request_id, span);
         }
     }
 
@@ -2180,98 +2204,6 @@ mod tests {
         );
     }
 
-    /// Verifies that `client_computed_stats` set on a context via `add_tracer_span` is
-    /// propagated all the way through `send_ctx_spans` to the `aws.lambda` payload sent
-    /// to the backend, so the extension does not generate duplicate stats.
-    #[tokio::test]
-    #[allow(clippy::unwrap_used)]
-    async fn test_client_computed_stats_propagated_to_aws_lambda_span() {
-        use crate::traces::stats_concentrator_service::StatsConcentratorService;
-        use crate::traces::stats_generator::StatsGenerator;
-        use libdd_trace_obfuscation::obfuscation_config::ObfuscationConfig;
-        use tokio::sync::mpsc;
-
-        let config = Arc::new(config::Config {
-            apm_dd_url: "https://trace.agent.datadoghq.com".to_string(),
-            ..config::Config::default()
-        });
-        let tags_provider = Arc::new(provider::Provider::new(
-            Arc::clone(&config),
-            LAMBDA_RUNTIME_SLUG.to_string(),
-            &HashMap::from([("function_arn".to_string(), "test-arn".to_string())]),
-        ));
-        let aws_config = Arc::new(AwsConfig {
-            region: "us-east-1".into(),
-            aws_lwa_proxy_lambda_runtime_api: None,
-            function_name: "test-function".into(),
-            sandbox_init_time: Instant::now(),
-            runtime_api: "***".into(),
-            exec_wrapper: None,
-            initialization_type: "on-demand".into(),
-        });
-        let (aggregator_service, aggregator_handle) =
-            AggregatorService::new(EMPTY_TAGS, 1024).expect("failed to create aggregator service");
-        tokio::spawn(aggregator_service.run());
-        let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(&config)));
-        let (durable_context_tx, _) = tokio::sync::mpsc::channel(1);
-        let mut p = Processor::new(
-            Arc::clone(&tags_provider),
-            Arc::clone(&config),
-            aws_config,
-            aggregator_handle,
-            propagator,
-            durable_context_tx,
-        );
-
-        let (trace_tx, mut trace_rx) = mpsc::channel(10);
-        let (stats_concentrator_service, stats_concentrator_handle) =
-            StatsConcentratorService::new(Arc::clone(&config));
-        tokio::spawn(stats_concentrator_service.run());
-        let trace_sender = Arc::new(SendingTraceProcessor {
-            appsec: None,
-            processor: Arc::new(trace_processor::ServerlessTraceProcessor {
-                obfuscation_config: Arc::new(
-                    ObfuscationConfig::new().expect("Failed to create ObfuscationConfig"),
-                ),
-            }),
-            trace_tx,
-            stats_generator: Arc::new(StatsGenerator::new(stats_concentrator_handle)),
-        });
-
-        let mut context = Context::from_request_id("req-1");
-        context.invocation_span.trace_id = 1;
-        context.invocation_span.span_id = 2;
-        context.client_computed_stats = true;
-
-        p.send_ctx_spans(&tags_provider, &trace_sender, context)
-            .await;
-
-        let payload = trace_rx
-            .recv()
-            .await
-            .expect("expected payload from trace_tx");
-        assert!(
-            payload.header_tags.client_computed_stats,
-            "client_computed_stats must be propagated to the aws.lambda span payload"
-        );
-
-        // Verify _dd.compute_stats is "0" in the built payload tags: client_computed_stats=true
-        // means the tracer has already computed stats, so neither extension nor backend should.
-        let send_data = payload.builder.build();
-        let libdd_trace_utils::tracer_payload::TracerPayloadCollection::V07(payloads) =
-            send_data.get_payloads()
-        else {
-            panic!("expected V07 payload");
-        };
-        for p in payloads {
-            assert_eq!(
-                p.tags.get(crate::tags::lambda::tags::COMPUTE_STATS_KEY),
-                Some(&"0".to_string()),
-                "_dd.compute_stats must be 0 when client_computed_stats is true"
-            );
-        }
-    }
-
     fn make_trace_sender(config: Arc<config::Config>) -> Arc<SendingTraceProcessor> {
         use libdd_trace_obfuscation::obfuscation_config::ObfuscationConfig;
         let (stats_concentrator_service, stats_concentrator_handle) =

bottlecap/src/lifecycle/invocation/processor_service.rs

@@ -109,7 +109,6 @@ pub enum ProcessorCommand {
     },
     AddTracerSpan {
         span: Box<Span>,
-        client_computed_stats: bool,
     },
     ForwardDurableContext {
         request_id: String,
@@ -378,12 +377,10 @@ impl InvocationProcessorHandle {
     pub async fn add_tracer_span(
         &self,
         span: Span,
-        client_computed_stats: bool,
     ) -> Result<(), mpsc::error::SendError<ProcessorCommand>> {
         self.sender
             .send(ProcessorCommand::AddTracerSpan {
                 span: Box::new(span),
-                client_computed_stats,
             })
             .await
     }
@@ -612,11 +609,8 @@ impl InvocationProcessorService {
                     let result = Ok(self.processor.set_cold_start_span_trace_id(trace_id));
                     let _ = response.send(result);
                 }
-                ProcessorCommand::AddTracerSpan {
-                    span,
-                    client_computed_stats,
-                } => {
-                    self.processor.add_tracer_span(&span, client_computed_stats);
+                ProcessorCommand::AddTracerSpan { span } => {
+                    self.processor.add_tracer_span(&span);
                 }
                 ProcessorCommand::ForwardDurableContext {
                     request_id,