feat: eventbridge to sqs support (#412)

* feat: eventbridge to sqs support

* feat: Dry up SNS span creation

* fix: Remap nonconforming variables

* feat: Specs for eb -> sqs events

Author: astuyve

event_samples/eventbridge-sqs.json

@@ -0,0 +1,21 @@
+{
+  "Records": [
+    {
+      "messageId": "e995e54f-1724-41fa-82c0-8b81821f854e",
+      "receiptHandle": "AQEB4mIfRcyqtzn1X5Ss+ConhTejVGc+qnAcmu3/Z9ZvbNkaPcpuDLX/bzvPD/ZkAXJUXZcemGSJmd7L3snZHKMP2Ck8runZiyl4mubiLb444pZvdiNPuGRJ6a3FvgS/GQPzho/9nNMyOi66m8Viwh70v4EUCPGO4JmD3TTDAUrrcAnqU4WSObjfC/NAp9bI6wH2CEyAYEfex6Nxplbl/jBf9ZUG0I3m3vQd0Q4l4gd4jIR4oxQUglU2Tldl4Kx5fMUAhTRLAENri6HsY81avBkKd9FAuxONlsITB5uj02kOkvLlRGEcalqsKyPJ7AFaDLrOLaL3U+yReroPEJ5R5nwhLOEbeN5HROlZRXeaAwZOIN8BjqdeooYTIOrtvMEVb7a6OPLMdH1XB+ddevtKAH8K9Tm2ZjpaA7dtBGh1zFVHzBk=",
+      "body": "{\"version\":\"0\",\"id\":\"af718b2a-b987-e8c0-7a2b-a188fad2661a\",\"detail-type\":\"my.Detail\",\"source\":\"my.Source\",\"account\":\"425362996713\",\"time\":\"2023-08-03T22:49:03Z\",\"region\":\"us-east-1\",\"resources\":[],\"detail\":{\"text\":\"Hello, world!\",\"_datadog\":{\"x-datadog-trace-id\":\"7379586022458917877\",\"x-datadog-parent-id\":\"2644033662113726488\",\"x-datadog-sampling-priority\":\"1\",\"x-datadog-tags\":\"_dd.p.dm=-0\",\"traceparent\":\"00-000000000000000066698e63821a03f5-24b17e9b6476c018-01\",\"tracestate\":\"dd=t.dm:-0;s:1\"}}}",
+      "attributes": {
+        "ApproximateReceiveCount": "1",
+        "AWSTraceHeader": "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0",
+        "SentTimestamp": "1691102943638",
+        "SenderId": "AIDAJXNJGGKNS7OSV23OI",
+        "ApproximateFirstReceiveTimestamp": "1691102943647"
+      },
+      "messageAttributes": {},
+      "md5OfBody": "93d9f0cd8886d1e000a1a0b7007bffc4",
+      "eventSource": "aws:sqs",
+      "eventSourceARN": "arn:aws:sqs:us-east-1:425362996713:lambda-eb-sqs-lambda-dev-demo-queue",
+      "awsRegion": "us-east-1"
+    }
+  ]
+}

src/trace/context.spec.ts

@@ -251,6 +251,38 @@ describe("readTraceFromEvent", () => {
     });
   });
 
