extrace trace context from AWSTraceHeader in SQS java upstream case

Author: joeyzhao2018

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",

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");
+    });
   });
 });

src/trace/context/extractors/sqs.ts

@@ -3,18 +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;
+      var prepared_headers;
+      const headers = event?.Records?.[0]?.messageAttributes?._datadog?.stringValue;
+      if (headers !== undefined) {
+        prepared_headers = JSON.parse(headers);
+      } else {
+        if (event?.Records?.[0]?.attributes?.AWSTraceHeader) {
+          prepared_headers = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader);
+        }
+      }
 
+      if (!prepared_headers) return null;
+      const traceContext = this.tracerWrapper.extract(prepared_headers);
+      if (traceContext === null) {
+        logDebug("Unable to extract the injected trace context from event");
+        return null;
+      }
       logDebug(`Extracted trace context from SQS event`, { traceContext, event });
       return traceContext;
     } catch (error) {

src/trace/xray-service.ts

@@ -1,10 +1,15 @@
 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";
@@ -70,16 +75,8 @@ 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(";");
+  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 +91,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 +181,11 @@ 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 (_) {
+      return undefined;
+    }
   }
 
   private convertToTraceId(xrayTraceId: string): string | undefined {
@@ -187,13 +196,29 @@ 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);
+    } catch (_) {
+      return undefined;
+    }
+  }
+
+  public static extraceDDContextFromAWSTraceHeader(amznTraceId: string): DatadogTraceHeaders | null {
+    var aws_context = XrayService.parseAWSTraceHeader(amznTraceId);
+    if (!aws_context) {
+      return null;
+    }
+    const trace_id_parts = aws_context.traceId.split("-");
+    if (trace_id_parts && trace_id_parts.length > 2 && trace_id_parts[2].startsWith("00000000")) {
+      // This AWSTraceHeader contains Datadog injected trace context
+      return {
+        [DATADOG_TRACE_ID_HEADER]: hexStrToDecimalStr(trace_id_parts[2].substring(8)),
+        [DATADOG_PARENT_ID_HEADER]: hexStrToDecimalStr(aws_context.parentId),
+        [DATADOG_SAMPLING_PRIORITY_HEADER]: aws_context.sampled,
+      };
+    }
+    return null;
   }
 }
+
+const hexStrToDecimalStr = (hexString: string): string => BigInt("0x" + hexString).toString(10);