feat: Support async trace context extraction (#387)

Co-authored-by: Liza Pressman <35850765+lizapressman@users.noreply.github.com>

Author: astuyve

src/index.ts

@@ -138,7 +138,7 @@ export function datadog<TEvent, TResult>(
     currentTraceListener = traceListener;
 
     try {
-      traceListener.onStartInvocation(event, context);
+      await traceListener.onStartInvocation(event, context);
       await metricsListener.onStartInvocation(event);
       if (finalConfig.enhancedMetrics) {
         incrementInvocationsMetric(metricsListener, context);

src/trace/context.spec.ts

@@ -817,10 +817,10 @@ describe("extractTraceContext", () => {
     process.env["_X_AMZN_TRACE_ID"] = undefined;
     process.env[awsXrayDaemonAddressEnvVar] = undefined;
   });
-  it("returns trace read from header as highest priority with no extractor", () => {
+  it("returns trace read from header as highest priority with no extractor", async () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         headers: {
           "x-datadog-parent-id": "797643193680388251",
@@ -837,16 +837,53 @@ describe("extractTraceContext", () => {
       source: Source.Event,
     });
   });
-  it("returns an empty context when headers are null", () => {
-    const result = extractTraceContext(
+  it("returns an empty context when headers are null", async () => {
+    const result = await extractTraceContext(
       {
         headers: null,
       },
       {} as Context,
     );
     expect(result).toEqual(undefined);
   });
-  it("returns trace read from event with the extractor as the highest priority", () => {
+
+  it("returns trace read from event with an async extractor as the highest priority", async () => {
+    process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
+
+    const extractor = async (event: any, context: Context) => {
+      const traceID = event.foo[traceIDHeader];
+      const parentID = event.foo[parentIDHeader];
+      const sampledHeader = event.foo[samplingPriorityHeader];
+      const sampleMode = parseInt(sampledHeader, 10);
+
+      return {
+        parentID,
+        sampleMode,
+        source: Source.Event,
+        traceID,
+      };
+    };
+
+    const result = await extractTraceContext(
+      {
+        foo: {
+          "x-datadog-parent-id": "797643193680388251",
+          "x-datadog-sampling-priority": "2",
+          "x-datadog-trace-id": "4110911582297405551",
+        },
+      },
+      {} as Context,
+      extractor,
+    );
+    expect(result).toEqual({
+      parentID: "797643193680388251",
+      sampleMode: SampleMode.USER_KEEP,
+      traceID: "4110911582297405551",
+      source: Source.Event,
+    });
+  });
+
+  it("returns trace read from event with a synchronous extractor as the highest priority", async () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
     const extractor = (event: any, context: Context) => {
@@ -863,7 +900,7 @@ describe("extractTraceContext", () => {
       };
     };
 
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         foo: {
           "x-datadog-parent-id": "797643193680388251",
@@ -882,14 +919,14 @@ describe("extractTraceContext", () => {
     });
   });
 
-  it("handles gracefully errors in extractors", () => {
+  it("handles gracefully errors in extractors", async () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
     const extractor = (event: any, context: Context) => {
       throw new Error("test");
     };
 
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         foo: {
           "x-datadog-parent-id": "797643193680388251",
@@ -907,10 +944,10 @@ describe("extractTraceContext", () => {
       source: "xray",
     });
   });
-  it("returns trace read from SQS metadata as second highest priority", () => {
+  it("returns trace read from SQS metadata as second highest priority", async () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         Records: [
           {
@@ -945,7 +982,7 @@ describe("extractTraceContext", () => {
       source: Source.Event,
     });
   });
-  it("returns trace read from Lambda Context as third highest priority", () => {
+  it("returns trace read from Lambda Context as third highest priority", async () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
     const lambdaContext: Context = {
       clientContext: {
@@ -959,7 +996,7 @@ describe("extractTraceContext", () => {
         },
       },
     } as any;
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         Records: [
           {
@@ -993,34 +1030,34 @@ describe("extractTraceContext", () => {
       source: Source.Event,
     });
   });
-  it("returns trace read from env if no headers present", () => {
+  it("returns trace read from env if no headers present", async () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
-    const result = extractTraceContext({}, {} as Context);
+    const result = await extractTraceContext({}, {} as Context);
     expect(result).toEqual({
       parentID: "797643193680388254",
       sampleMode: SampleMode.USER_KEEP,
       traceID: "4110911582297405557",
       source: "xray",
     });
   });
-  it("returns trace read from env if no headers present", () => {
+  it("returns trace read from env if no headers present", async () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
-    const result = extractTraceContext({}, {} as Context);
+    const result = await extractTraceContext({}, {} as Context);
     expect(result).toEqual({
       parentID: "797643193680388254",
       sampleMode: SampleMode.USER_KEEP,
       traceID: "4110911582297405557",
       source: "xray",
     });
   });
-  it("adds datadog metadata segment to xray when trace context is in event", () => {
+  it("adds datadog metadata segment to xray when trace context is in event", async () => {
     jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
     process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
     process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000";
 
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         headers: {
           "x-datadog-parent-id": "797643193680388251",
@@ -1039,11 +1076,11 @@ describe("extractTraceContext", () => {
       {\\"id\\":\\"11111\\",\\"trace_id\\":\\"1-5e272390-8c398be037738dc042009320\\",\\"parent_id\\":\\"94ae789b969f1cc5\\",\\"name\\":\\"datadog-metadata\\",\\"start_time\\":1487076708,\\"end_time\\":1487076708,\\"type\\":\\"subsegment\\",\\"metadata\\":{\\"datadog\\":{\\"trace\\":{\\"parent-id\\":\\"797643193680388251\\",\\"sampling-priority\\":\\"2\\",\\"trace-id\\":\\"4110911582297405551\\"}}}}"
     `);
   });
-  it("skips adding datadog metadata to x-ray when daemon isn't present", () => {
+  it("skips adding datadog metadata to x-ray when daemon isn't present", async () => {
     jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
     process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
 
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         headers: {
           "x-datadog-parent-id": "797643193680388251",
@@ -1057,7 +1094,7 @@ describe("extractTraceContext", () => {
     expect(sentSegment).toBeUndefined();
   });
 
-  it("returns trace read from step functions event with the extractor as the highest priority", () => {
+  it("returns trace read from step functions event with the extractor as the highest priority", async () => {
     const stepFunctionEvent = {
       MyInput: "MyValue",
       Execution: {
@@ -1080,7 +1117,7 @@ describe("extractTraceContext", () => {
       },
     };
 
-    const result = extractTraceContext(stepFunctionEvent, {} as Context, undefined);
+    const result = await extractTraceContext(stepFunctionEvent, {} as Context, undefined);
     expect(result).toEqual({
       parentID: "4602916161841036335",
       sampleMode: 1,
@@ -1089,12 +1126,12 @@ describe("extractTraceContext", () => {
     });
   });
 
-  it("skips adding datadog metadata to x-ray when x-ray trace isn't sampled", () => {
+  it("skips adding datadog metadata to x-ray when x-ray trace isn't sampled", async () => {
     jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
     process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0";
     process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000";
 
-    const result = extractTraceContext(
+    const result = await extractTraceContext(
       {
         headers: {
           "x-datadog-parent-id": "797643193680388251",
@@ -1108,7 +1145,7 @@ describe("extractTraceContext", () => {
     expect(sentSegment).toBeUndefined();
   });
 
-  it("adds step function metadata to xray", () => {
+  it("adds step function metadata to xray", async () => {
     const stepFunctionEvent = {
       Execution: {
         Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
@@ -1134,7 +1171,7 @@ describe("extractTraceContext", () => {
     process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
     process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000";
 
-    extractTraceContext(stepFunctionEvent, {} as Context);
+    await extractTraceContext(stepFunctionEvent, {} as Context);
     expect(sentSegment instanceof Buffer).toBeTruthy();
 
     expect(closedSocket).toBeTruthy();

src/trace/context.ts

@@ -102,17 +102,17 @@ export function deterministicMd5HashToBigIntString(s: string): string {
  * Reads the trace context from either an incoming lambda event, or the current xray segment.
  * @param event An incoming lambda event. This must have incoming trace headers in order to be read.
  */
-export function extractTraceContext(
+export async function extractTraceContext(
   event: any,
   context: Context,
   extractor?: TraceExtractor,
   decodeAuthorizerContext: boolean = true,
-): TraceContext | undefined {
+): Promise<TraceContext | undefined> {
   let trace;
 
   if (extractor) {
     try {
-      trace = extractor(event, context);
+      trace = await extractor(event, context);
       logDebug(`extracted trace context from the custom extractor`, { trace });
     } catch (error) {
       if (error instanceof Error) {

src/trace/listener.spec.ts

@@ -106,7 +106,7 @@ describe("TraceListener", () => {
 
   it("wraps dd-trace span around invocation", async () => {
     const listener = new TraceListener(defaultConfig);
-    listener.onStartInvocation({}, context as any);
+    await listener.onStartInvocation({}, context as any);
     const unwrappedFunc = () => {};
     const wrappedFunc = listener.onWrap(unwrappedFunc);
     wrappedFunc();
@@ -141,7 +141,7 @@ describe("TraceListener", () => {
       "x-datadog-trace-id": "4110911582297405551",
     };
     mockTraceSource = Source.Event;
-    listener.onStartInvocation({}, context as any);
+    await listener.onStartInvocation({}, context as any);
     const unwrappedFunc = () => {};
     const wrappedFunc = listener.onWrap(unwrappedFunc);
     wrappedFunc();
@@ -179,7 +179,7 @@ describe("TraceListener", () => {
     };
     mockTraceSource = Source.Xray;
 
-    listener.onStartInvocation({}, context as any);
+    await listener.onStartInvocation({}, context as any);
     const unwrappedFunc = () => {};
     const wrappedFunc = listener.onWrap(unwrappedFunc);
     wrappedFunc();
@@ -215,7 +215,7 @@ describe("TraceListener", () => {
     };
     mockTraceSource = Source.Xray;
 
-    listener.onStartInvocation({}, context as any);
+    await listener.onStartInvocation({}, context as any);
     const unwrappedFunc = () => {};
     const wrappedFunc = listener.onWrap(unwrappedFunc);
     wrappedFunc();
@@ -246,7 +246,7 @@ describe("TraceListener", () => {
 
   it("wraps dd-trace span around invocation, with function alias", async () => {
     const listener = new TraceListener(defaultConfig);
-    listener.onStartInvocation({}, contextWithFunctionAlias as any);
+    await listener.onStartInvocation({}, contextWithFunctionAlias as any);
     const unwrappedFunc = () => {};
     const wrappedFunc = listener.onWrap(unwrappedFunc);
     wrappedFunc();
@@ -275,7 +275,7 @@ describe("TraceListener", () => {
 
   it("wraps dd-trace span around invocation, with function version", async () => {
     const listener = new TraceListener(defaultConfig);
-    listener.onStartInvocation({}, contextWithFunctionVersion as any);
+    await listener.onStartInvocation({}, contextWithFunctionVersion as any);
     const unwrappedFunc = () => {};
     const wrappedFunc = listener.onWrap(unwrappedFunc);
     wrappedFunc();

src/trace/listener.ts

@@ -20,7 +20,7 @@ import { SpanContext, TraceOptions, TracerWrapper } from "./tracer-wrapper";
 import { SpanInferrer } from "./span-inferrer";
 import { SpanWrapper } from "./span-wrapper";
 import { getTraceTree } from "../runtime/index";
-export type TraceExtractor = (event: any, context: Context) => TraceContext;
+export type TraceExtractor = (event: any, context: Context) => Promise<TraceContext> | TraceContext;
 
 export interface TraceConfig {
   /**
@@ -88,7 +88,7 @@ export class TraceListener {
     this.inferrer = new SpanInferrer(this.tracerWrapper);
   }
 
-  public onStartInvocation(event: any, context: Context) {
+  public async onStartInvocation(event: any, context: Context) {
     const tracerInitialized = this.tracerWrapper.isTracerAvailable;
     if (this.config.injectLogContext) {
       patchConsole(console, this.contextService);
@@ -104,7 +104,7 @@ export class TraceListener {
     } else {
       logDebug("Not patching HTTP libraries", { autoPatchHTTP: this.config.autoPatchHTTP, tracerInitialized });
     }
-    const rootTraceHeaders = this.contextService.extractHeadersFromContext(
+    const rootTraceHeaders = await this.contextService.extractHeadersFromContext(
       event,
       context,
       this.config.traceExtractor,

src/trace/trace-context-service.ts

@@ -21,13 +21,13 @@ export class TraceContextService {
 
   constructor(private tracerWrapper: TracerWrapper) {}
 
-  extractHeadersFromContext(
+  async extractHeadersFromContext(
     event: any,
     context: Context,
     extractor?: TraceExtractor,
     decodeAuthorizerContext: boolean = true,
-  ): Partial<TraceHeaders> | undefined {
-    this.rootTraceContext = extractTraceContext(event, context, extractor, decodeAuthorizerContext);
+  ): Promise<Partial<TraceHeaders> | undefined> {
+    this.rootTraceContext = await extractTraceContext(event, context, extractor, decodeAuthorizerContext);
     return this.currentTraceHeaders;
   }