clean up AI complicated stuff...

Author: joeyzhao2018

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

@@ -12,7 +12,6 @@
  * so we just hand the resulting header dict back to `tracer.extract` here.
  */
 
-import { randomBytes } from "crypto";
 import { logDebug } from "../../../utils";
 import { SpanContextWrapper } from "../../span-context-wrapper";
 import { TracerWrapper } from "../../tracer-wrapper";
@@ -107,21 +106,6 @@ export function getCheckpointToken(event: unknown): string | undefined {
   return event.CheckpointToken;
 }
 
-function generateRandomPositiveId(): string {
-  const bytes = randomBytes(8);
-  bytes[0] = bytes[0] & 0x7f; // keep it positive int64
-  const value = bufferToBigInt(bytes);
-  return value === 0n ? "1" : value.toString(10);
-}
-
-function bufferToBigInt(buf: Buffer): bigint {
-  let result = 0n;
-  for (let i = 0; i < buf.length; i++) {
-    result = (result << 8n) | BigInt(buf[i]);
-  }
-  return result;
-}
-
 // Terminal operation statuses that indicate an operation has completed
 const TERMINAL_STATUSES = new Set(["SUCCEEDED", "FAILED", "CANCELLED", "STOPPED", "TIMED_OUT"]);
 
@@ -362,25 +346,21 @@ export function createDurableExecutionRootSpan(
     return null;
   }
 
-  const rootSpanId = extractedRootContext?.toSpanId() || generateRandomPositiveId();
-
-  // Determine consistent start_time from the first operation's StartTimestamp
-  // StartTimestamp is unix milliseconds from the durable execution SDK
+  // Use the first operation's StartTimestamp (unix milliseconds) so the root
+  // span's start time matches the actual start of the durable execution.
   let startTime: number | undefined;
-  const replayOperations = event.InitialExecutionState?.Operations;
-  if (replayOperations && replayOperations.length > 0) {
-    const firstStartTs = replayOperations[0].StartTimestamp;
+  if (operations && operations.length > 0) {
+    const firstStartTs = operations[0].StartTimestamp;
     if (firstStartTs != null) {
       const parsed = Number(firstStartTs);
       if (!isNaN(parsed)) {
-        startTime = parsed;  // already in millis, dd-trace startSpan expects millis
+        startTime = parsed;
       }
     }
   }
 
   try {
     const tracer = require("dd-trace");
-    const id = require("dd-trace/packages/dd-trace/src/id");
 
     const serviceName = process.env.DD_DURABLE_EXECUTION_SERVICE || "aws.durable-execution";
     const resourceName = executionArn.includes(":") ? executionArn.split(":").pop() : executionArn;
@@ -392,48 +372,22 @@ export function createDurableExecutionRootSpan(
         "resource.name": resourceName,
         "durable.execution_arn": executionArn,
         "durable.is_root_span": true,
-        "durable.invocation_count": replayOperations?.length ?? 0,
+        "durable.invocation_count": operations?.length ?? 0,
       },
     };
 
     if (startTime !== undefined) {
       spanOptions.startTime = startTime;
     }
     if (extractedRootContext?.spanContext) {
-      // Ensure the durable root span stays in the same trace as the extracted
-      // durable invocation context even when there is no active scope.
+      // 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;
     }
 
     const span = tracer.startSpan("aws.durable-execution", spanOptions);
 
-    // Use the extracted durable root span_id when available to keep the
-    // durable root identity stable with propagated checkpoint context.
-    try {
-      if (rootSpanId) {
-        span.context()._spanId = id(rootSpanId, 10);
-      }
-    } catch (e) {
-      logDebug(`Failed to set durable root span_id: ${e}`);
-    }
-
-    // Fix parent_id: when an extracted span context exists, tracer.startSpan()
-    // inherits its span_id as parent_id and we just overwrote our own span_id
-    // to match — that would self-parent. The root span's parent should be the
-    // upstream caller (if any) or 0 (true root).
-    try {
-      const upstreamHeaders = findUpstreamHeaders(event as DurableExecutionEvent);
-      const upstreamParentId = upstreamHeaders?.["x-datadog-parent-id"];
-      if (upstreamParentId) {
-        span.context()._parentId = id(String(upstreamParentId), 10);
-      } else {
-        span.context()._parentId = id("0", 10);
-      }
-    } catch (e) {
-      logDebug(`Failed to set root span parent_id: ${e}`);
-    }
-
-    logDebug(`Created root execution span: span_id=${rootSpanId ?? "auto"}, start_time=${startTime}`);
+    logDebug(`Created root execution span: start_time=${startTime}`);
 
     return {
       span,

src/trace/listener.ts

@@ -148,19 +148,22 @@ export class TraceListener {
     // 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.
+    // 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,
-        parentSpanContext,
+        durableRootSpanContext || parentSpanContext,
         this.config.encodeAuthorizerContext,
       );
     }
 
-    this.lambdaSpanParentContext = this.inferredSpan?.span || parentSpanContext;
+    this.lambdaSpanParentContext = this.inferredSpan?.span || durableRootSpanContext || parentSpanContext;
     this.context = context;
     const eventSource = parseEventSource(event);
     this.triggerTags = extractTriggerTags(event, context, eventSource);