Summary

• Fixes aws.lambda.cold_start spans having a different trace ID than the main aws.lambda span

Root Cause

tracer.current_trace_context() was being called after span.finish(), which could return a stale or empty context since the span was no longer active in the context provider.

How Node.js Correctly Handles This

The Node.js tracer library (datadog-lambda-js) passes the live parent span object to ColdStartTracer before calling finish():

// datadog-lambda-js/src/trace/listener.ts (lines 196-199)
const coldStartConfig: ColdStartTracerConfig = {
  parentSpan: didFunctionColdStart() ? this.inferredSpan || this.wrappedCurrentSpan : this.wrappedCurrentSpan,
  // ...
};

Then in cold-start-tracer.ts, the cold_start span is created as a child of this live span:

// datadog-lambda-js/src/trace/cold-start-tracer.ts (lines 49-64)
if (parentSpan) {
  options.childOf = parentSpan.span;  // Inherits correct trace_id
}

Changes

Moved trace context capture to before any spans are finished, ensuring aws.lambda.cold_start and aws.lambda.load spans inherit the correct trace ID. This mirrors the Node.js approach of capturing context while spans are still active.

Test plan

• Integration tests in datadog-lambda-extension validate this fix (all 22 tests passing)
• Python cold_start span tests now pass with the correct trace ID

🤖 Generated with Claude Code