[SVLS-8230] Fix SnapStart cold_start tag using restore_time

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

@@ -89,6 +89,8 @@
     /// Tracks whether if first invocation after init has been received in Managed Instance mode.
     /// Used to determine if we should search for the empty context on an invocation.
     awaiting_first_invocation: bool,
+    /// Time of the `SnapStart` restore event, set when `PlatformRestoreStart` is received.
+    restore_time: Option<DateTime<Utc>>,
 }
 
 impl Processor {
@@ -128,6 +130,7 @@
             dynamic_tags: HashMap::new(),
             active_invocations: 0,
             awaiting_first_invocation: false,
+            restore_time: None,
         }
     }
 
@@ -243,12 +246,34 @@
 
         // 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 = Utc::now();
+                        let time_since_restore = now.signed_duration_since(restore_time);
+                        if time_since_restore.num_seconds() > 10
+                        {
+                            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
@@ -374,6 +399,8 @@
     /// 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(time);
+
         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', () => {