Extract trace context from AWSTraceHeader in SQS java upstream case (#…

…511) * extrace trace context from AWSTraceHeader in SQS java upstream case --------- Co-authored-by: jordan gonzález <30836115+duncanista@users.noreply.github.com>
DataDog · Mar 18, 2024 · 4160e9a · 4160e9a
1 parent cb2656d
commit 4160e9a
Show file tree

Hide file tree

Showing 5 changed files with 164 additions and 32 deletions.
diff --git a/package.json b/package.json
@@ -39,7 +39,6 @@
     "typescript": "^4.3.2"
   },
   "dependencies": {
-    "bignumber.js": "^9.0.1",
     "dc-polyfill": "^0.1.3",
     "hot-shots": "8.5.0",
     "promise-retry": "^2.0.1",

diff --git a/src/trace/context/extractors/sqs.spec.ts b/src/trace/context/extractors/sqs.spec.ts
@@ -126,5 +126,53 @@ describe("SQSEventTraceExtractor", () => {
       const traceContext = extractor.extract(payload);
       expect(traceContext).toBeNull();
     });
+
+    it("extracts trace context from AWSTraceHeader with valid payload", () => {
+      mockSpanContext = {
+        toTraceId: () => "625397077193750208",
+        toSpanId: () => "6538302989251745223",
+        _sampling: {
+          priority: "1",
+        },
+      };
+      const tracerWrapper = new TracerWrapper();
+      const payload: SQSEvent = {
+        Records: [
+          {
+            body: "Hello world",
+            attributes: {
+              ApproximateReceiveCount: "1",
+              SentTimestamp: "1605544528092",
+              SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer",
+              ApproximateFirstReceiveTimestamp: "1605544528094",
+              AWSTraceHeader: "Root=1-65f2f78c-0000000008addb5405b376c0;Parent=5abcb7ed643995c7;Sampled=1",
+            },
+            messageAttributes: {},
+            eventSource: "aws:sqs",
+            eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue",
+            awsRegion: "eu-west-1",
+            messageId: "foo",
+            md5OfBody: "x",
+            receiptHandle: "x",
+          },
+        ],
+      };
+
+      const extractor = new SQSEventTraceExtractor(tracerWrapper);
+
+      const traceContext = extractor.extract(payload);
+      expect(traceContext).not.toBeNull();
+
+      expect(spyTracerWrapper).toHaveBeenCalledWith({
+        "x-datadog-parent-id": "6538302989251745223",
+        "x-datadog-sampling-priority": "1",
+        "x-datadog-trace-id": "625397077193750208",
+      });
+
+      expect(traceContext?.toTraceId()).toBe("625397077193750208");
+      expect(traceContext?.toSpanId()).toBe("6538302989251745223");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+    });
   });
 });
diff --git a/src/trace/context/extractors/sqs.ts b/src/trace/context/extractors/sqs.ts
@@ -3,19 +3,29 @@ import { EventTraceExtractor } from "../extractor";
 import { TracerWrapper } from "../../tracer-wrapper";
 import { logDebug } from "../../../utils";
 import { SpanContextWrapper } from "../../span-context-wrapper";
