diff --git a/.gitlab/datasources/test-suites.yaml b/.gitlab/datasources/test-suites.yaml index 56ab4026f..d4eb86259 100644 --- a/.gitlab/datasources/test-suites.yaml +++ b/.gitlab/datasources/test-suites.yaml @@ -7,3 +7,4 @@ test_suites: - name: oom - name: lmi-oom - name: payload-size + - name: durable-cold-start diff --git a/AGENTS.md b/AGENTS.md index bebb7c640..2e27f380a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,6 +10,11 @@ The Datadog Lambda Extension ("Bottlecap") is an AWS Lambda Extension written in All Rust code lives under `bottlecap/`. Run these from that directory unless noted. +## Code Style + +- Do not be overly defensive and add unnecessary checks or tests unless you have a good reason. +- Keep comments concise. Do not explain the history of a change or reference past code unless necessary. + ## Pull Requests - When creating a PR, follow the PR template at `.github/pull_request_template.md`. diff --git a/bottlecap/src/lifecycle/invocation/processor.rs b/bottlecap/src/lifecycle/invocation/processor.rs index 7d6943bdf..307484b56 100644 --- a/bottlecap/src/lifecycle/invocation/processor.rs +++ b/bottlecap/src/lifecycle/invocation/processor.rs @@ -111,6 +111,27 @@ pub struct Processor { /// on `platform.report`. This flag ensures whichever event arrives first wins and the other is skipped, /// preventing double counting. init_duration_metric_emitted: bool, + /// How the cold-start invocation metric in On-Demand mode has been resolved relative to + /// `platform.initStart`, since either the first invocation or `platform.initStart` can + /// arrive first (the Extensions API's `next()` call is unbuffered, while Telemetry API + /// delivery lags, so we've observed the invocation arriving first in production, but that's + /// not a hard ordering guarantee). + first_invocation_metric_status: FirstInvocationMetricStatus, +} + +/// See `Processor::first_invocation_metric_status`. +enum FirstInvocationMetricStatus { + /// Neither the first invocation nor `platform.initStart` has been seen yet. + Pending, + /// `platform.initStart` arrived first, so the durable tag is already known; the first + /// invocation's metric can be emitted immediately once it arrives. + InitStartSeen, + /// The first invocation arrived before `platform.initStart`; its metric is held back until + /// `platform.initStart` flushes it. + HeldBack(i64), + /// The first invocation's metric has been emitted, either immediately or after being held + /// back. Nothing left to track. + Resolved, } impl Processor { @@ -154,12 +175,25 @@ impl Processor { durable_context_tx, restore_time: None, init_duration_metric_emitted: false, + first_invocation_metric_status: FirstInvocationMetricStatus::Pending, } } /// Given a `request_id`, creates the context and adds the enhanced metric offsets to the context buffer. /// pub fn on_invoke_event(&mut self, request_id: String) { + // See `first_invocation_metric_status`. + let is_first_on_demand_invocation = !self.aws_config.is_managed_instance_mode() + && matches!( + self.first_invocation_metric_status, + FirstInvocationMetricStatus::Pending | FirstInvocationMetricStatus::InitStartSeen + ); + // Only hold back if `platform.initStart` hasn't already resolved the durable tag. + let should_hold_back_invocation_metric = matches!( + self.first_invocation_metric_status, + FirstInvocationMetricStatus::Pending + ) && is_first_on_demand_invocation; + // In Managed Instance mode, if awaiting the first invocation after init, find and update the empty context created on init start if self.aws_config.is_managed_instance_mode() && self.awaiting_first_invocation { if self @@ -224,8 +258,15 @@ impl Processor { .add_enhanced_metric_data(&request_id, enhanced_metric_offsets); } - // Increment the invocation metric - self.enhanced_metrics.increment_invocation_metric(timestamp); + // Hold back the metric for a cold start in On-Demand mode; see `first_invocation_metric_status`. + if should_hold_back_invocation_metric { + self.first_invocation_metric_status = FirstInvocationMetricStatus::HeldBack(timestamp); + } else { + self.enhanced_metrics.increment_invocation_metric(timestamp); + if is_first_on_demand_invocation { + self.first_invocation_metric_status = FirstInvocationMetricStatus::Resolved; + } + } self.enhanced_metrics.set_invoked_received(); // Both modes: check for buffered UniversalInstrumentationStart by request_id first. @@ -257,6 +298,16 @@ impl Processor { } } + /// Flushes the metric held back by `on_invoke_event`, if any. + fn flush_pending_first_invocation_metric(&mut self) { + if let FirstInvocationMetricStatus::HeldBack(timestamp) = + self.first_invocation_metric_status + { + self.enhanced_metrics.increment_invocation_metric(timestamp); + self.first_invocation_metric_status = FirstInvocationMetricStatus::Resolved; + } + } + /// On the first invocation, determine if it's a cold start or proactive init. /// /// For every other invocation, it will always be warm start. @@ -339,6 +390,15 @@ impl Processor { { self.enhanced_metrics.set_durable_function_tag(); } + if matches!( + self.first_invocation_metric_status, + FirstInvocationMetricStatus::Pending + ) { + self.first_invocation_metric_status = FirstInvocationMetricStatus::InitStartSeen; + } + // Flush the held-back metric now that the durable tag (if any) has been set. + self.flush_pending_first_invocation_metric(); + let start_time: i64 = SystemTime::from(time) .duration_since(UNIX_EPOCH) .expect("time went backwards") @@ -990,6 +1050,9 @@ impl Processor { } pub fn on_shutdown_event(&mut self) { + // Fallback flush in case the sandbox shuts down before it's otherwise flushed. + self.flush_pending_first_invocation_metric(); + let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .expect("unable to poll clock, unrecoverable") @@ -2095,6 +2158,111 @@ mod tests { ); } + #[tokio::test] + async fn test_on_invoke_event_cold_start_holds_invocation_metric_until_init_start() { + let mut processor = setup(); + + // Simulate the On-Demand race: the Invoke event reaches the processor before + // `platform.initStart`. The invocation metric must not be emitted yet. + processor.on_invoke_event("req-1".to_string()); + let FirstInvocationMetricStatus::HeldBack(held_back_timestamp) = + processor.first_invocation_metric_status + else { + panic!("Expected the cold start invocation metric to be held back") + }; + + // `platform.initStart` arrives afterwards and reports a durable runtime. + processor.on_platform_init_start( + Utc::now(), + Some("python:3.14.DurableFunction.v6".to_string()), + ); + + assert!( + matches!( + processor.first_invocation_metric_status, + FirstInvocationMetricStatus::Resolved + ), + "Expected the held-back invocation metric to be flushed by on_platform_init_start" + ); + + let runtime = processor + .runtime + .clone() + .expect("runtime should have been resolved by on_invoke_event"); + let ts = (held_back_timestamp / 10) * 10; + let expected_tags = dogstatsd::metric::SortedTags::parse(&format!( + "cold_start:true,durable_function:true,runtime:{runtime}" + )) + .ok(); + let entry = processor + .enhanced_metrics + .aggr_handle + .get_entry_by_id( + crate::metrics::enhanced::constants::INVOCATIONS_METRIC.into(), + expected_tags, + ts, + ) + .await + .unwrap(); + assert!( + entry.is_some(), + "Expected the flushed invocation metric to carry the durable_function:true tag" + ); + } + + #[tokio::test] + async fn test_on_invoke_event_does_not_hold_metric_if_init_start_already_processed() { + let mut processor = setup(); + + // Simulate the reverse order: `platform.initStart` arrives before the Invoke event + // and reports a durable runtime. + processor.on_platform_init_start( + Utc::now(), + Some("python:3.14.DurableFunction.v6".to_string()), + ); + + // The durable tag is already known, so the metric must be emitted immediately + // instead of being held back. + let now: i64 = std::time::UNIX_EPOCH + .elapsed() + .expect("unable to poll clock, unrecoverable") + .as_secs() + .try_into() + .unwrap_or_default(); + processor.on_invoke_event("req-1".to_string()); + assert!( + matches!( + processor.first_invocation_metric_status, + FirstInvocationMetricStatus::Resolved + ), + "Expected the invocation metric to be emitted immediately, not held back" + ); + + let runtime = processor + .runtime + .clone() + .expect("runtime should have been resolved by on_invoke_event"); + let ts = (now / 10) * 10; + let expected_tags = dogstatsd::metric::SortedTags::parse(&format!( + "cold_start:true,durable_function:true,runtime:{runtime}" + )) + .ok(); + let entry = processor + .enhanced_metrics + .aggr_handle + .get_entry_by_id( + crate::metrics::enhanced::constants::INVOCATIONS_METRIC.into(), + expected_tags, + ts, + ) + .await + .unwrap(); + assert!( + entry.is_some(), + "Expected the invocation metric to carry the durable_function:true tag" + ); + } + #[tokio::test] async fn test_init_duration_emitted_from_init_report() { let mut processor = setup(); diff --git a/integration-tests/bin/app.ts b/integration-tests/bin/app.ts index b34f6b910..9c0ee5e76 100644 --- a/integration-tests/bin/app.ts +++ b/integration-tests/bin/app.ts @@ -10,6 +10,7 @@ import {Oom} from '../lib/stacks/oom'; import {LmiOom} from '../lib/stacks/lmi-oom'; import {CustomMetrics} from '../lib/stacks/custom-metrics'; import {PayloadSize} from '../lib/stacks/payload-size'; +import {DurableColdStart} from '../lib/stacks/durable-cold-start'; import {AuthRoleStack} from '../lib/auth-role'; import {ACCOUNT, IDENTIFIER, REGION} from '../config'; import {CapacityProviderStack} from "../lib/capacity-provider"; @@ -54,6 +55,9 @@ const stacks = [ new PayloadSize(app, `${IDENTIFIER}-payload-size`, { env, }), + new DurableColdStart(app, `${IDENTIFIER}-durable-cold-start`, { + env, + }), ] // Tag all stacks so we can easily clean them up diff --git a/integration-tests/lib/stacks/durable-cold-start.ts b/integration-tests/lib/stacks/durable-cold-start.ts new file mode 100644 index 000000000..84ad87d09 --- /dev/null +++ b/integration-tests/lib/stacks/durable-cold-start.ts @@ -0,0 +1,66 @@ +import * as cdk from 'aws-cdk-lib'; +import * as lambda from 'aws-cdk-lib/aws-lambda'; +import { Construct } from 'constructs'; +import { + createLogGroup, + defaultDatadogEnvVariables, + defaultDatadogSecretPolicy, + getExtensionLayer, + getDefaultNodeLayer, + defaultNodeRuntime, +} from '../util'; + +// Deploys a durable and a non-durable function to test the cold-start durable_function tag. +export class DurableColdStart extends cdk.Stack { + constructor(scope: Construct, id: string, props: cdk.StackProps) { + super(scope, id, props); + + const extensionLayer = getExtensionLayer(this); + const nodeLayer = getDefaultNodeLayer(this); + + const durableFunctionName = `${id}-durable-lambda`; + const durableFunction = new lambda.Function(this, durableFunctionName, { + runtime: defaultNodeRuntime, + architecture: lambda.Architecture.ARM_64, + handler: '/opt/nodejs/node_modules/datadog-lambda-js/handler.handler', + code: lambda.Code.fromAsset('./lambda/default-node'), + functionName: durableFunctionName, + timeout: cdk.Duration.seconds(30), + memorySize: 256, + durableConfig: { + executionTimeout: cdk.Duration.minutes(5), + }, + environment: { + ...defaultDatadogEnvVariables, + DD_SERVICE: durableFunctionName, + DD_TRACE_ENABLED: 'true', + DD_LAMBDA_HANDLER: 'index.handler', + }, + logGroup: createLogGroup(this, durableFunctionName), + }); + durableFunction.addToRolePolicy(defaultDatadogSecretPolicy); + durableFunction.addLayers(extensionLayer); + durableFunction.addLayers(nodeLayer); + + const nonDurableFunctionName = `${id}-non-durable-lambda`; + const nonDurableFunction = new lambda.Function(this, nonDurableFunctionName, { + runtime: defaultNodeRuntime, + architecture: lambda.Architecture.ARM_64, + handler: '/opt/nodejs/node_modules/datadog-lambda-js/handler.handler', + code: lambda.Code.fromAsset('./lambda/default-node'), + functionName: nonDurableFunctionName, + timeout: cdk.Duration.seconds(30), + memorySize: 256, + environment: { + ...defaultDatadogEnvVariables, + DD_SERVICE: nonDurableFunctionName, + DD_TRACE_ENABLED: 'true', + DD_LAMBDA_HANDLER: 'index.handler', + }, + logGroup: createLogGroup(this, nonDurableFunctionName), + }); + nonDurableFunction.addToRolePolicy(defaultDatadogSecretPolicy); + nonDurableFunction.addLayers(extensionLayer); + nonDurableFunction.addLayers(nodeLayer); + } +} diff --git a/integration-tests/tests/durable-cold-start.test.ts b/integration-tests/tests/durable-cold-start.test.ts new file mode 100644 index 000000000..21d2d866e --- /dev/null +++ b/integration-tests/tests/durable-cold-start.test.ts @@ -0,0 +1,81 @@ +// Verifies the invocations metric has durable_function:true only for a durable function, and +// that the tag persists across a second (warm) invocation instead of only being set on cold start. +import { getMetricCount, getMetricCountWithTag } from './utils/datadog'; +import { forceColdStart, invokeLambda } from './utils/lambda'; +import { IDENTIFIER, DEFAULT_DATADOG_INDEXING_WAIT_MS } from '../config'; + +const stackName = `${IDENTIFIER}-durable-cold-start`; + +const INVOCATIONS_METRIC = 'aws.lambda.enhanced.invocations'; + +describe('Durable Function Cold-Start Metric Tag Integration Tests', () => { + let windowStart: number; + let windowEnd: number; + + const durableFunctionName = `${stackName}-durable-lambda`; + const nonDurableFunctionName = `${stackName}-non-durable-lambda`; + + // Durable functions reject unqualified-ARN invokes; `:$LATEST` avoids publishing a version. + const invokeBoth = () => + Promise.all([ + invokeLambda(`${durableFunctionName}:$LATEST`), + invokeLambda(nonDurableFunctionName), + ]); + + beforeAll(async () => { + const functionNames = [durableFunctionName, nonDurableFunctionName]; + await Promise.all(functionNames.map((fn) => forceColdStart(fn))); + + // Back up 60s so the metric's rollup bucket falls inside the query range. + windowStart = Date.now() - 60_000; + + // Invoke twice back-to-back (cold start, then warm) and wait for indexing once, rather than + // waiting after each invocation. Whether the tag was set on both invocations (rather than + // only the first) is checked below by comparing tagged vs. total invocation counts. + await invokeBoth(); + await invokeBoth(); + + await new Promise((resolve) => + setTimeout(resolve, DEFAULT_DATADOG_INDEXING_WAIT_MS), + ); + windowEnd = Date.now(); + + console.log('Lambdas invoked twice and indexing wait complete'); + }, 1800000); + + it('durable function invocation metric should have durable_function:true on every invocation', async () => { + const totalCount = await getMetricCount( + INVOCATIONS_METRIC, + durableFunctionName, + windowStart, + windowEnd, + ); + const taggedCount = await getMetricCountWithTag( + INVOCATIONS_METRIC, + durableFunctionName, + 'durable_function:true', + windowStart, + windowEnd, + ); + expect(totalCount).toBe(2); + expect(taggedCount).toBe(totalCount); + }); + + it('non-durable function invocation metric should NOT have durable_function:true on any invocation', async () => { + const totalCount = await getMetricCount( + INVOCATIONS_METRIC, + nonDurableFunctionName, + windowStart, + windowEnd, + ); + const taggedCount = await getMetricCountWithTag( + INVOCATIONS_METRIC, + nonDurableFunctionName, + 'durable_function:true', + windowStart, + windowEnd, + ); + expect(totalCount).toBe(2); + expect(taggedCount).toBe(0); + }); +}); diff --git a/integration-tests/tests/utils/datadog.ts b/integration-tests/tests/utils/datadog.ts index f8d2bcb31..cd72d1fed 100644 --- a/integration-tests/tests/utils/datadog.ts +++ b/integration-tests/tests/utils/datadog.ts @@ -370,6 +370,40 @@ export async function getMetricCount( return pointlist.reduce((acc, [, value]) => acc + (value || 0), 0); } +/** + * Same as `getMetricCount`, but filtered by an additional tag (e.g. "durable_function:true"). + * Used to check whether every emission of a metric over a window carried a given tag, by + * comparing against the untagged `getMetricCount` total for the same window. + */ +export async function getMetricCountWithTag( + metricName: string, + functionName: string, + tagFilter: string, + fromTime: number, + toTime: number, +): Promise { + const baseFunctionName = getServiceName(functionName).toLowerCase(); + const query = `sum:${metricName}{functionname:${baseFunctionName},${tagFilter}}.as_count()`; + + console.log(`Querying metric count with tag filter: ${query}`); + + const response = await requestWithRetry(() => datadogClient.get('/api/v1/query', { + params: { + query, + from: Math.floor(fromTime / 1000), + to: Math.floor(toTime / 1000), + }, + }), query); + + const series = response.data.series || []; + if (series.length === 0) { + return 0; + } + + const pointlist: [number, number][] = series[0].pointlist || []; + return pointlist.reduce((acc, [, value]) => acc + (value || 0), 0); +} + async function getMetrics( metricName: string, functionName: string,