diff --git a/package.json b/package.json index 56916873..46979721 100644 --- 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 index dae99007..ce74cdfc 100644 --- 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 index fd89e7eb..bf7d5a3b 100644 --- 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 index ad75b29e..3d7428df 100644 --- 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 index dd47ae84..f0308487 100644 --- 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);