+import { XrayService } from "../../xray-service";
 
 export class SQSEventTraceExtractor implements EventTraceExtractor {
   constructor(private tracerWrapper: TracerWrapper) {}
 
   extract(event: SQSEvent): SpanContextWrapper | null {
-    const headers = event?.Records?.[0]?.messageAttributes?._datadog?.stringValue;
-    if (headers === undefined) return null;
-
     try {
-      const traceContext = this.tracerWrapper.extract(JSON.parse(headers));
-      if (traceContext === null) return null;
+      let parsedHeaders;
+      const headers = event?.Records?.[0]?.messageAttributes?._datadog?.stringValue;
+      if (headers !== undefined) {
+        parsedHeaders = JSON.parse(headers);
+      } else if (event?.Records?.[0]?.attributes?.AWSTraceHeader !== undefined) {
+        parsedHeaders = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader);
+      }
+      if (!parsedHeaders) return null;
+
+      const traceContext = this.tracerWrapper.extract(parsedHeaders);
+      if (traceContext === null) {
+        logDebug("Failed to extract trace context from parsed headers", { parsedHeaders, event });
+        return null;
+      }
 
-      logDebug(`Extracted trace context from SQS event`, { traceContext, event });
+      logDebug("Extracted trace context from SQS event", { traceContext, event });
       return traceContext;
     } catch (error) {
       if (error instanceof Error) {

diff --git a/src/trace/xray-service.spec.ts b/src/trace/xray-service.spec.ts
@@ -1,3 +1,8 @@
+import {
+  DATADOG_SAMPLING_PRIORITY_HEADER,
+  DATADOG_TRACE_ID_HEADER,
+  DATADOG_PARENT_ID_HEADER,
+} from "./context/extractor";
 import { SampleMode } from "./trace-context-service";
 import { XrayService } from "./xray-service";
 
@@ -340,4 +345,46 @@ describe("XrayService", () => {
       expect(traceId).toBeUndefined();
     });
   });
+
+  describe("parseAWSTraceHeader", () => {
+    it("parses AWS trace header correctly", () => {
+      const awsTraceHeader = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
+      const xrayHeaders = XrayService.parseAWSTraceHeader(awsTraceHeader);
+      expect(xrayHeaders).toEqual({
+        parentId: "94ae789b969f1cc5",
+        sampled: "1",
+        traceId: "1-5e272390-8c398be037738dc042009320",
+      });
+    });
+    it.each(["Root=1-5e272390-8c398be037738dc042009320", "Root=1-65f2f78c-0000000008addb5405b376c0;Parent;Sampled"])(
+      "returns undefined when AWS trace header is malformatted",
+      (awsTraceHeader) => {
+        const xrayHeaders = XrayService.parseAWSTraceHeader(awsTraceHeader);
+        expect(xrayHeaders).toBeUndefined();
+      },
+    );
+  });
+  describe("extraceDDContextFromAWSTraceHeader", () => {
+    it("extracts Datadog trace context from AWS trace header", () => {
+      const awsTraceId = "Root=1-65f2f78c-0000000008addb5405b376c0;Parent=5abcb7ed643995c7;Sampled=1";
+      const ddTraceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceId);
+
+      expect(ddTraceContext).toEqual({
+        [DATADOG_TRACE_ID_HEADER]: "625397077193750208",
+        [DATADOG_PARENT_ID_HEADER]: "6538302989251745223",
+        [DATADOG_SAMPLING_PRIORITY_HEADER]: "1",
+      });
+    });
+
+    it("returns null when AWS trace header is NOT injected by dd-trace", () => {
+      const awsTraceId = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
+      const ddTraceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceId);
+      expect(ddTraceContext).toBeNull();
+    });
+    it("returns null when AWS trace header cannot be parsed", () => {
+      const awsTraceId = "Root=1-5e272390-8c398be037738dc042009320;;";
+      const ddTraceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceId);
+      expect(ddTraceContext).toBeNull();
+    });
+  });
 });
diff --git a/src/trace/xray-service.ts b/src/trace/xray-service.ts
@@ -1,14 +1,19 @@
 import { randomBytes } from "crypto";
 import { logDebug } from "../utils";
 import { SampleMode, TraceContext, TraceSource } from "./trace-context-service";
-import BigNumber from "bignumber.js";
 import { Socket, createSocket } from "dgram";
 import { SpanContextWrapper } from "./span-context-wrapper";
 import { StepFunctionContext } from "./step-function-service";
+import {
+  DATADOG_TRACE_ID_HEADER,
+  DATADOG_PARENT_ID_HEADER,
+  DATADOG_SAMPLING_PRIORITY_HEADER,
+  DatadogTraceHeaders,
+} from "./context/extractor";
 
 const AMZN_TRACE_ID_ENV_VAR = "_X_AMZN_TRACE_ID";
 const AWS_XRAY_DAEMON_ADDRESS_ENV_VAR = "AWS_XRAY_DAEMON_ADDRESS";
