Propagate pubsub context when lambda event is an sqs event

Author: wpessers

packages/instrumentation-aws-lambda/src/instrumentation.ts

@@ -39,11 +39,20 @@ import {
   ROOT_CONTEXT,
   Attributes,
 } from '@opentelemetry/api';
+import { pubsubPropagation } from '@opentelemetry/propagation-utils';
 import {
   ATTR_URL_FULL,
   SEMATTRS_FAAS_EXECUTION,
   SEMRESATTRS_CLOUD_ACCOUNT_ID,
   SEMRESATTRS_FAAS_ID,
+  MESSAGINGDESTINATIONKINDVALUES_QUEUE,
+  MESSAGINGOPERATIONVALUES_PROCESS,
+  SEMATTRS_MESSAGING_DESTINATION,
+  SEMATTRS_MESSAGING_DESTINATION_KIND,
+  SEMATTRS_MESSAGING_MESSAGE_ID,
+  SEMATTRS_MESSAGING_OPERATION,
+  SEMATTRS_MESSAGING_SYSTEM,
+  SEMATTRS_MESSAGING_URL,
 } from '@opentelemetry/semantic-conventions';
 import { ATTR_FAAS_COLDSTART } from './semconv';
 
@@ -68,6 +77,15 @@ const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
   },
 };
 
+const sqsContextGetter = {
+  keys(carrier: any): string[] {
+    return Object.keys(carrier || {});
+  },
+  get(carrier: any, key: string) {
+    return carrier?.[key]?.stringValue;
+  },
+};
+
 export const lambdaMaxInitInMilliseconds = 10_000;
 
 export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstrumentationConfig> {
@@ -265,6 +283,42 @@ export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstr
       }
 
       return otelContext.with(trace.setSpan(parent, span), () => {
+        if (event.Records) {
+          const messages = event.Records;
+          const queueArn = messages[0]?.eventSourceARN;
+          const queueName = queueArn?.split(':').pop() ?? 'unknown';
+
+          pubsubPropagation.patchMessagesArrayToStartProcessSpans({
+            messages,
+            parentContext: trace.setSpan(otelContext.active(), span),
+            tracer: plugin.tracer,
+            messageToSpanDetails: (message: any) => ({
+              name: queueName,
+              parentContext: propagation.extract(
+                ROOT_CONTEXT,
+                message.messageAttributes || {},
+                sqsContextGetter
+              ),
+              attributes: {
+                [SEMATTRS_MESSAGING_SYSTEM]: 'aws.sqs',
+                [SEMATTRS_MESSAGING_DESTINATION]: queueName,
+                [SEMATTRS_MESSAGING_DESTINATION_KIND]:
+                  MESSAGINGDESTINATIONKINDVALUES_QUEUE,
+                [SEMATTRS_MESSAGING_MESSAGE_ID]: message.messageId,
+                [SEMATTRS_MESSAGING_URL]: queueArn,
+                [SEMATTRS_MESSAGING_OPERATION]:
+                  MESSAGINGOPERATIONVALUES_PROCESS,
+              },
+            }),
+          });
+
+          pubsubPropagation.patchArrayForProcessSpans(
+            messages,
+            plugin.tracer,
+            otelContext.active()
+          );
+        }
+
         // Lambda seems to pass a callback even if handler is of Promise form, so we wrap all the time before calling
         // the handler and see if the result is a Promise or not. In such a case, the callback is usually ignored. If
         // the handler happened to both call the callback and complete a returned Promise, whichever happens first will

packages/instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts

@@ -30,7 +30,7 @@ import {
   ReadableSpan,
 } from '@opentelemetry/sdk-trace-base';
 import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
-import { Context } from 'aws-lambda';
+import { APIGatewayProxyEvent, Context, SQSEvent } from 'aws-lambda';
 import * as assert from 'assert';
 import {
   ATTR_URL_FULL,
@@ -827,4 +827,44 @@ describe('lambda handler', () => {
       );
     });
   });
+
+  describe('sqs test', () => {
+    it('creates process span for sqs record, with lambda invocation span as parent and span link to the producer traceId and spanId', async () => {
+      initializeHandler('lambda-test/sync.sqshandler');
+      const producerTraceId = '1df415edd0ad7f83e573f6504381dcec';
+      const producerSpanId = '83b7424a259945cb';
+      const event = {
+        Records: [
+          {
+            messageAttributes: {
+              traceparent: {
+                stringValue: `00-${producerTraceId}-${producerSpanId}-01`,
+                dataType: 'String',
+              },
+            },
+            eventSource: 'aws:sqs',
+            eventSourceARN:
+              'arn:aws:sqs:eu-central-1:783764587482:launch-queue',
+          },
+        ],
+      };
+
+      await lambdaRequire('lambda-test/sync').sqshandler(event, ctx, () => {});
+      const spans = memoryExporter.getFinishedSpans();
+
+      assert.strictEqual(spans.length, 2);
+
+      assert.equal(
+        spans[0].parentSpanContext?.traceId,
+        spans[1].spanContext().traceId
+      );
+      assert.equal(
+        spans[0].parentSpanContext?.spanId,
+        spans[1].spanContext().spanId
+      );
+
+      assert.equal(spans[0].links[0]?.context.traceId, producerTraceId);
+      assert.equal(spans[0].links[0].context.spanId, producerSpanId);
+    });
+  });
 });

packages/instrumentation-aws-lambda/test/lambda-test/sync.js

@@ -19,21 +19,27 @@ exports.handler = function (event, context, callback) {
   callback(null, 'ok');
 };
 
+exports.sqshandler = function (event, context, callback) {
+  // Dummy forEach loop, to trigger sqs instrumentation
+  event.Records.forEach(_r => {});
+  callback(null, 'ok');
+};
+
 exports.error = function (event, context, callback) {
   throw new Error('handler error');
-}
+};
 
 exports.callbackerror = function (event, context, callback) {
   callback(new Error('handler error'));
-}
+};
 
 exports.stringerror = function (event, context, callback) {
   throw 'handler error';
-}
+};
 
 exports.callbackstringerror = function (event, context, callback) {
   callback('handler error');
-}
+};
 
 exports.context = function (event, context, callback) {
   callback(null, api.trace.getSpan(api.context.active()).spanContext().traceId);