fix: close any existing active span when extracting the tracecontext (#688)

* close any existing active span when extracting the tracecontext

Author: joeyzhao2018

src/trace/trace-context-service.spec.ts

@@ -14,6 +14,7 @@ describe("TraceContextService", () => {
     mockXRayShouldThrow = false;
     const traceWrapper = {
       traceContext: () => spanContextWrapper,
+      closeScope: jest.fn(),
     };
     traceContextService = new TraceContextService(traceWrapper as any, {} as any);
   });
@@ -126,4 +127,54 @@ describe("TraceContextService", () => {
     expect(result?.toTraceId()).toBe("newTraceId");
     expect(traceContextService.traceSource).toBe("event");
   });
+
+  it("should not leak dd-trace context from previous invocation when extract is called", async () => {
+    // Simulate dd-trace having a stale active span from a previous invocation (warm start scenario)
+    const staleDdTraceContext = {
+      toTraceId: () => "staleTraceId_999",
+      toSpanId: () => "staleSpanId_888",
+      sampleMode: () => 1,
+      source: TraceSource.DdTrace,
+      spanContext: {},
+    };
+
+    // Mock tracerWrapper that returns stale context initially, then null after closeScope is called
+    let traceContextValue: any = staleDdTraceContext;
+    const mockCloseScopeFn = jest.fn(() => {
+      // After closeScope is called, traceContext should return null
+      traceContextValue = null;
+    });
+
+    const mockTracerWrapper = {
+      traceContext: jest.fn(() => traceContextValue),
+      closeScope: mockCloseScopeFn,
+    };
+
+    const service = new TraceContextService(mockTracerWrapper as any, {} as any);
+
+    // Mock the extractor to return a NEW context for the current invocation
+    const newEventContext = {
+      toTraceId: () => "newTraceId_123",
+      toSpanId: () => "newSpanId_456",
+      sampleMode: () => 2,
+      source: TraceSource.Event,
+      spanContext: {},
+    };
+    const mockExtract = jest.fn().mockResolvedValue(newEventContext);
+    service["traceExtractor"] = { extract: mockExtract } as any;
+
+    // Call extract for the new invocation
+    await service.extract({}, {} as any);
+
+    // Verify that closeScope was called to clear the stale context
+    expect(mockCloseScopeFn).toHaveBeenCalled();
+
+    // After the fix: currentTraceHeaders should return the NEW context from the event
+    // not the stale dd-trace context from the previous invocation
+    const headers = service.currentTraceHeaders;
+
+    expect(headers["x-datadog-trace-id"]).toBe("newTraceId_123");
+    expect(headers["x-datadog-parent-id"]).toBe("newSpanId_456");
+    expect(headers["x-datadog-sampling-priority"]).toBe("2");
+  });
 });

src/trace/trace-context-service.ts

@@ -50,12 +50,13 @@ export class TraceContextService {
   }
 
   async extract(event: any, context: Context): Promise<SpanContextWrapper | null> {
-    // Reset trace context from previous invocation to prevent caching
+    // Reset trace context and close dd-trace scope to prevent stale context from previous invocation due to unfinished spans
     this.rootTraceContext = null;
+    this.tracerWrapper.closeScope();
 
     this.rootTraceContext = await this.traceExtractor?.extract(event, context);
-
-    return this.currentTraceContext;
+    // Return the extracted context, not the current context which may not be related to the event or context
+    return this.rootTraceContext;
   }
 
   get currentTraceHeaders(): Partial<DatadogTraceHeaders> {
@@ -70,8 +71,6 @@ export class TraceContextService {
   }
 
   get currentTraceContext(): SpanContextWrapper | null {
-    if (this.rootTraceContext === null) return null;
-
     const traceContext = this.rootTraceContext;
     const currentDatadogContext = this.tracerWrapper.traceContext();
     if (currentDatadogContext) {

src/trace/tracer-wrapper.spec.ts

@@ -127,4 +127,35 @@ describe("TracerWrapper", () => {
 
     expect(mockDataStreamsCheckpointer.setConsumeCheckpoint).toHaveBeenCalledWith(eventType, arn, contextJson, false);
   });
+
+  it("should finish active span when closing scope to prevent context leakage", () => {
+    const mockFinishFn = jest.fn();
+    mockSpan = {
+      context: () => ({
+        toSpanId: () => "1234",
+        toTraceId: () => "45678",
+        _sampling: {
+          priority: "2",
+        },
+      }),
+      finish: mockFinishFn,
+    };
+
+    const wrapper = new TracerWrapper();
+    wrapper.closeScope();
+
+    expect(mockFinishFn).toHaveBeenCalled();
+  });
+
+  it("should not error when closing scope with no active span", () => {
+    mockSpan = null;
+    const wrapper = new TracerWrapper();
+    expect(() => wrapper.closeScope()).not.toThrow();
+  });
+
+  it("should not error when closing scope with tracer unavailable", () => {
+    mockNoTracer = true;
+    const wrapper = new TracerWrapper();
+    expect(() => wrapper.closeScope()).not.toThrow();
+  });
 });

src/trace/tracer-wrapper.ts

@@ -1,4 +1,4 @@
-import { logDebug } from "../utils";
+import { logDebug, logWarning } from "../utils";
 import { SpanContextWrapper } from "./span-context-wrapper";
 import { TraceSource } from "./trace-context-service";
 
@@ -94,6 +94,20 @@ export class TracerWrapper {
     return new SpanContextWrapper(span.context(), TraceSource.DdTrace);
   }
 
+  public closeScope(): void {
+    try {
+      const activeSpan = this.currentSpan;
+      if (activeSpan && typeof activeSpan.finish === "function") {
+        logDebug("Detected stale span from previous invocation, finishing it to prevent trace context leakage");
+        activeSpan.finish();
+      }
+    } catch (err) {
+      if (err instanceof Object || err instanceof Error) {
+        logDebug("Failed to close dd-trace scope", err);
+      }
+    }
+  }
+
   public injectSpan(span: SpanContext): any {
     const dest = {};
     this.tracer.inject(span, "text_map", dest);