-
+const DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000";
 interface XrayTraceHeader {
   traceId: string;
   parentId: string;
@@ -70,16 +75,9 @@ export class XrayService {
     });
   }
 
-  private parseTraceContextHeader(): XrayTraceHeader | undefined {
-    const header = process.env[AMZN_TRACE_ID_ENV_VAR];
-    if (header === undefined) {
-      logDebug("Couldn't read Xray trace header from env");
-      return;
-    }
-
-    // Example: Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
-    logDebug(`Reading Xray trace context from env var ${header}`);
-    const [root, parent, _sampled] = header.split(";");
+  // Example: Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
+  public static parseAWSTraceHeader(awsTraceHeader: string): XrayTraceHeader | undefined {
+    const [root, parent, _sampled] = awsTraceHeader.split(";");
     if (parent === undefined || _sampled === undefined) return;
 
     const [, traceId] = root.split("=");
@@ -94,6 +92,18 @@ export class XrayService {
     };
   }
 
+  private parseTraceContextHeader(): XrayTraceHeader | undefined {
+    const header = process.env[AMZN_TRACE_ID_ENV_VAR];
+    if (header === undefined) {
+      logDebug("Couldn't read Xray trace header from env");
+      return;
+    }
+
+    // Example: Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
+    logDebug(`Reading Xray trace context from env var ${header}`);
+    return XrayService.parseAWSTraceHeader(header);
+  }
+
   private convertToSampleMode(xraySampled: number): SampleMode {
     return xraySampled === 1 ? SampleMode.USER_KEEP : SampleMode.USER_REJECT;
   }
@@ -172,11 +182,12 @@ export class XrayService {
 
   private convertToParentId(xrayParentId: string): string | undefined {
     if (xrayParentId.length !== 16) return;
-
-    const hex = new BigNumber(xrayParentId, 16);
-    if (hex.isNaN()) return;
-
-    return hex.toString(10);
+    try {
+      return BigInt("0x" + xrayParentId).toString(10);
+    } catch (_) {
+      logDebug(`Failed to convert Xray Parent Id ${xrayParentId}`);
+      return undefined;
+    }
   }
 
   private convertToTraceId(xrayTraceId: string): string | undefined {
@@ -187,13 +198,30 @@ export class XrayService {
     if (lastPart.length !== 24) return;
 
     // We want to turn the last 63 bits into a decimal number in a string representation
-    // Unfortunately, all numbers in javascript are represented by float64 bit numbers, which
-    // means we can't parse 64 bit integers accurately.
-    const hex = new BigNumber(lastPart, 16);
-    if (hex.isNaN()) return;
-
-    // Toggle off the 64th bit
-    const last63Bits = hex.mod(new BigNumber("8000000000000000", 16));
-    return last63Bits.toString(10);
+    try {
+      return (BigInt("0x" + lastPart) % BigInt("0x8000000000000000")).toString(10); // mod by 2^63 will leave us with the last 63 bits
+    } catch (_) {
+      logDebug(`Failed to convert Xray Trace Id ${lastPart}`);
+      return undefined;
+    }
+  }
+
+  public static extraceDDContextFromAWSTraceHeader(amznTraceId: string): DatadogTraceHeaders | null {
+    const awsContext = XrayService.parseAWSTraceHeader(amznTraceId);
+    if (awsContext === undefined) {
+      return null;
+    }
+    const traceIdParts = awsContext.traceId.split("-");
+    if (traceIdParts && traceIdParts.length > 2 && traceIdParts[2].startsWith(DD_TRACE_JAVA_TRACE_ID_PADDING)) {
+      // This AWSTraceHeader contains Datadog injected trace context
+      return {
+        [DATADOG_TRACE_ID_HEADER]: hexStrToDecimalStr(traceIdParts[2].substring(8)),
+        [DATADOG_PARENT_ID_HEADER]: hexStrToDecimalStr(awsContext.parentId),
+        [DATADOG_SAMPLING_PRIORITY_HEADER]: awsContext.sampled,
+      };
+    }
+    return null;
   }
 }
+
+const hexStrToDecimalStr = (hexString: string): string => BigInt("0x" + hexString).toString(10);