rearrange and enhancements

Author: joeyzhao2018

src/trace/context/extractors/durable-execution.spec.ts

@@ -116,12 +116,12 @@ describe("DurableExecutionEventTraceExtractor", () => {
       CheckpointToken: "t-replay",
       InitialExecutionState: {
         Operations: [
-          {
-            Id: "op-1",
-            Name: "_dd_trace_context_0",
-            Status: "SUCCEEDED",
-            StepDetails: {
-              Result: JSON.stringify({
+            {
+              Id: "op-1",
+              Name: "_datadog_0",
+              Status: "SUCCEEDED",
+              StepDetails: {
+                Result: JSON.stringify({
                 "x-datadog-trace-id": "149750110124521191",
                 "x-datadog-parent-id": "538591322263933970",
                 "x-datadog-sampling-priority": "1",

src/trace/context/extractors/durable-execution.ts

@@ -110,19 +110,12 @@ export function getCheckpointToken(event: unknown): string | undefined {
 const TERMINAL_STATUSES = new Set(["SUCCEEDED", "FAILED", "CANCELLED", "STOPPED", "TIMED_OUT"]);
 
 const TRACE_CHECKPOINT_NAME_PREFIX = "_datadog_";
-const LEGACY_TRACE_CHECKPOINT_NAME_PREFIX = "_dd_trace_context_";
-const TRACE_CHECKPOINT_NAME_PREFIXES = [
-  TRACE_CHECKPOINT_NAME_PREFIX,
-  LEGACY_TRACE_CHECKPOINT_NAME_PREFIX,
-];
 
 function parseTraceCheckpointNumber(name: unknown): number | null {
   if (typeof name !== "string") return null;
 
-  const prefix = TRACE_CHECKPOINT_NAME_PREFIXES.find((candidate) => name.startsWith(candidate));
-  if (!prefix) return null;
-
-  const suffix = name.slice(prefix.length);
+  if (!name.startsWith(TRACE_CHECKPOINT_NAME_PREFIX)) return null;
+  const suffix = name.slice(TRACE_CHECKPOINT_NAME_PREFIX.length);
   const n = Number.parseInt(suffix, 10);
   if (Number.isNaN(n) || String(n) !== suffix) return null;
   return n;
@@ -141,7 +134,6 @@ function isTraceCheckpointName(name: unknown): boolean {
  * dd-trace-js plugin via `tracer.inject(span, 'http_headers', headers)` so the
  * payload is a standard HTTP-style header dict.
  *
- * Also accepts legacy `_dd_trace_context_{N}` names for compatibility.
  */
 function findLatestCheckpointHeaders(event: DurableExecutionEvent): Record<string, string> | null {
   const operations = event.InitialExecutionState?.Operations;
@@ -319,12 +311,14 @@ export function getCompletedOperationCount(event: unknown): number {
  * - no operation has terminal status
  * - operation count is <= 1
  *
+ * The created span is parented to the current aws.lambda span context.
+ *
  * Returns an object with { span, finish() } or null if not a durable execution.
  * Caller must call finish() when the invocation ends.
  */
 export function createDurableExecutionRootSpan(
   event: unknown,
-  extractedRootContext?: SpanContextWrapper | null,
+  parentSpanContext?: unknown,
 ): { span: any; finish: () => void } | null {
   if (!isDurableExecutionEvent(event)) {
     return null;
@@ -379,10 +373,9 @@ export function createDurableExecutionRootSpan(
     if (startTime !== undefined) {
       spanOptions.startTime = startTime;
     }
-    if (extractedRootContext?.spanContext) {
-      // Stay in the same trace as the upstream caller even when there is no
-      // active scope yet. aws.lambda will be parented to this span downstream.
-      spanOptions.childOf = extractedRootContext.spanContext;
+    if (parentSpanContext) {
+      // Root span is modeled as a child of aws.lambda.
+      spanOptions.childOf = parentSpanContext;
     }
 
     const span = tracer.startSpan("aws.durable-execution", spanOptions);

src/trace/listener.ts

@@ -27,7 +27,6 @@ import {
 } from "./durable-function-context";
 import { XrayService } from "./xray-service";
 import { AUTHORIZING_REQUEST_ID_HEADER } from "./context/extractors/http";
-import { createDurableExecutionRootSpan } from "./context/extractors/durable-execution";
 import { getSpanPointerAttributes, SpanPointerAttributes } from "../utils/span-pointers";
 export type TraceExtractor = (event: any, context: Context) => Promise<TraceContext> | TraceContext;
 
@@ -103,7 +102,6 @@ export class TraceListener {
   private wrappedCurrentSpan?: SpanWrapper;
   private triggerTags?: { [key: string]: string };
   private lambdaSpanParentContext?: SpanContext;
-  private durableRootSpan?: { span: any; finish: () => void };
   private spanPointerAttributesList: SpanPointerAttributes[] | undefined;
 
   public get currentTraceHeaders() {
@@ -145,25 +143,16 @@ export class TraceListener {
         traceSource: this.contextService.traceSource,
       });
     }
-    // Create the durable execution root span before everything else so later
-    // spans can parent correctly. Root creation is gated in
-    // createDurableExecutionRootSpan() and only happens for likely first
-    // invocations; replay invocations return null. The root span inherits
-    // trace_id and sampling from spanContextWrapper, so anything parented to
-    // it (aws.lambda below) joins the same trace automatically.
-    this.durableRootSpan = createDurableExecutionRootSpan(event, spanContextWrapper) ?? undefined;
-    const durableRootSpanContext = this.durableRootSpan?.span?.context();
-
     if (this.config.createInferredSpan) {
       this.inferredSpan = this.inferrer.createInferredSpan(
         event,
         context,
-        durableRootSpanContext || parentSpanContext,
+        parentSpanContext,
         this.config.encodeAuthorizerContext,
       );
     }
 
-    this.lambdaSpanParentContext = this.inferredSpan?.span || durableRootSpanContext || parentSpanContext;
+    this.lambdaSpanParentContext = this.inferredSpan?.span || parentSpanContext;
     this.context = context;
     const eventSource = parseEventSource(event);
     this.triggerTags = extractTriggerTags(event, context, eventSource);
@@ -320,12 +309,6 @@ export class TraceListener {
       this.injectAuthorizerSpan(result, event?.requestContext?.requestId, finishTime || Date.now());
     }
 
-    // Finish the durable execution root span after all other spans.
-    if (this.durableRootSpan) {
-      this.durableRootSpan.finish();
-      this.durableRootSpan = undefined;
-    }
-
     // Reset singletons and trace context
     this.stepFunctionContext = undefined;
     this.durableFunctionContext = undefined;