fix(plugin-aws-sdk): set span hierarchy for sqs processing

Author: Amir Blum

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

@@ -69,10 +69,7 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
     return target;
   }
 
-  private _startAwsSpan(
-    request: AWS.Request<any, any>,
-    additionalAttributes?: Attributes
-  ): Span {
+  private _startAwsSpan(request: AWS.Request<any, any>): Span {
     const operation = (request as any).operation;
     const service = (request as any).service;
 
@@ -86,7 +83,6 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         [AttributeNames.AWS_SERVICE_API]: service?.api?.className,
         [AttributeNames.AWS_SERVICE_IDENTIFIER]: service?.serviceIdentifier,
         [AttributeNames.AWS_SERVICE_NAME]: service?.api?.abbreviation,
-        ...additionalAttributes,
       },
     });
 
@@ -118,11 +114,12 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
 
       const responseAttributes = getResponseServiceAttributes(
         response,
+        span,
         thisPlugin._tracer
       );
+
       span.setAttributes({
         [AttributeNames.AWS_REQUEST_ID]: response.requestId,
-        ...responseAttributes,
       });
       span.end();
     });
@@ -145,11 +142,8 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         return original.apply(this, arguments);
       }
 
-      const requestMetadata = getRequestServiceAttributes(awsRequest);
-      const span = thisPlugin._startAwsSpan(
-        awsRequest,
-        requestMetadata?.attributes
-      );
+      const span = thisPlugin._startAwsSpan(awsRequest);
+      const requestMetadata = getRequestServiceAttributes(awsRequest, span);
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 
@@ -175,11 +169,8 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         return original.apply(this, arguments);
       }
 
-      const requestMetadata = getRequestServiceAttributes(awsRequest);
-      const span = thisPlugin._startAwsSpan(
-        awsRequest,
-        requestMetadata?.attributes
-      );
+      const span = thisPlugin._startAwsSpan(awsRequest);
+      const requestMetadata = getRequestServiceAttributes(awsRequest, span);
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 

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

