refacotr(plugin-aws-sdk): call request hook before span is started

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,6 +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,
+        ...additionalAttributes,
       },
     });
 
@@ -138,8 +142,13 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         return original.apply(this, arguments);
       }
 
-      const span = thisPlugin._startAwsSpan(awsRequest);
-      thisPlugin.servicesExtensions.requestHook(awsRequest, span);
+      const requestMetadata = thisPlugin.servicesExtensions.requestHook(
+        awsRequest
+      );
+      const span = thisPlugin._startAwsSpan(
+        awsRequest,
+        requestMetadata.spanAttributes
+      );
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 
@@ -165,10 +174,12 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         return original.apply(this, arguments);
       }
 
-      const span = thisPlugin._startAwsSpan(awsRequest);
       const requestMetadata = thisPlugin.servicesExtensions.requestHook(
+        awsRequest
+      );
+      const span = thisPlugin._startAwsSpan(
         awsRequest,
-        span
+        requestMetadata.spanAttributes
       );
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);

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

@@ -1,13 +1,15 @@
-import { Span } from "@opentelemetry/api";
+import { Span, Attributes } from "@opentelemetry/api";
 
 /*
 isIncoming - if true, then the operation callback / promise should be bind with the operation's span 
 */
 export interface RequestMetadata {
   isIncoming: boolean;
+  spanAttributes?: Attributes;
 }
 
 export interface ServiceExtension {
-  requestHook: (request: AWS.Request<any, any>, span: Span) => RequestMetadata;
+  // called before request is sent, and before span is created
+  requestHook: (request: AWS.Request<any, any>) => RequestMetadata;
   responseHook: (response: AWS.Response<any, any>, span: Span) => void;
 }

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

@@ -10,11 +10,14 @@ export class ServicesExtensions implements ServiceExtension {
     this.services.set("sqs", new SqsServiceExtension(tracer));
   }
 
-  requestHook(request: AWS.Request<any, any>, span: Span): RequestMetadata {
+  requestHook(request: AWS.Request<any, any>): RequestMetadata {
     const serviceId = (request as any)?.service?.serviceIdentifier;
     const serviceExtension = this.services.get(serviceId);
-    if (!serviceExtension) return { isIncoming: false };
-    return serviceExtension.requestHook(request, span);
+    if (!serviceExtension)
+      return {
+        isIncoming: false,
+      };
+    return serviceExtension.requestHook(request);
   }
 
   responseHook(response: AWS.Response<any, any>, span: Span) {

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

@@ -27,27 +27,30 @@ export class SqsServiceExtension implements ServiceExtension {
     this.tracer = tracer;
   }
 
-  requestHook(request: AWS.Request<any, any>, span: Span): RequestMetadata {
+  requestHook(request: AWS.Request<any, any>): RequestMetadata {
     const queueUrl = this.extractQueueUrl(request);
     const queueName = this.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);
+    const spanAttributes = {
+      [SqsAttributeNames.MESSAGING_SYSTEM]: "aws.sqs",
+      [SqsAttributeNames.MESSAGING_DESTINATIONKIND]: "queue",
+      [SqsAttributeNames.MESSAGING_DESTINATION]: queueName,
+      [SqsAttributeNames.MESSAGING_URL]: queueUrl,
+    };
 
     let isIncoming = false;
 
     const operation = (request as any)?.operation;
     switch (operation) {
       case "receiveMessage":
         isIncoming = true;
-        span.setAttribute(SqsAttributeNames.MESSAGING_OPERATION, "receive");
+        spanAttributes[SqsAttributeNames.MESSAGING_OPERATION] = "receive";
         break;
     }
 
     return {
       isIncoming,
+      spanAttributes,
     };
   }