diff --git a/src/trace/context/extractors/sns-sqs.spec.ts b/src/trace/context/extractors/sns-sqs.spec.ts index 950f8bb8..33d673eb 100644 --- a/src/trace/context/extractors/sns-sqs.spec.ts +++ b/src/trace/context/extractors/sns-sqs.spec.ts @@ -20,6 +20,7 @@ describe("SNSSQSEventTraceExtractor", () => { describe("extract", () => { beforeEach(() => { mockSpanContext = null; + spyTracerWrapper.mockClear(); }); afterEach(() => { @@ -42,7 +43,7 @@ describe("SNSSQSEventTraceExtractor", () => { messageId: "64812b68-4d9b-4dca-b3fb-9b18f255ee51", receiptHandle: "AQEBER6aRkfG8092GvkL7FRwCwbQ7LLDW9Tlk/CembqHe+suS2kfFxXiukomvaIN61QoyQMoRgWuV52SDkiQno2u+5hP64BDbmw+e/KR9ayvIfHJ3M6RfyQLaWNWm3hDFBCKTnBMVIxtdx0N9epZZewyokjKcrNYtmCghFgTCvZzsQkowi5rnoHAVHJ3je1c3bDnQ1KLrZFgajDnootYXDwEPuMq5FIxrf4EzTe0S7S+rnRm+GaQfeBLBVAY6dASL9usV3/AFRqDtaI7GKI+0F2NCgLlqj49VlPRz4ldhkGknYlKTZTluAqALWLJS62/J1GQo53Cs3nneJcmu5ajB2zzmhhRXoXINEkLhCD5ujZfcsw9H4xqW69Or4ECvlqx14bUU2rtMIW0QM2p7pEeXnyocymQv6m1te113eYWTVmaJ4I=", - body: '{\n "Type" : "Notification",\n "MessageId" : "0a0ab23e-4861-5447-82b7-e8094ff3e332",\n "TopicArn" : "arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA",\n "Message" : "{\\"hello\\":\\"harv\\",\\"nice of you to join us\\":\\"david\\",\\"anotherThing\\":{\\"foo\\":\\"bar\\",\\"blah\\":null,\\"harv\\":123},\\"vals\\":[{\\"thingOne\\":1},{\\"thingTwo\\":2}],\\"ajTimestamp\\":1639777617957}",\n "Timestamp" : "2021-12-17T21:46:58.040Z",\n "SignatureVersion" : "1",\n "Signature" : "FR35/7E8C3LHEVk/rC4XxXlXwV/5mNkFNPgDhHSnJ2I6hIoSrTROAm7h5xm1PuBkAeFDvq0zofw91ouk9zZyvhdrMLFIIgrjEyNayRmEffmoEAkzLFUsgtQX7MmTl644r4NuWiM0Oiz7jueRvIcKXcZr7Nc6GJcWV1ymec8oOmuHNMisnPMxI07LIQVYSyAfv6P9r2jEWMVIukRoCzwTnRk4bUUYhPSGHI7OC3AsxxXBbv8snqTrLM/4z2rXCf6jHCKNxWeLlm9/45PphCkEyx5BWS4/71KaoMWUWy8+6CCsy+uF3XTCVmvSEYLyEwTSzOY+vCUjazrRW93498i70g==",\n "SigningCertUrl" : "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-7ff5318490ec183fbaddaa2a969abfda.pem",\n "UnsubscribeUrl" : "https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA:1290f550-9a8a-4e8f-a900-8f5f96dcddda",\n "MessageAttributes" : {\n "_datadog" : {"Type":"String","Value":"{\\"x-datadog-trace-id\\":\\"2776434475358637757\\",\\"x-datadog-parent-id\\":\\"4493917105238181843\\",\\"x-datadog-sampled\\":\\"1\\",\\"x-datadog-sampling-priority\\":\\"1\\"}"}\n }\n}', + body: '{\n "Type" : "Notification",\n "MessageId" : "0a0ab23e-4861-5447-82b7-e8094ff3e332",\n "TopicArn" : "arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA",\n "Message" : "{\\"hello\\":\\"harv\\",\\"nice of you to join us\\":\\"david\\",\\"anotherThing\\":{\\"foo\\":\\"bar\\",\\"blah\\":null,\\"harv\\":123},\\"vals\\":[{\\"thingOne\\":1},{\\"thingTwo\\":2}],\\"ajTimestamp\\":1639777617957}",\n "Timestamp" : "2021-12-17T21:46:58.040Z",\n "SignatureVersion" : "1",\n "Signature" : "FR35/7E8C3LHEVk/rC4XxXlXwV/5mNkFNPgDhHSnJ2I6hIoSrTROAm7h5xm1PuBkAeFDvq0zofw91ouk9zZyvhdrMLFIIgrjEyNayRmEffmoEAkzLFUsgtQX7MmTl644r4NuWiM0Oiz7jueRvIcKXcZr7Nc6GJcWV1ymec8oOmuHNMisnPMxI07LIQVYSyAfv6P9r2jEWMVIukRoCzwTnRk4bUUYhPSGHI7OC3AsxxXBbv8snqTrLM/4z2rXCf6jHCKNxWeLlm9/45PphCkEyx5BWS4/71KaoMWUWy8+6CCsy+uF3XTCVmvSEYLyEwTSzOY+vCUjazrRW93498i70g==",\n "SigningCertUrl" : "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-************************33ab7e69.pem",\n "UnsubscribeUrl" : "https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:601427279990:js-library-test-dev-demoTopic-15WGUVRCBMPAA:1290f550-9a8a-4e8f-a900-8f5f96dcddda",\n "MessageAttributes" : {\n "_datadog" : {"Type":"String","Value":"{\\"x-datadog-trace-id\\":\\"2776434475358637757\\",\\"x-datadog-parent-id\\":\\"4493917105238181843\\",\\"x-datadog-sampled\\":\\"1\\",\\"x-datadog-sampling-priority\\":\\"1\\"}"}\n }\n}', attributes: { ApproximateReceiveCount: "1", SentTimestamp: "1639777618130", @@ -171,5 +172,45 @@ describe("SNSSQSEventTraceExtractor", () => { const traceContext = extractor.extract(payload); expect(traceContext).toBeNull(); }); + it("extracts trace context from AWSTraceHeader with valid payload", () => { + const tracerWrapper = new TracerWrapper(); + const payload = { + Records: [ + { + messageId: "40304cba-041c-4284-808e-8d8692c6cba1", + receiptHandle: + "AQEBisC4xIjrp5W72r64fIf86ilRqadl6sbkmqbx41BI7P5ov0dJSNaNG97dnmOqSMD+MhXTvvC7HL2i+3viGrC8iNacZXSr9zzZvpFLFYO6jphDRTNJdkEZqSbKUoN5c/Nri5FjA4X52q4pCvW1esADJ2ZcMQQuQ19gsKAEQ0VGkHpH+BeceQtoFc3XT7uJboykfUA6iWT4TyNdJS+O4119ZBdN3U1jZ3PUn8mmTSi+SkTiXPBD9ywu6X8VzkGahueT+P7tJQTZ27mbPKhfrt3kvbFD6z7lqBNQyAPoqHzAThGC3VbZOxth3iqf7kjsFccmSJsxvsBzcVpF6nmobf6dpxwnZTEIrlNpQGrBgoePIHrpWfC6UG6aRTc4zWc30VY6hcg09WjCNGI81KwDfNMDAdJknOhsbY3HtvRQkQncbXgsYXgDDJG3PdIUoI2YScLeWEBMwE/HPCWk0X3K6McczRIHw3PfLaS2eVpjzNlq9I4=", + body: '{\n "Type" : "Notification",\n "MessageId" : "d968a20b-73c3-5389-ae7a-fcb844faf1c7",\n "TopicArn" : "arn:aws:sns:us-west-2:425362996713:DdTraceXLambda-snssqschecksNestedStacksnssqschecksNestedStackResource58F786C6-11NORKTA1JFML-snsProducerJavaForPythonNonRawsnssqsproducerjavaforpythonnonrawtopicDDBAB6EA-ZBb8uCZzkS0S",\n "Message" : "hello from DdTraceXLambda-snssqschec-snssqsproducerjavaforpyt-z0t7yDk3zWt1",\n "Timestamp" : "2024-05-06T19:52:25.181Z",\n "SignatureVersion" : "1",\n "Signature" : "pZIIa0Ae49vUPVSZcR4XCt9K2gMWYBjIJCZbQJo6URKLJOcC4yJNXVzOQAb80tG10lgOq+gMahLaTcuJ5+yFr3LtK/8nl7mdeP7aH6V2VoRubmJuc7P2WUixhubve577MfFMjp1LrkQaa5D/ken6yOjjgxRy32GazYAUEeQ9duldSAuu3omfsljWnZSHoeHkpbVkCrp/KyNGDQKrf+pFxxuxb9yqUzbHa8H80zS9fwOEsBuSqlbyK2Mj68wneqSeuRcZ30l5xyJ82vVjyXEukNcSkt5OcZOYFGqxotIY7MKTr2nrkiFOJAiRsOK34eQyk7eVdWRRfyoxCHVpnImT1Q==",\n "SigningCertURL" : "https://sns.us-west-2.amazonaws.com/SimpleNotificationService-************************33ab7e69",\n "UnsubscribeURL" : "https://sns.us-west-2.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:us-west-2:425362996713:DdTraceXLambda-snssqschecksNestedStacksnssqschecksNestedStackResource58F786C6-11NORKTA1JFML-snsProducerJavaForPythonNonRawsnssqsproducerjavaforpythonnonrawtopicDDBAB6EA-ZBb8uCZzkS0S:d44c53e8-538c-472f-89e1-89c131d9cd26",\n "MessageAttributes" : {\n "JOEYTEST" : {"Type":"String","Value":"test"},\n "JOEYTEST2" : {"Type":"String","Value":"test2"}\n }\n}', + attributes: { + ApproximateReceiveCount: "1", + AWSTraceHeader: "Root=1-663934f8-0000000045c5da17458c9910;Parent=18f37c8f541d052f;Sampled=1", + SentTimestamp: "1715025145210", + SenderId: "AIDAIYLAVTDLUXBIEIX46", + ApproximateFirstReceiveTimestamp: "1715025145220", + }, + messageAttributes: {}, + md5OfBody: "************************33ab7e69", + eventSource: "aws:sqs", + eventSourceARN: + "arn:aws:sqs:us-west-2:425362996713:DdTraceXLambda-snssqschecksNested-snsProducerJavaForPythonNonRawsns-1vws7QPqW8e6", + awsRegion: "us-west-2", + }, + ], + }; + + const extractor = new SNSSQSEventTraceExtractor(tracerWrapper); + + const traceContext = extractor.extract(payload); + expect(traceContext).not.toBeNull(); + + // Should not use ddtracer extractor. Because 1. it's an unnecessary extra step and + // 2. More importantly, DD_TRACE_PROPAGATION_STYLE could cause extraction fail + expect(spyTracerWrapper).not.toHaveBeenCalled(); + + expect(traceContext?.toTraceId()).toBe("5027664352514971920"); + expect(traceContext?.toSpanId()).toBe("1797917631284315439"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); }); }); diff --git a/src/trace/context/extractors/sns-sqs.ts b/src/trace/context/extractors/sns-sqs.ts index 73586227..e68ee6a2 100644 --- a/src/trace/context/extractors/sns-sqs.ts +++ b/src/trace/context/extractors/sns-sqs.ts @@ -3,32 +3,46 @@ 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 SNSSQSEventTraceExtractor implements EventTraceExtractor { constructor(private tracerWrapper: TracerWrapper) {} extract(event: SQSEvent): SpanContextWrapper | null { - const body = event?.Records?.[0]?.body; - if (body === undefined) return null; - try { - const parsedBody = JSON.parse(body) as SNSMessage; - const messageAttribute = parsedBody?.MessageAttributes?._datadog; - if (messageAttribute.Value === undefined) return null; + // First try to extract trace context from message attributes + if (event?.Records?.[0]?.body) { + const parsedBody = JSON.parse(event?.Records?.[0]?.body) as SNSMessage; + const messageAttribute = parsedBody?.MessageAttributes?._datadog; + if (messageAttribute?.Value) { + let headers; + if (messageAttribute.Type === "String") { + headers = JSON.parse(messageAttribute.Value); + } else { + const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii"); + headers = JSON.parse(decodedValue); + } - let headers; - if (messageAttribute.Type === "String") { - headers = JSON.parse(messageAttribute.Value); - } else { - const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii"); - headers = JSON.parse(decodedValue); + const traceContext = this.tracerWrapper.extract(headers); + if (traceContext) { + logDebug("Extracted trace context from SNS-SQS event"); + return traceContext; + } else { + logDebug("Failed to extract trace context from SNS-SQS event"); + } + } + } + // Then try to extract trace context from attributes.AWSTraceHeader. (Upstream Java apps can + // pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case) + if (event?.Records?.[0]?.attributes?.AWSTraceHeader !== undefined) { + const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader); + if (traceContext) { + logDebug("Extracted trace context from SNS-SQS event attributes.AWSTraceHeader"); + return traceContext; + } else { + logDebug("No Datadog trace context found from SNS-SQS event attributes.AWSTraceHeader"); + } } - - const traceContext = this.tracerWrapper.extract(headers); - if (traceContext === null) return null; - - logDebug("Extracted trace context from SNS-SQS event", { traceContext, event }); - return traceContext; } catch (error) { if (error instanceof Error) { logDebug("Unable to extract trace context from SNS-SQS event", error); diff --git a/src/trace/context/extractors/sns.spec.ts b/src/trace/context/extractors/sns.spec.ts index bd2c69c3..aadc29f6 100644 --- a/src/trace/context/extractors/sns.spec.ts +++ b/src/trace/context/extractors/sns.spec.ts @@ -21,6 +21,7 @@ describe("SNSEventTraceExtractor", () => { describe("extract", () => { beforeEach(() => { mockSpanContext = null; + spyTracerWrapper.mockClear(); }); afterEach(() => { @@ -55,7 +56,7 @@ describe("SNSEventTraceExtractor", () => { Signature: "mzp2Ou0fASw4LYRxY6SSww7qFfofn4luCJBRaTjLpQ5uhwhsAUKdyLz9VPD+/dlRbi1ImsWtIZ7A+wxj1oV7Z2Gyu/N4RpGalae37+jTluDS7AhjgcD7Bs4bgQtFkCfMFEwbhICQfukLLzbwbgczZ4NTPn6zj5o28c5NBKSJMYSnLz82ohw77GgnZ/m26E32ZQNW4+VCEMINg9Ne2rHstwPWRXPr5xGTrx8jH8CNUZnVpFVfhU8o+OSeAdpzm2l99grHIo7qPhekERxANz6QHynMlhdzD3UNSgc3oZkamZban/NEKd4MKJzgNQdNOYVj3Kw6eF2ZweEoBQ5sSFK5fQ==", SigningCertUrl: - "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-7ff5318490ec183fbaddaa2a969abfda.pem", + "https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-************************33ab7e69.pem", UnsubscribeUrl: "https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:601427279990:aj-js-library-test-dev-solo-topic:1bd19208-a99a-46d9-8398-f90f8699c641", MessageAttributes: { @@ -203,5 +204,57 @@ describe("SNSEventTraceExtractor", () => { const traceContext = extractor.extract(payload); expect(traceContext).toBeNull(); }); + + it("extracts trace context from AWSTraceHeader when no tracecontext found from payload", () => { + const tracerWrapper = new TracerWrapper(); + const payload: SNSEvent = { + Records: [ + { + EventSource: "aws:sns", + EventVersion: "1.0", + EventSubscriptionArn: + "arn:aws:sns:us-west-2:425362996713:DdTraceXLambda-snsjavachecksNestedStacksnsjavachecksNestedStackResource569D7DFD-VGBRIQ3RKOFR-snsJavaProducer2snsjavaproducerforPythontopic2EDA1CA9-A4bHj7LHgBQh:f3864261-c160-443d-b21e-8effb1899456", + Sns: { + Type: "Notification", + MessageId: "f0cae95c-225d-5718-896c-b7ef1e3a7108", + TopicArn: + "arn:aws:sns:us-west-2:425362996713:DdTraceXLambda-snsjavachecksNestedStacksnsjavachecksNestedStackResource569D7DFD-VGBRIQ3RKOFR-snsJavaProducer2snsjavaproducerforPythontopic2EDA1CA9-A4bHj7LHgBQh", + Subject: "", + Message: "hello from DdTraceXLambda-snsjavache-snsjavaproducerforPython-f1dKTIyh8pN6", + Timestamp: "2024-05-06T20:14:31.519Z", + SignatureVersion: "1", + Signature: + "gkBg+e7h+DSrgXzehe+rWCJmH0UWXwsaAz/dnngWivJmYKbdH0pNHNpM0ttFko2OItYzDEE+DdEm9dfSF7MYzl36UaFpijr1VI6YIwRcYAy6wpwwV23uX1XD7PrIoE/4/bonL6WJQLhdKiO3SD4eKd1RoIk4e1SjDzfMUtz6/4UaXC56wwQ/mj9QwglZ867xM0icLJjyT2LWIrIkHjyXFNFrQUn973sRcgfZDs92kOfacF0/Zqm6I/pHzoDeoQU4MapuX8jWSA45bPjxtWXglpez0z7Yne4eLujPSSEwzJRyPJeRaxdyL1NFF43X0SAZ5C7gWgeExN7GJKs1zW2sFA==", + SigningCertUrl: + "https://sns.us-west-2.amazonaws.com/SimpleNotificationService-************************33ab7e69", + UnsubscribeUrl: + "https://sns.us-west-2.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:us-west-2:425362996713:DdTraceXLambda-snsjavachecksNestedStacksnsjavachecksNestedStackResource569D7DFD-VGBRIQ3RKOFR-snsJavaProducer2snsjavaproducerforPythontopic2EDA1CA9-A4bHj7LHgBQh:f3864261-c160-443d-b21e-8effb1899456", + MessageAttributes: { + test: { + Type: "String", + Value: "test", + }, + }, + }, + }, + ], + }; + + const extractor = new SNSEventTraceExtractor(tracerWrapper); + process.env["_X_AMZN_TRACE_ID"] = + "Root=1-66393a26-0000000017acacbad335fb99;Parent=12f5e70cc905dfb7;Sampled=1;Lineage=48e79d5f:0"; + const traceContext = extractor.extract(payload); + process.env["_X_AMZN_TRACE_ID"] = ""; + expect(traceContext).not.toBeNull(); + + // Should not use ddtracer extractor. Because 1. it's an unnecessary extra step and + // 2. More importantly, DD_TRACE_PROPAGATION_STYLE could cause extraction fail + expect(spyTracerWrapper).not.toHaveBeenCalled(); + + expect(traceContext?.toTraceId()).toBe("1705928277274000281"); + expect(traceContext?.toSpanId()).toBe("1366252104075042743"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); }); }); diff --git a/src/trace/context/extractors/sns.ts b/src/trace/context/extractors/sns.ts index 6f881d54..33a009db 100644 --- a/src/trace/context/extractors/sns.ts +++ b/src/trace/context/extractors/sns.ts @@ -3,29 +3,44 @@ import { TracerWrapper } from "../../tracer-wrapper"; import { logDebug } from "../../../utils"; import { EventTraceExtractor } from "../extractor"; import { SpanContextWrapper } from "../../span-context-wrapper"; +import { XrayService, AMZN_TRACE_ID_ENV_VAR } from "../../xray-service"; export class SNSEventTraceExtractor implements EventTraceExtractor { constructor(private tracerWrapper: TracerWrapper) {} extract(event: SNSEvent): SpanContextWrapper | null { - const messageAttribute = event?.Records?.[0]?.Sns?.MessageAttributes?._datadog; - if (messageAttribute?.Value === undefined) return null; - try { - let headers; - if (messageAttribute.Type === "String") { - headers = JSON.parse(messageAttribute.Value); - } else { - // Try decoding base64 values - const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii"); - headers = JSON.parse(decodedValue); - } - - const traceContext = this.tracerWrapper.extract(headers); - if (traceContext === null) return null; + // First try to extract trace context from message attributes + const messageAttribute = event?.Records?.[0]?.Sns?.MessageAttributes?._datadog; + if (messageAttribute?.Value) { + let headers; + if (messageAttribute.Type === "String") { + headers = JSON.parse(messageAttribute.Value); + } else { + // Try decoding base64 values + const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii"); + headers = JSON.parse(decodedValue); + } - logDebug(`Extracted trace context from SNS event`, { traceContext, event }); - return traceContext; + const traceContext = this.tracerWrapper.extract(headers); + if (traceContext) { + logDebug("Extracted trace context from SNS event"); + return traceContext; + } else { + logDebug("Failed to extract trace context from SNS event"); + } + } + // Then try to extract trace context from _X_AMZN_TRACE_ID header (Upstream Java apps can + // pass down Datadog trace id (parent id wrong) in the env in SNS case) + if (process.env[AMZN_TRACE_ID_ENV_VAR]) { + const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(process.env[AMZN_TRACE_ID_ENV_VAR]); + if (traceContext) { + logDebug("Extracted Datadog trace context from _X_AMZN_TRACE_ID"); + return traceContext; + } else { + logDebug("No Datadog trace context found from _X_AMZN_TRACE_ID"); + } + } } catch (error) { if (error instanceof Error) { logDebug("Unable to extract trace context from SNS event", error); diff --git a/src/trace/context/extractors/sqs.spec.ts b/src/trace/context/extractors/sqs.spec.ts index ce74cdfc..2da249e0 100644 --- a/src/trace/context/extractors/sqs.spec.ts +++ b/src/trace/context/extractors/sqs.spec.ts @@ -21,6 +21,7 @@ describe("SQSEventTraceExtractor", () => { describe("extract", () => { beforeEach(() => { mockSpanContext = null; + spyTracerWrapper.mockClear(); }); afterEach(() => { @@ -128,13 +129,6 @@ describe("SQSEventTraceExtractor", () => { }); 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: [ @@ -163,11 +157,9 @@ describe("SQSEventTraceExtractor", () => { 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", - }); + // Should not use ddtracer extractor. Because 1. it's an unnecessary extra step and + // 2. More importantly, DD_TRACE_PROPAGATION_STYLE could cause extraction fail + expect(spyTracerWrapper).not.toHaveBeenCalled(); expect(traceContext?.toTraceId()).toBe("625397077193750208"); expect(traceContext?.toSpanId()).toBe("6538302989251745223"); diff --git a/src/trace/context/extractors/sqs.ts b/src/trace/context/extractors/sqs.ts index bf7d5a3b..b6485f69 100644 --- a/src/trace/context/extractors/sqs.ts +++ b/src/trace/context/extractors/sqs.ts @@ -10,23 +10,28 @@ export class SQSEventTraceExtractor implements EventTraceExtractor { extract(event: SQSEvent): SpanContextWrapper | null { try { - let parsedHeaders; + // First try to extract trace context from message attributes 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); + const traceContext = this.tracerWrapper.extract(JSON.parse(headers)); + if (traceContext) { + logDebug("Extracted trace context from SQS event messageAttributes"); + return traceContext; + } else { + logDebug("Failed to extract trace context from messageAttributes"); + } } - 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; + // Then try to extract trace context from attributes.AWSTraceHeader. (Upstream Java apps can + // pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case) + if (event?.Records?.[0]?.attributes?.AWSTraceHeader !== undefined) { + const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader); + if (traceContext) { + logDebug("Extracted trace context from SQS event attributes AWSTraceHeader"); + return traceContext; + } else { + logDebug("No Datadog trace context found from SQS event attributes AWSTraceHeader"); + } } - - logDebug("Extracted trace context from SQS event", { traceContext, event }); - return traceContext; } catch (error) { if (error instanceof Error) { logDebug("Unable to extract trace context from SQS event", error); diff --git a/src/trace/xray-service.spec.ts b/src/trace/xray-service.spec.ts index 3d7428df..e746d1e0 100644 --- a/src/trace/xray-service.spec.ts +++ b/src/trace/xray-service.spec.ts @@ -369,11 +369,9 @@ describe("XrayService", () => { 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", - }); + expect(ddTraceContext?.toTraceId()).toEqual("625397077193750208"); + expect(ddTraceContext?.toSpanId()).toEqual("6538302989251745223"); + expect(ddTraceContext?.sampleMode()).toEqual("1"); }); it("returns null when AWS trace header is NOT injected by dd-trace", () => { diff --git a/src/trace/xray-service.ts b/src/trace/xray-service.ts index f0308487..a0158475 100644 --- a/src/trace/xray-service.ts +++ b/src/trace/xray-service.ts @@ -4,14 +4,8 @@ import { SampleMode, TraceContext, TraceSource } from "./trace-context-service"; 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"; + +export 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 { @@ -206,7 +200,7 @@ export class XrayService { } } - public static extraceDDContextFromAWSTraceHeader(amznTraceId: string): DatadogTraceHeaders | null { + public static extraceDDContextFromAWSTraceHeader(amznTraceId: string): SpanContextWrapper | null { const awsContext = XrayService.parseAWSTraceHeader(amznTraceId); if (awsContext === undefined) { return null; @@ -214,11 +208,12 @@ export class XrayService { 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 SpanContextWrapper.fromTraceContext({ + traceId: hexStrToDecimalStr(traceIdParts[2].substring(8)), + parentId: hexStrToDecimalStr(awsContext.parentId), + sampleMode: parseInt(awsContext.sampled, 10), + source: TraceSource.Event, // This is still an event-source trace although we are getting it from xray trace header + }); } return null; }