+  it("can parse an EventBridge message in an SQS queue", () => {
+    const result = readTraceFromEvent({
+      Records: [
+        {
+          messageId: "e995e54f-1724-41fa-82c0-8b81821f854e",
+          receiptHandle:
+            "AQEB4mIfRcyqtzn1X5Ss+ConhTejVGc+qnAcmu3/Z9ZvbNkaPcpuDLX/bzvPD/ZkAXJUXZcemGSJmd7L3snZHKMP2Ck8runZiyl4mubiLb444pZvdiNPuGRJ6a3FvgS/GQPzho/9nNMyOi66m8Viwh70v4EUCPGO4JmD3TTDAUrrcAnqU4WSObjfC/NAp9bI6wH2CEyAYEfex6Nxplbl/jBf9ZUG0I3m3vQd0Q4l4gd4jIR4oxQUglU2Tldl4Kx5fMUAhTRLAENri6HsY81avBkKd9FAuxONlsITB5uj02kOkvLlRGEcalqsKyPJ7AFaDLrOLaL3U+yReroPEJ5R5nwhLOEbeN5HROlZRXeaAwZOIN8BjqdeooYTIOrtvMEVb7a6OPLMdH1XB+ddevtKAH8K9Tm2ZjpaA7dtBGh1zFVHzBk=",
+          body: '{"version":"0","id":"af718b2a-b987-e8c0-7a2b-a188fad2661a","detail-type":"my.Detail","source":"my.Source","account":"425362996713","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!","_datadog":{"x-datadog-trace-id":"7379586022458917877","x-datadog-parent-id":"2644033662113726488","x-datadog-sampling-priority":"1","x-datadog-tags":"_dd.p.dm=-0","traceparent":"00-000000000000000066698e63821a03f5-24b17e9b6476c018-01","tracestate":"dd=t.dm:-0;s:1"}}}',
+          attributes: {
+            ApproximateReceiveCount: "1",
+            AWSTraceHeader: "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0",
+            SentTimestamp: "1691102943638",
+            SenderId: "AIDAJXNJGGKNS7OSV23OI",
+            ApproximateFirstReceiveTimestamp: "1691102943647",
+          },
+          messageAttributes: {},
+          md5OfBody: "93d9f0cd8886d1e000a1a0b7007bffc4",
+          eventSource: "aws:sqs",
+          eventSourceARN: "arn:aws:sqs:us-east-1:425362996713:lambda-eb-sqs-lambda-dev-demo-queue",
+          awsRegion: "us-east-1",
+        },
+      ],
+    });
+
+    expect(result).toEqual({
+      parentID: "2644033662113726488",
+      sampleMode: 1,
+      source: "event",
+      traceID: "7379586022458917877",
+    });
+  });
+
   it("can parse an SNS message source", () => {
     const result = readTraceFromEvent({
       Records: [

src/trace/context.ts

@@ -10,6 +10,7 @@ import {
   isSNSEvent,
   isSNSSQSEvent,
   isSQSEvent,
+  isEBSQSEvent,
 } from "../utils/event-type-guards";
 import {
   authorizingRequestIdHeader,
@@ -309,6 +310,25 @@ export function readTraceFromSNSSQSEvent(event: SQSEvent): TraceContext | undefi
   }
 }
 
+export function readTraceFromEBSQSEvent(event: SQSEvent): TraceContext | undefined {
+  if (event?.Records?.[0]?.body) {
+    try {
+      const parsedBody = JSON.parse(event.Records[0].body) as EventBridgeEvent<any, any>;
+      if (parsedBody?.detail?._datadog) {
+        const trace = exportTraceData(parsedBody.detail._datadog);
+
+        logDebug(`extracted trace context from EventBridge SQS event`, { trace, event });
+        return trace;
+      }
+    } catch (err) {
+      if (err instanceof Error) {
+        logDebug("Error parsing EventBridge SQS message trace data", err as Error);
+      }
+      return;
+    }
+  }
+}
+
 export function readTraceFromKinesisEvent(event: KinesisStreamEvent): TraceContext | undefined {
   if (event?.Records?.[0]?.kinesis?.data) {
     try {
@@ -463,6 +483,10 @@ export function readTraceFromEvent(event: any, decodeAuthorizerContext: boolean
     return readTraceFromSNSSQSEvent(event);
   }
 
+  if (isEBSQSEvent(event)) {
+    return readTraceFromEBSQSEvent(event);
+  }
+
   if (isAppSyncResolverEvent(event)) {
     return readTraceFromAppSyncEvent(event);
   }

src/trace/span-inferrer.spec.ts

@@ -8,6 +8,7 @@ const sqsEvent = require("../../event_samples/sqs.json");
 const ddbEvent = require("../../event_samples/dynamodb.json");
 const kinesisEvent = require("../../event_samples/kinesis.json");
 const eventBridgeEvent = require("../../event_samples/eventbridge.json");
+const eventBridgeSQSEvent = require("../../event_samples/eventbridge-sqs.json");
 const webSocketEvent = require("../../event_samples/api-gateway-wss.json");
 const apiGatewayV1 = require("../../event_samples/api-gateway-v1.json");
 const apiGatewayV2 = require("../../event_samples/api-gateway-v2.json");
@@ -568,6 +569,55 @@ describe("SpanInferrer", () => {
     });
   });
 
+  it("creates an inferred span for eventbridge sqs events", () => {
+    const inferrer = new SpanInferrer(mockWrapper as unknown as TracerWrapper);
+    inferrer.createInferredSpan(eventBridgeSQSEvent, {} as any, {} as SpanContext);
+
+    expect(mockWrapper.startSpan.mock.calls).toEqual([
+      [
+        "aws.eventbridge",
+        {
+          childOf: {},
+          startTime: 1691102943000,
+          tags: {
+            _inferred_span: { synchronicity: "async", tag_source: "self" },
+            operation_name: "aws.eventbridge",
+            "peer.service": "mock-lambda-service",
+            request_id: undefined,
+            "resource.name": "my.Source",
+            resource_names: "my.Source",
+            service: "eventbridge",
+            "span.type": "web",
+          },
+        },
+      ],
+      [
+        "aws.sqs",
+        {
+          childOf: undefined,
+          startTime: 1691102943638,
+          tags: {
+            _inferred_span: { synchronicity: "async", tag_source: "self" },
+            event_source_arn: "arn:aws:sqs:us-east-1:425362996713:lambda-eb-sqs-lambda-dev-demo-queue",
+            operation_name: "aws.sqs",
+            "peer.service": "mock-lambda-service",
+            queuename: "lambda-eb-sqs-lambda-dev-demo-queue",
+            receipt_handle:
+              "AQEB4mIfRcyqtzn1X5Ss+ConhTejVGc+qnAcmu3/Z9ZvbNkaPcpuDLX/bzvPD/ZkAXJUXZcemGSJmd7L3snZHKMP2Ck8runZiyl4mubiLb444pZvdiNPuGRJ6a3FvgS/GQPzho/9nNMyOi66m8Viwh70v4EUCPGO4JmD3TTDAUrrcAnqU4WSObjfC/NAp9bI6wH2CEyAYEfex6Nxplbl/jBf9ZUG0I3m3vQd0Q4l4gd4jIR4oxQUglU2Tldl4Kx5fMUAhTRLAENri6HsY81avBkKd9FAuxONlsITB5uj02kOkvLlRGEcalqsKyPJ7AFaDLrOLaL3U+yReroPEJ5R5nwhLOEbeN5HROlZRXeaAwZOIN8BjqdeooYTIOrtvMEVb7a6OPLMdH1XB+ddevtKAH8K9Tm2ZjpaA7dtBGh1zFVHzBk=",
+            request_id: undefined,
+            "resource.name": "lambda-eb-sqs-lambda-dev-demo-queue",
+            resource_names: "lambda-eb-sqs-lambda-dev-demo-queue",
+            retry_count: 1,
+            sender_id: "AIDAJXNJGGKNS7OSV23OI",
+            service: "sqs",
+            "service.name": "lambda-eb-sqs-lambda-dev-demo-queue",
+            "span.type": "web",
+          },
+        },
+      ],
+    ]);
+  });
+
   it("creates an inferred span for websocket events", () => {
     const inferrer = new SpanInferrer(mockWrapper as unknown as TracerWrapper);
     inferrer.createInferredSpan(webSocketEvent, {} as any, {} as SpanContext);

src/trace/span-inferrer.ts

@@ -280,57 +280,26 @@ export class SpanInferrer {
     parentSpanContext: SpanContext | undefined,
   ): SpanWrapper {
     const options: SpanOptions = {};
-    const { Records } = event as SNSEvent;
-    const referenceRecord = Records[0];
-    const {
-      EventSubscriptionArn,
-      Sns: { TopicArn, Timestamp, Type, Subject, MessageId },
-    } = referenceRecord;
-    const topicName = TopicArn?.split(":").pop() || "";
-    const resourceName = topicName;
-    const serviceName = SpanInferrer.determineServiceName(topicName, "lambda_sns", "sns");
-    options.tags = {
-      operation_name: "aws.sns",
-      resource_names: resourceName,
-      request_id: context?.awsRequestId,
-      "span.type": "sns",
-      "resource.name": resourceName,
-      "peer.service": this.service,
-      service: serviceName,
-      _inferred_span: {
-        tag_source: "self",
-        synchronicity: "async",
-      },
-      type: Type,
-      subject: Subject,
-      message_id: MessageId,
-      topicname: topicName,
-      topic_arn: TopicArn,
-      event_subscription_arn: EventSubscriptionArn,
-    };
-    if (parentSpanContext) {
-      options.childOf = parentSpanContext;
+
+    let referenceRecord: SNSMessage;
+    let eventSubscriptionArn = "";
+    if (event.Records) {
+      // Full SNS Event into Lambda
+      const { Records } = event as SNSEvent;
+      ({ Sns: referenceRecord, EventSubscriptionArn: eventSubscriptionArn } = Records[0]);
+    } else {
+      // SNS message wrapping an SQS message
+      referenceRecord = event;
     }
-    options.startTime = Date.parse(Timestamp);
-    const spanWrapperOptions = {
-      isAsync: true,
-    };
-    return new SpanWrapper(this.traceWrapper.startSpan("aws.sns", options), spanWrapperOptions);
-  }
+    const { TopicArn, Timestamp, Type, Subject, MessageId } = referenceRecord;
 
-  createInferredSpanForSqsSns(
-    event: SNSMessage,
-    context: Context | undefined,
-    parentSpanContext: SpanContext | undefined,
-  ): SpanWrapper {
-    const options: SpanOptions = {};
-    const { TopicArn, Timestamp, Type, Subject, MessageId } = event;
     const topicName = TopicArn?.split(":").pop() || "";
     const resourceName = topicName;
     const serviceName = SpanInferrer.determineServiceName(topicName, "lambda_sns", "sns");
     options.tags = {
       operation_name: "aws.sns",
       resource_names: resourceName,
+      request_id: context?.awsRequestId,
       "span.type": "sns",
       "resource.name": resourceName,
       "peer.service": this.service,
@@ -345,6 +314,11 @@ export class SpanInferrer {
       topicname: topicName,
       topic_arn: TopicArn,
     };
+
+    // EventSubscriptionARN not available for direct integrations to SQS from SNS.
+    if (eventSubscriptionArn !== "") {
+      options.tags.event_subscription_arn = eventSubscriptionArn;
+    }
     if (parentSpanContext) {
       options.childOf = parentSpanContext;
     }
@@ -396,18 +370,21 @@ export class SpanInferrer {
     // Check if sqs message was from sns
     // If so, unpack and look at timestamp
     // create further upstream sns span and finish/attach it here
-    let upstreamSnsSpan: SpanWrapper | null = null;
+    let upstreamSpan: SpanWrapper | null = null;
     try {
-      let upstreamSnsMessage: SNSMessage;
-      upstreamSnsMessage = JSON.parse(body);
-      if (upstreamSnsMessage && upstreamSnsMessage.TopicArn && upstreamSnsMessage.Timestamp) {
-        upstreamSnsSpan = this.createInferredSpanForSqsSns(upstreamSnsMessage, context, parentSpanContext);
-        upstreamSnsSpan.finish(Number(SentTimestamp));
+      let upstreamMessage: any;
+      upstreamMessage = JSON.parse(body);
+      if (upstreamMessage && upstreamMessage.TopicArn && upstreamMessage.Timestamp) {
+        upstreamSpan = this.createInferredSpanForSns(upstreamMessage, context, parentSpanContext);
+        upstreamSpan.finish(Number(SentTimestamp));
+      } else if (upstreamMessage?.detail?._datadog) {
+        upstreamSpan = this.createInferredSpanForEventBridge(upstreamMessage, context, parentSpanContext);
+        upstreamSpan.finish(Number(SentTimestamp));
       }
     } catch (e) {
       // Pass, it's a raw SQS message
     }
-    options.childOf = upstreamSnsSpan ? upstreamSnsSpan.span : parentSpanContext;
+    options.childOf = upstreamSpan ? upstreamSpan.span : parentSpanContext;
 
     options.startTime = Number(SentTimestamp);
 
@@ -512,7 +489,7 @@ export class SpanInferrer {
   }
 
   createInferredSpanForEventBridge(
-    event: any,
+    event: EventBridgeEvent<any, any>,
     context: Context | undefined,
     parentSpanContext: SpanContext | undefined,
   ): SpanWrapper {

src/utils/event-type-guards.ts

@@ -83,6 +83,18 @@ export function isSNSSQSEvent(event: any): event is SQSEvent {
   return false;
 }
 
+export function isEBSQSEvent(event: any): event is SQSEvent {
+  if (Array.isArray(event.Records) && event.Records.length > 0 && event.Records[0].eventSource === "aws:sqs") {
+    try {
+      const body = JSON.parse(event.Records[0].body) as EventBridgeEvent<any, any>;
+      return body["detail-type"] !== undefined;
+    } catch (e) {
+      return false;
+    }
+  }
+  return false;
+}
+
 export function isAppSyncResolverEvent(event: any): event is AppSyncResolverEvent<any> {
   return event.info !== undefined && event.info.selectionSetGraphQL !== undefined;
 }