extract trace context from client context

Author: Czechh

src/trace/context.spec.ts

@@ -19,6 +19,7 @@ import {
   readStepFunctionContextFromEvent,
   readTraceFromSQSEvent,
   readTraceFromHTTPEvent,
+  readTraceFromLambdaContext,
 } from "./context";
 
 let sentSegment: any;
@@ -388,6 +389,41 @@ describe("readTraceFromSQSEvent", () => {
   });
 });
 
+describe("readTraceFromLambdaContext", () => {
+  it("can read from lambda context source", () => {
+    const result = readTraceFromLambdaContext({
+      clientContext: {
+        custom: {
+          _datadog: {
+            "x-datadog-trace-id":"666",
+            "x-datadog-parent-id":"777",
+            "x-datadog-sampled":"1",
+            "x-datadog-sampling-priority":"1"
+          }
+        }
+      }
+    });
+    expect(result).toEqual({
+      parentID: "777",
+      sampleMode: SampleMode.AUTO_KEEP,
+      traceID: "666",
+      source: Source.Event,
+    });
+  });
+  it("can handle no `custom` key", () => {
+    const result = readTraceFromLambdaContext({
+      clientContext: {
+        foo: 'bar'
+      }
+    });
+    expect(result).toBeUndefined();
+  });
+  it("can handle no context", () => {
+    const result = readTraceFromLambdaContext(undefined);
+    expect(result).toBeUndefined();
+  });
+});
+
 describe("readStepFunctionContextFromEvent", () => {
   const stepFunctionEvent = {
     dd: {
@@ -532,7 +568,7 @@ describe("extractTraceContext", () => {
         "x-datadog-sampling-priority": "2",
         "x-datadog-trace-id": "4110911582297405551",
       },
-    });
+    },  {});
     expect(result).toEqual({
       parentID: "797643193680388251",
       sampleMode: SampleMode.USER_KEEP,
@@ -567,7 +603,53 @@ describe("extractTraceContext", () => {
           awsRegion: "sa-east-1",
         },
       ],
+    }, {});
+    expect(result).toEqual({
+      parentID: "3369753143434738315",
+      sampleMode: SampleMode.AUTO_KEEP,
+      traceID: "4555236104497098341",
+      source: Source.Event,
     });
+  });
+  it("returns trace read from Lambda Context as third highest priority", () => {
+    process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
+    const lambdaContext = {
+      clientContext: {
+        custom: {
+          _datadog: {
+            "x-datadog-trace-id":"4555236104497098341",
+            "x-datadog-parent-id":"3369753143434738315",
+            "x-datadog-sampled":"1",
+            "x-datadog-sampling-priority":"1"
+          }
+        }
+      }
+    };
+    const result = extractTraceContext({
+      Records: [
+        {
+          body: "Hello world",
+          attributes: {
+            ApproximateReceiveCount: "1",
+            SentTimestamp: "1605544528092",
+            SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer",
+            ApproximateFirstReceiveTimestamp: "1605544528094",
+          },
+          messageAttributes: {
+            _datadog: {
+              stringValue:
+                '{"x-datadog-parent-id":"666","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}',
+              stringListValues: [],
+              binaryListValues: [],
+              dataType: "String",
+            },
+          },
+          eventSource: "aws:sqs",
+          eventSourceARN: "arn:aws:sqs:sa-east-1:601427279990:metal-queue",
+          awsRegion: "sa-east-1",
+        },
+      ],
+    }, lambdaContext);
     expect(result).toEqual({
       parentID: "3369753143434738315",
       sampleMode: SampleMode.AUTO_KEEP,
@@ -578,7 +660,7 @@ describe("extractTraceContext", () => {
   it("returns trace read from env if no headers present", () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
-    const result = extractTraceContext({});
+    const result = extractTraceContext({}, {});
     expect(result).toEqual({
       parentID: "797643193680388254",
       sampleMode: SampleMode.USER_KEEP,
@@ -589,7 +671,7 @@ describe("extractTraceContext", () => {
   it("returns trace read from env if no headers present", () => {
     process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1";
 
-    const result = extractTraceContext({});
+    const result = extractTraceContext({}, {});
     expect(result).toEqual({
       parentID: "797643193680388254",
       sampleMode: SampleMode.USER_KEEP,
@@ -608,7 +690,7 @@ describe("extractTraceContext", () => {
         "x-datadog-sampling-priority": "2",
         "x-datadog-trace-id": "4110911582297405551",
       },
-    });
+    }, {});
 
     expect(sentSegment instanceof Buffer).toBeTruthy();
     expect(closedSocket).toBeTruthy();
@@ -628,7 +710,7 @@ describe("extractTraceContext", () => {
         "x-datadog-sampling-priority": "2",
         "x-datadog-trace-id": "4110911582297405551",
       },
-    });
+    }, {});
 
     expect(sentSegment).toBeUndefined();
   });
@@ -654,7 +736,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);
+    extractTraceContext(stepFunctionEvent, {});
 
     expect(sentSegment instanceof Buffer).toBeTruthy();
 

src/trace/context.ts

@@ -1,3 +1,4 @@
+import { Context } from "aws-lambda";
 import { BigNumber } from "bignumber.js";
 import { randomBytes } from "crypto";
 import { createSocket, Socket } from "dgram";
@@ -47,8 +48,13 @@ function isSQSEvent(event: any): event is SQSEvent {
  * 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(event: any) {
-  const trace = readTraceFromEvent(event);
+export function extractTraceContext(event: any, context: any) {
+  let trace = readTraceFromEvent(event);
+
+  if (trace === undefined) {
+    trace = readTraceFromLambdaContext(context)
+  }
+
   const stepFuncContext = readStepFunctionContextFromEvent(event);
   if (stepFuncContext) {
     try {
@@ -188,6 +194,39 @@ export function readTraceFromSQSEvent(event: SQSEvent): TraceContext | undefined
   return;
 }
 
+export function readTraceFromLambdaContext(context: any): TraceContext | undefined {
+  if (!context || typeof context !== "object") {
+    return;
+  }
+
+  const traceData = context.clientContext?.custom?._datadog;
+
+  if (!traceData || typeof traceData !== "object") {
+    return;
+  }
+
+  const traceID = traceData[traceIDHeader];
+  if (typeof traceID !== "string") {
+    return;
+  }
+  const parentID = traceData[parentIDHeader];
+  if (typeof parentID !== "string") {
+    return;
+  }
+  const sampledHeader = traceData[samplingPriorityHeader];
+  if (typeof sampledHeader !== "string") {
+    return;
+  }
+  const sampleMode = parseInt(sampledHeader, 10);
+
+  return {
+    parentID,
+    sampleMode,
+    source: Source.Event,
+    traceID,
+  };
+}
+
 export function readTraceFromHTTPEvent(event: any): TraceContext | undefined {
   const headers = event.headers;
   const lowerCaseHeaders: { [key: string]: string } = {};

src/trace/listener.ts

@@ -60,7 +60,7 @@ export class TraceListener {
     }
 
     this.context = context;
-    this.contextService.rootTraceContext = extractTraceContext(event);
+    this.contextService.rootTraceContext = extractTraceContext(event, context);
     this.stepFunctionContext = readStepFunctionContextFromEvent(event);
   }