Merge branch 'main' into zarir/base_service

Author: duncanista

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/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")

integration-tests/tests/snapstart.test.ts

@@ -101,6 +101,23 @@ describe('Snapstart Integration Tests', () => {
         );
         expect(coldStartSpan).toBeUndefined();
       });
+
+      it('should have aws.lambda span with cold_start=true', () => {
+        const result = getRestoreInvocation();
+        expect(result).toBeDefined();
+        const trace = result.traces![0];
+        const awsLambdaSpan = trace.spans.find((span: any) =>
+          span.attributes.operation_name === 'aws.lambda'
+        );
+        expect(awsLambdaSpan).toBeDefined();
+        expect(awsLambdaSpan).toMatchObject({
+          attributes: {
+            custom: {
+              cold_start: 'true'
+            }
+          }
+        });
+      });
     });
 
     describe('second invocation (warm)', () => {
@@ -146,6 +163,23 @@ describe('Snapstart Integration Tests', () => {
         );
         expect(coldStartSpan).toBeUndefined();
       });
+
+      it('should have aws.lambda span with cold_start=false', () => {
+        const result = getWarmInvocation();
+        expect(result).toBeDefined();
+        const trace = result.traces![0];
+        const awsLambdaSpan = trace.spans.find((span: any) =>
+          span.attributes.operation_name === 'aws.lambda'
+        );
+        expect(awsLambdaSpan).toBeDefined();
+        expect(awsLambdaSpan).toMatchObject({
+          attributes: {
+            custom: {
+              cold_start: 'false'
+            }
+          }
+        });
+      });
     });
 
     describe('trace isolation', () => {