Merge pull request #15 from aspecto-io/request-hook-refactor

feat(plugin-aws-sdk): request hook enhancment

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, SpanKind } from "@opentelemetry/api";
 import * as shimmer from "shimmer";
 import AWS from "aws-sdk";
 import { AttributeNames } from "./enums";
@@ -69,11 +69,18 @@ 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,
+    spanKind?: SpanKind,
+    spanName?: string
+  ): Span {
     const operation = (request as any).operation;
     const service = (request as any).service;
+    const name = spanName ?? this._getSpanName(request);
 
-    const newSpan = this._tracer.startSpan(this._getSpanName(request), {
+    const newSpan = this._tracer.startSpan(name, {
+      kind: spanKind,
       attributes: {
         [AttributeNames.COMPONENT]: this.moduleName,
         [AttributeNames.AWS_OPERATION]: operation,
@@ -83,6 +90,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 +146,15 @@ 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,
+        requestMetadata.spanKind,
+        requestMetadata.spanName
+      );
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 
@@ -165,10 +180,14 @@ 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,
+        requestMetadata.spanKind,
+        requestMetadata.spanName
       );
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);

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

@@ -1,13 +1,17 @@
-import { Span } from "@opentelemetry/api";
+import { Span, Attributes, SpanKind } 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;
+  spanKind?: SpanKind;
+  spanName?: string;
 }
 
 export interface ServiceExtension {
-  requestHook: (request: AWS.Request<any, any>, span: Span) => RequestMetadata;
+  // called before request is sent, and before span is started
+  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,32 +27,47 @@ 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);
+    let spanKind: SpanKind = SpanKind.CLIENT;
+    let spanName: string;
+
+    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");
+        spanKind = SpanKind.CONSUMER;
+        spanName = queueName;
+        spanAttributes[SqsAttributeNames.MESSAGING_OPERATION] = "receive";
+        break;
+
+      case "sendMessage":
+      case "sendMessageBatch":
+        spanKind = SpanKind.PRODUCER;
+        spanName = queueName;
         break;
     }
 
     return {
       isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
     };
   }
 
   responseHook = (response: AWS.Response<any, any>, span: Span) => {
-    const messages: AWS.SQS.Message[] = response.data.Messages;
+    const messages: AWS.SQS.Message[] = response?.data?.Messages;
     if (messages) {
       const queueUrl = this.extractQueueUrl((response as any)?.request);
       const queueName = this.extractQueueNameFromUrl(queueUrl);