Merge pull request #19 from aspecto-io/sqs-inject-span-bug-fix

fix(opentelemetry-plugin-aws-sdk): set aws span as the propagated context

Author: Amir Blum

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

@@ -163,6 +163,7 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
 
       const callbackWithContext = thisPlugin._tracer.bind(callback, span);
       return thisPlugin._tracer.withSpan(span, () => {
+        thisPlugin.servicesExtensions.requestPostSpanHook(awsRequest);
         return original.call(awsRequest, callbackWithContext);
       });
     };
@@ -198,6 +199,7 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
       const origPromise: Promise<any> = thisPlugin._tracer.withSpan(
         span,
         () => {
+          thisPlugin.servicesExtensions.requestPostSpanHook(awsRequest);
           return original.apply(awsRequest, arguments);
         }
       );

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

@@ -13,5 +13,9 @@ export interface RequestMetadata {
 export interface ServiceExtension {
   // called before request is sent, and before span is started
   requestHook: (request: AWS.Request<any, any>) => RequestMetadata;
+
+  // called before request is sent, and after span is started
+  requestPostSpanHook?: (request: AWS.Request<any, any>) => void;
+
   responseHook: (response: AWS.Response<any, any>, span: Span) => void;
 }

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

@@ -20,6 +20,13 @@ export class ServicesExtensions implements ServiceExtension {
     return serviceExtension.requestHook(request);
   }
 
+  requestPostSpanHook(request: AWS.Request<any, any>) {
+    const serviceId = (request as any)?.service?.serviceIdentifier;
+    const serviceExtension = this.services.get(serviceId);
+    if (!serviceExtension?.requestPostSpanHook) return;
+    return serviceExtension.requestPostSpanHook(request);
+  }
+
   responseHook(response: AWS.Response<any, any>, span: Span) {
     const serviceId = (response as any)?.request?.service?.serviceIdentifier;
     const serviceExtension = this.services.get(serviceId);

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

@@ -100,10 +100,26 @@ export class SqsServiceExtension implements ServiceExtension {
         break;
 
       case "sendMessage":
-        {
-          spanKind = SpanKind.PRODUCER;
-          spanName = queueName;
+      case "sendMessageBatch":
+        spanKind = SpanKind.PRODUCER;
+        spanName = queueName;
+        break;
+    }
+
+    return {
+      isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
 
+  requestPostSpanHook = (request: AWS.Request<any, any>) => {
+    const operation = (request as any)?.operation;
+    console.log(operation);
+    switch (operation) {
+      case "sendMessage":
+        {
           const params: SendMessageRequest = (request as any).params;
           params.MessageAttributes = this.InjectPropagationContext(
             params.MessageAttributes
@@ -113,9 +129,6 @@ export class SqsServiceExtension implements ServiceExtension {
 
       case "sendMessageBatch":
         {
-          spanKind = SpanKind.PRODUCER;
-          spanName = queueName;
-
           const params: SendMessageBatchRequest = (request as any).params;
           params.Entries.forEach(
             (messageParams: SendMessageBatchRequestEntry) => {
@@ -127,14 +140,7 @@ export class SqsServiceExtension implements ServiceExtension {
         }
         break;
     }
-
-    return {
-      isIncoming,
-      spanAttributes,
-      spanKind,
-      spanName,
-    };
-  }
+  };
 
   responseHook = (response: AWS.Response<any, any>, span: Span) => {
     const messages: AWS.SQS.Message[] = response?.data?.Messages;