feat(plugin-aws-sdk): bind promise only on incoming operations

Author: Amir Blum

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

@@ -9,7 +9,7 @@
     callback    |       1           |       2   
  */
 import { BasePlugin } from "@opentelemetry/core";
-import { Span, CanonicalCode } from "@opentelemetry/api";
+import { Span, CanonicalCode, Attributes } from "@opentelemetry/api";
 import * as shimmer from "shimmer";
 import AWS from "aws-sdk";
 import { AttributeNames } from "./enums";
@@ -69,7 +69,10 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
     return target;
   }
 
-  private _startAwsSpan(request: AWS.Request<any, any>): Span {
+  private _startAwsSpan(
+    request: AWS.Request<any, any>,
+    additionalAttributes?: Attributes
+  ): Span {
     const operation = (request as any).operation;
     const service = (request as any).service;
 
@@ -83,7 +86,7 @@ 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,
-        ...getRequestServiceAttributes(request),
+        ...additionalAttributes,
       },
     });
 
@@ -138,7 +141,11 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         return original.apply(this, arguments);
       }
 
-      const span = thisPlugin._startAwsSpan(awsRequest);
+      const requestMetadata = getRequestServiceAttributes(awsRequest);
+      const span = thisPlugin._startAwsSpan(
+        awsRequest,
+        requestMetadata?.attributes
+      );
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 
@@ -164,17 +171,24 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         return original.apply(this, arguments);
       }
 
-      const span = thisPlugin._startAwsSpan(awsRequest);
+      const requestMetadata = getRequestServiceAttributes(awsRequest);
+      const span = thisPlugin._startAwsSpan(
+        awsRequest,
+        requestMetadata?.attributes
+      );
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 
-      const promiseToReturn: Promise<any> = thisPlugin._tracer.withSpan(
+      const origPromise: Promise<any> = thisPlugin._tracer.withSpan(
         span,
         () => {
           return original.apply(awsRequest, arguments);
         }
       );
-      return thisPlugin._bindPromise(promiseToReturn, span);
+
+      return requestMetadata.isIncoming
+        ? thisPlugin._bindPromise(origPromise, span)
+        : origPromise;
     };
   }
 

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

@@ -1,9 +1,13 @@
 import { Attributes } from "@opentelemetry/api";
+import { RequestMetadata } from "./service-attributes";
 
 export function getS3RequestSpanAttributes(
   request: AWS.Request<any, any>
-): Attributes {
-  return {};
+): RequestMetadata {
+  return {
+    attributes: {},
+    isIncoming: false,
+  };
 }
 
 export function getS3ResponseSpanAttributes(

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

@@ -1,7 +1,20 @@
 import { Attributes } from "@opentelemetry/api";
 import { getS3RequestSpanAttributes, getS3ResponseSpanAttributes } from "./s3";
+import {
+  getSqsRequestSpanAttributes,
+  getSqsResponseSpanAttributes,
+} from "./sqs";
 
-type RequestAttrProcessor = (request: AWS.Request<any, any>) => Attributes;
+/*
+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 ResponseAttrProcessor = (response: AWS.Response<any, any>) => Attributes;
 
 class ServiceAttributes {
@@ -12,7 +25,7 @@ class ServiceAttributes {
     };
   } = {};
 
-  public initAttributeProcessor = (
+  public addAttributeProcessor = (
     serviceId: string,
     requestProcessor: RequestAttrProcessor,
     responseProcessor: ResponseAttrProcessor
@@ -26,7 +39,7 @@ class ServiceAttributes {
 
 export function getRequestServiceAttributes(
   request: AWS.Request<any, any>
-): Attributes {
+): RequestMetadata {
   const serviceId = (request as any)?.service?.serviceIdentifier;
   if (serviceId) {
     return serviceAttributes.attributeProcessors[serviceId]?.request(request);
@@ -37,15 +50,21 @@ export function getResponseServiceAttributes(
   response: AWS.Response<any, any>
 ): Attributes {
   const serviceId = (response as any)?.request?.service?.serviceIdentifier;
-  if (serviceId) {
-    return serviceAttributes.attributeProcessors[serviceId]?.response(response);
-  }
+  if (!serviceId) return;
+
+  return serviceAttributes.attributeProcessors[serviceId]?.response(response);
 }
 
 export const serviceAttributes = new ServiceAttributes();
 
-serviceAttributes.initAttributeProcessor(
+serviceAttributes.addAttributeProcessor(
   "s3",
   getS3RequestSpanAttributes,
   getS3ResponseSpanAttributes
 );
+
+serviceAttributes.addAttributeProcessor(
+  "sqs",
+  getSqsRequestSpanAttributes,
+  getSqsResponseSpanAttributes
+);

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

@@ -0,0 +1,26 @@
+import { Attributes } from "@opentelemetry/api";
+import { RequestMetadata } from "./service-attributes";
+
+export function getSqsRequestSpanAttributes(
+  request: AWS.Request<any, any>
+): RequestMetadata {
+  const operation = (request as any)?.operation;
+  switch (operation) {
+    case "receiveMessage":
+      return {
+        attributes: {},
+        isIncoming: true,
+      };
+  }
+
+  return {
+    attributes: {},
+    isIncoming: false,
+  };
+}
+
+export function getSqsResponseSpanAttributes(
+  response: AWS.Response<any, any>
+): Attributes {
+  return {};
+}