Merge pull request #17 from aspecto-io/sqs-context-propagation

feat: SQS context propagation

Author: Amir Blum

packages/plugin-aws-sdk/docs/sqs.md

@@ -5,17 +5,17 @@ SQS is amazon's managed message queue. Thus, it should follow the [Open Telemetr
 The following methods are automatically enhanced:
 
 ### sendMessage / sendMessageBatch
-- [Message Attributes](https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/trace/semantic_conventions/messaging.md#messaging-attributes) are added by this plugin according to the spec.
-- TODO: Inject trace context as SQS MessageAttributes, so the service receiving the message can link cascading spans to the trace which created the message. 
+- [Messaging Attributes](https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/trace/semantic_conventions/messaging.md#messaging-attributes) are added by this plugin according to the spec.
+- Open Telemetry trace context is injected as SQS MessageAttributes, so the service receiving the message can link cascading spans to the trace which created the message. 
 
 ### receiveMessage
-- [Message Attributes](https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/trace/semantic_conventions/messaging.md#messaging-attributes) are added by this plugin according to the spec.
+- [Messaging Attributes](https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/trace/semantic_conventions/messaging.md#messaging-attributes) are added by this plugin according to the spec.
 - Additional "processing spans" are created for each message received by the application.   
 If an application invoked `receiveMessage`, and received a 10 messages batch, a single `messaging.operation` = `receive` span will be created for the `receiveMessage` operation, and 10 `messaging.operation` = `process` spans will be created, one for each message.  
 Those processing spans are created by the library. This behavior is partially implemented, [See discussion below](#processing-spans).
 - Sets the inter process context correctly, so that additional spans created through the process will be linked to parent spans correctly.  
 This behavior is partially implemented, [See discussion below](#processing-spans).
-- TODO: Extract trace context from SQS MessageAttributes, and set span's `parent` and `links` correctly according to the spec.
+- Extract trace context from SQS MessageAttributes, and set span's `parent` and `links` correctly according to the spec.
 
 #### Processing Spans
 According to open telemetry specification (and to reasonable expectation for trace structure), user of this library would expect to see one span for the operation of receiving messages batch from SQS, and then, for each message, a span with it's own sub-tree for the processing of this specific message. 

packages/plugin-aws-sdk/package.json

@@ -35,8 +35,8 @@
     "url": "https://github.com/aspecto-io/opentelemetry-ext-js/issues"
   },
   "devDependencies": {
-    "@opentelemetry/node": "^0.8.3",
-    "@opentelemetry/tracing": "^0.8.3",
+    "@opentelemetry/node": "^0.9.0",
+    "@opentelemetry/tracing": "^0.9.0",
     "@types/jest": "^25.2.3",
     "@types/shimmer": "^1.0.1",
     "aws-sdk": "^2.694.0",
@@ -48,7 +48,7 @@
     "typescript": "^3.9.5"
   },
   "dependencies": {
-    "@opentelemetry/core": "^0.8.3",
+    "@opentelemetry/core": "^0.9.0",
     "shimmer": "^1.2.1"
   },
   "jest": {

packages/plugin-aws-sdk/src/aws-sdk.ts

@@ -29,7 +29,10 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
   }
 
   protected patch() {
-    this.servicesExtensions = new ServicesExtensions(this._tracer);
+    this.servicesExtensions = new ServicesExtensions(
+      this._tracer,
+      this._logger
+    );
 
     this._logger.debug(
       "applying patch to %s@%s",

packages/plugin-aws-sdk/src/services/ServicesExtensions.ts

@@ -1,13 +1,13 @@
-import { Tracer, Span } from "@opentelemetry/api";
+import { Tracer, Span, Logger } from "@opentelemetry/api";
 import { ServiceExtension, RequestMetadata } from "./ServiceExtension";
 import { SqsServiceExtension } from "./sqs";
 import * as AWS from "aws-sdk";
 
 export class ServicesExtensions implements ServiceExtension {
   services: Map<string, ServiceExtension> = new Map();
 
-  constructor(tracer: Tracer) {
-    this.services.set("sqs", new SqsServiceExtension(tracer));
+  constructor(tracer: Tracer, logger: Logger) {
+    this.services.set("sqs", new SqsServiceExtension(tracer, logger));
   }
 
   requestHook(request: AWS.Request<any, any>): RequestMetadata {

packages/plugin-aws-sdk/src/services/sqs.ts

@@ -1,6 +1,24 @@
-import { Tracer, SpanKind, Span } from "@opentelemetry/api";
+import {
+  Tracer,
+  SpanKind,
+  Span,
+  propagation,
+  Context,
+  Link,
+  Logger,
+} from "@opentelemetry/api";
 import { RequestMetadata, ServiceExtension } from "./ServiceExtension";
 import * as AWS from "aws-sdk";
+import {
+  getExtractedSpanContext,
+  TRACE_PARENT_HEADER,
+} from "@opentelemetry/core";
+import {
+  MessageBodyAttributeMap,
+  SendMessageRequest,
+  SendMessageBatchRequest,
+  SendMessageBatchRequestEntry,
+} from "aws-sdk/clients/sqs";
 
 export enum SqsAttributeNames {
   // https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/trace/semantic_conventions/messaging.md
@@ -20,11 +38,34 @@ export const END_SPAN_FUNCTION = Symbol(
   "opentelemetry.plugin.aws-sdk.sqs.end_span"
 );
 
+// https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-quotas.html
+const SQS_MAX_MESSAGE_ATTRIBUTES = 10;
+
+const contextSetterFunc = (
+  messageAttributes: AWS.SQS.MessageBodyAttributeMap,
+  key: string,
+  value: unknown
+) => {
+  messageAttributes[key] = {
+    DataType: "String",
+    StringValue: value as string,
+  };
+};
+
+const contextGetterFunc = (
+  messageAttributes: AWS.SQS.MessageBodyAttributeMap,
+  key: string
+) => {
+  return messageAttributes?.[key]?.StringValue;
+};
+
 export class SqsServiceExtension implements ServiceExtension {
   tracer: Tracer;
+  logger: Logger;
 
-  constructor(tracer: Tracer) {
+  constructor(tracer: Tracer, logger: Logger) {
     this.tracer = tracer;
+    this.logger = logger;
   }
 
   requestHook(request: AWS.Request<any, any>): RequestMetadata {
@@ -45,16 +86,45 @@ export class SqsServiceExtension implements ServiceExtension {
     const operation = (request as any)?.operation;
     switch (operation) {
       case "receiveMessage":
-        isIncoming = true;
-        spanKind = SpanKind.CONSUMER;
-        spanName = queueName;
-        spanAttributes[SqsAttributeNames.MESSAGING_OPERATION] = "receive";
+        {
+          isIncoming = true;
+          spanKind = SpanKind.CONSUMER;
+          spanName = queueName;
+          spanAttributes[SqsAttributeNames.MESSAGING_OPERATION] = "receive";
+
+          const params: Record<string, any> = (request as any).params;
+          const attributesNames = params.MessageAttributeNames || [];
+          attributesNames.push(TRACE_PARENT_HEADER);
+          params.MessageAttributeNames = attributesNames;
+        }
         break;
 
       case "sendMessage":
+        {
+          spanKind = SpanKind.PRODUCER;
+          spanName = queueName;
+
+          const params: SendMessageRequest = (request as any).params;
+          params.MessageAttributes = this.InjectPropagationContext(
+            params.MessageAttributes
+          );
+        }
+        break;
+
       case "sendMessageBatch":
-        spanKind = SpanKind.PRODUCER;
-        spanName = queueName;
+        {
+          spanKind = SpanKind.PRODUCER;
+          spanName = queueName;
+
+          const params: SendMessageBatchRequest = (request as any).params;
+          params.Entries.forEach(
+            (messageParams: SendMessageBatchRequestEntry) => {
+              messageParams.MessageAttributes = this.InjectPropagationContext(
+                messageParams.MessageAttributes
+              );
+            }
+          );
+        }
         break;
     }
 
@@ -73,9 +143,18 @@ export class SqsServiceExtension implements ServiceExtension {
       const queueName = this.extractQueueNameFromUrl(queueUrl);
 
       messages.forEach((message: AWS.SQS.Message) => {
+        const parentContext: Context = propagation.extract(
+          message.MessageAttributes,
+          contextGetterFunc
+        );
         message[START_SPAN_FUNCTION] = () => {
           return this.tracer.withSpan(span, () =>
-            this.startSingleMessageSpan(queueUrl, queueName, message)
+            this.startSingleMessageSpan(
+              queueUrl,
+              queueName,
+              message,
+              parentContext
+            )
           );
         };
         message[END_SPAN_FUNCTION] = () =>
@@ -105,8 +184,16 @@ export class SqsServiceExtension implements ServiceExtension {
   startSingleMessageSpan(
     queueUrl: string,
     queueName: string,
-    message: AWS.SQS.Message
+    message: AWS.SQS.Message,
+    propagtedContext: Context
   ): Span {
+    const links: Link[] = [];
+    if (propagtedContext) {
+      links.push({
+        context: getExtractedSpanContext(propagtedContext),
+      } as Link);
+    }
+
     const messageSpan = this.tracer.startSpan(queueName, {
       kind: SpanKind.CONSUMER,
       attributes: {
@@ -117,6 +204,7 @@ export class SqsServiceExtension implements ServiceExtension {
         [SqsAttributeNames.MESSAGING_URL]: queueUrl,
         [SqsAttributeNames.MESSAGING_OPERATION]: "process",
       },
+      links,
     });
 
     message[START_SPAN_FUNCTION] = () =>
@@ -147,4 +235,18 @@ export class SqsServiceExtension implements ServiceExtension {
       });
     };
   }
+
+  InjectPropagationContext(
+    attributesMap?: MessageBodyAttributeMap
+  ): MessageBodyAttributeMap {
+    const attributes = attributesMap ?? {};
+    if (Object.keys(attributes).length < SQS_MAX_MESSAGE_ATTRIBUTES) {
+      propagation.inject(attributes, contextSetterFunc);
+    } else {
+      this.logger.warn(
+        "OpenTelemetry aws-sdk plugin cannot set context propagation on SQS message due to maximum amount of MessageAttributes"
+      );
+    }
+    return attributes;
+  }
 }