diff --git a/bottlecap/src/config/aws.rs b/bottlecap/src/config/aws.rs index d1af40a8f..284f91b50 100644 --- a/bottlecap/src/config/aws.rs +++ b/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)] diff --git a/bottlecap/src/lifecycle/invocation/processor.rs b/bottlecap/src/lifecycle/invocation/processor.rs index b38ada8ca..46e974ea0 100644 --- a/bottlecap/src/lifecycle/invocation/processor.rs +++ b/bottlecap/src/lifecycle/invocation/processor.rs @@ -89,6 +89,8 @@ pub struct Processor { /// 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>, } impl Processor { @@ -128,6 +130,7 @@ impl Processor { dynamic_tags: HashMap::new(), active_invocations: 0, awaiting_first_invocation: false, + restore_time: None, } } @@ -243,12 +246,34 @@ 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 = 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 @@ 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) { + self.restore_time = Some(time); + let start_time: i64 = SystemTime::from(time) .duration_since(UNIX_EPOCH) .expect("time went backwards") diff --git a/integration-tests/tests/snapstart.test.ts b/integration-tests/tests/snapstart.test.ts index ccbd74111..7dbbc7f50 100644 --- a/integration-tests/tests/snapstart.test.ts +++ b/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', () => {