@@ -1,8 +1,9 @@
-import { Attributes } from "@opentelemetry/api";
+import { Attributes, Span } from "@opentelemetry/api";
 import { RequestMetadata } from "./service-attributes";
 
 export function getS3RequestSpanAttributes(
-  request: AWS.Request<any, any>
+  request: AWS.Request<any, any>,
+  span: Span
 ): RequestMetadata {
   return {
     attributes: {},
@@ -11,7 +12,8 @@ export function getS3RequestSpanAttributes(
 }
 
 export function getS3ResponseSpanAttributes(
-  response: AWS.Response<any, any>
+  response: AWS.Response<any, any>,
+  span: Span
 ): Attributes {
   return {};
 }

packages/plugin-aws-sdk/src/services/service-attributes.ts

@@ -1,24 +1,26 @@
-import { Attributes, Tracer } from "@opentelemetry/api";
+import { Attributes, Tracer, Span } from "@opentelemetry/api";
 import { getS3RequestSpanAttributes, getS3ResponseSpanAttributes } from "./s3";
 import {
   getSqsRequestSpanAttributes,
   getSqsResponseSpanAttributes,
 } from "./sqs";
 
 /*
-attributes are additional span attributes that should be added to the span.
 isIncoming - if true, then the operation callback / promise should be bind with the operation's span 
 */
 export interface RequestMetadata {
-  attributes: Attributes;
   isIncoming: boolean;
 }
 
-type RequestAttrProcessor = (request: AWS.Request<any, any>) => RequestMetadata;
+type RequestAttrProcessor = (
+  request: AWS.Request<any, any>,
+  span: Span
+) => RequestMetadata;
 type ResponseAttrProcessor = (
   response: AWS.Response<any, any>,
+  span: Span,
   tracer: Tracer
-) => Attributes;
+) => void;
 
 class ServiceAttributes {
   public attributeProcessors: {
@@ -41,23 +43,29 @@ class ServiceAttributes {
 }
 
 export function getRequestServiceAttributes(
-  request: AWS.Request<any, any>
+  request: AWS.Request<any, any>,
+  span: Span
 ): RequestMetadata {
   const serviceId = (request as any)?.service?.serviceIdentifier;
   if (serviceId) {
-    return serviceAttributes.attributeProcessors[serviceId]?.request(request);
+    return serviceAttributes.attributeProcessors[serviceId]?.request(
+      request,
+      span
+    );
   }
 }
 
 export function getResponseServiceAttributes(
   response: AWS.Response<any, any>,
+  span: Span,
   tracer: Tracer
-): Attributes {
+): void {
   const serviceId = (response as any)?.request?.service?.serviceIdentifier;
   if (!serviceId) return;
 
   return serviceAttributes.attributeProcessors[serviceId]?.response(
     response,
+    span,
     tracer
   );
 }

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

@@ -12,11 +12,11 @@ export enum SqsAttributeNames {
 }
 
 export const START_SPAN_FUNCTION = Symbol(
-  "opentelemetry.plugin.aws-sdk.start_span"
+  "opentelemetry.plugin.aws-sdk.sqs.start_span"
 );
 
 export const END_SPAN_FUNCTION = Symbol(
-  "opentelemetry.plugin.aws-sdk.end_span"
+  "opentelemetry.plugin.aws-sdk.sqs.end_span"
 );
 
 function extractQueueUrl(request: AWS.Request<any, any>): string {
@@ -73,52 +73,58 @@ function patchArrayFunction(
   messages[functionName] = function (callback) {
     return origFunc.call(this, function (message: AWS.SQS.Message) {
       const messageSpan = message[START_SPAN_FUNCTION]();
-      tracer.withSpan(messageSpan, () => callback.apply(this, arguments));
+      const res = tracer.withSpan(messageSpan, () =>
+        callback.apply(this, arguments)
+      );
       message[END_SPAN_FUNCTION]();
+      return res;
     });
   };
 }
 
 export function getSqsRequestSpanAttributes(
-  request: AWS.Request<any, any>
+  request: AWS.Request<any, any>,
+  span: Span
 ): RequestMetadata {
   const queueUrl = extractQueueUrl(request);
   const queueName = extractQueueNameFromUrl(queueUrl);
 
+  span.setAttribute(SqsAttributeNames.MESSAGING_SYSTEM, "aws.sqs");
+  span.setAttribute(SqsAttributeNames.MESSAGING_DESTINATIONKIND, "queue");
+  span.setAttribute(SqsAttributeNames.MESSAGING_DESTINATION, queueName);
+  span.setAttribute(SqsAttributeNames.MESSAGING_URL, queueUrl);
+
   let isIncoming = false;
-  const attributes = {
-    [SqsAttributeNames.MESSAGING_SYSTEM]: "aws.sqs",
-    [SqsAttributeNames.MESSAGING_DESTINATIONKIND]: "queue",
-    [SqsAttributeNames.MESSAGING_DESTINATION]: queueName,
-    [SqsAttributeNames.MESSAGING_URL]: queueUrl,
-  };
 
   const operation = (request as any)?.operation;
   switch (operation) {
     case "receiveMessage":
       isIncoming = true;
-      attributes[SqsAttributeNames.MESSAGING_OPERATION] = "receive";
+      span.setAttribute(SqsAttributeNames.MESSAGING_OPERATION, "receive");
       break;
   }
 
   return {
-    attributes,
     isIncoming,
   };
 }
 
 export function getSqsResponseSpanAttributes(
   response: AWS.Response<any, any>,
+  span: Span,
   tracer: Tracer
-): Attributes {
+) {
   const messages: AWS.SQS.Message[] = response.data.Messages;
   if (messages) {
     const queueUrl = extractQueueUrl((response as any)?.request);
     const queueName = extractQueueNameFromUrl(queueUrl);
 
     messages.forEach((message: AWS.SQS.Message) => {
-      message[START_SPAN_FUNCTION] = () =>
-        startSingleMessageSpan(tracer, queueUrl, queueName, message);
+      message[START_SPAN_FUNCTION] = () => {
+        return tracer.withSpan(span, () =>
+          startSingleMessageSpan(tracer, queueUrl, queueName, message)
+        );
+      };
       message[END_SPAN_FUNCTION] = () =>
         console.log(
           "open-telemetry aws-sdk plugin: end span called on sqs message which was not started"
@@ -128,5 +134,4 @@ export function getSqsResponseSpanAttributes(
     patchArrayFunction(messages, "forEach", tracer);
     patchArrayFunction(messages, "map", tracer);
   }
-  return {};
 }