Merge pull request #10 from aspecto-io/aws-cb-context

feat(plugin-aws-sdk): sqs proper context and semantic conventions

Author: Amir Blum

packages/plugin-aws-sdk/package.json

@@ -1,6 +1,6 @@
 {
   "name": "opentelemetry-plugin-aws-sdk",
-  "version": "0.0.3",
+  "version": "0.0.4",
   "description": "open telemetry instrumentation for the `aws-sdk` package",
   "keywords": [
     "aws",

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

@@ -13,24 +13,24 @@ import { Span, CanonicalCode } from "@opentelemetry/api";
 import * as shimmer from "shimmer";
 import AWS from "aws-sdk";
 import { AttributeNames } from "./enums";
-import {
-  getRequestServiceAttributes,
-  getResponseServiceAttributes,
-} from "./services";
+import { ServicesExtensions } from "./services";
 import { AwsSdkPluginConfig } from "./types";
 
-const VERSION = "0.0.3";
+const VERSION = "0.0.4";
 
 class AwsPlugin extends BasePlugin<typeof AWS> {
   readonly component: string;
   protected _config: AwsSdkPluginConfig;
-  private activeRequests: Set<AWS.Request<any, any>> = new Set();
+  private REQUEST_SPAN_KEY = Symbol("opentelemetry.plugin.aws-sdk.span");
+  private servicesExtensions: ServicesExtensions;
 
   constructor(readonly moduleName: string) {
     super(`opentelemetry-plugin-aws-sdk`, VERSION);
   }
 
   protected patch() {
+    this.servicesExtensions = new ServicesExtensions(this._tracer);
+
     this._logger.debug(
       "applying patch to %s@%s",
       this.moduleName,
@@ -40,12 +40,12 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
     shimmer.wrap(
       this._moduleExports?.Request.prototype,
       "send",
-      this._patchRequestMethod()
+      this._getRequestSendPatch.bind(this)
     );
     shimmer.wrap(
       this._moduleExports?.Request.prototype,
       "promise",
-      this._patchRequestMethod()
+      this._getRequestPromisePatch.bind(this)
     );
 
     return this._moduleExports;
@@ -56,73 +56,135 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
     shimmer.unwrap(this._moduleExports?.Request.prototype, "promise");
   }
 
-  private _patchRequestMethod = () => {
-    return this._getPatchedRequestMethod;
-  };
+  private _bindPromise(target: Promise<any>, span: Span) {
+    const thisPlugin = this;
+
+    const origThen = target.then;
+    target.then = function (onFulfilled, onRejected) {
+      const newOnFulfilled = thisPlugin._tracer.bind(onFulfilled, span);
+      const newOnRejected = thisPlugin._tracer.bind(onRejected, span);
+      return origThen.call(this, newOnFulfilled, newOnRejected);
+    };
+
+    return target;
+  }
+
+  private _startAwsSpan(request: AWS.Request<any, any>): Span {
+    const operation = (request as any).operation;
+    const service = (request as any).service;
+
+    const newSpan = this._tracer.startSpan(this._getSpanName(request), {
+      attributes: {
+        [AttributeNames.COMPONENT]: this.moduleName,
+        [AttributeNames.AWS_OPERATION]: operation,
+        [AttributeNames.AWS_SIGNATURE_VERSION]:
+          service?.config?.signatureVersion,
+        [AttributeNames.AWS_REGION]: service?.config?.region,
+        [AttributeNames.AWS_SERVICE_API]: service?.api?.className,
+        [AttributeNames.AWS_SERVICE_IDENTIFIER]: service?.serviceIdentifier,
+        [AttributeNames.AWS_SERVICE_NAME]: service?.api?.abbreviation,
+      },
+    });
+
+    request[this.REQUEST_SPAN_KEY] = newSpan;
+    return newSpan;
+  }
+
+  private _callPreRequestHooks(span: Span, request: AWS.Request<any, any>) {
+    if (this._config?.preRequestHook) {
+      this._safeExecute(
+        span,
+        () => this._config.preRequestHook(span, request),
+        false
+      );
+    }
+  }
 
-  private _getPatchedRequestMethod = (original: Function) => {
+  private _registerCompletedEvent(span: Span, request: AWS.Request<any, any>) {
     const thisPlugin = this;
-    return function () {
-      let span: Span | null = null;
+    request.on("complete", (response) => {
+      if (!request[thisPlugin.REQUEST_SPAN_KEY]) {
+        return;
+      }
+      request[thisPlugin.REQUEST_SPAN_KEY] = undefined;
+
+      if (response.error) {
+        span.setAttribute(AttributeNames.AWS_ERROR, response.error);
+      }
+
+      this.servicesExtensions.responseHook(response, span);
+
+      span.setAttributes({
+        [AttributeNames.AWS_REQUEST_ID]: response.requestId,
+      });
+      span.end();
+    });
+  }
+
+  private _getRequestSendPatch(
+    original: (callback?: (err: any, data: any) => void) => void
+  ) {
+    const thisPlugin = this;
+    return function (callback?: (err: any, data: any) => void) {
+      const awsRequest: AWS.Request<any, any> = this;
       /* 
         if the span was already started, we don't want to start a new one 
         when Request.promise() is called
       */
-
       if (
-        this._asm.currentState !== "complete" &&
-        !thisPlugin.activeRequests.has(this)
+        this._asm.currentState === "complete" ||
+        awsRequest[thisPlugin.REQUEST_SPAN_KEY]
       ) {
-        thisPlugin.activeRequests.add(this);
-
-        span = thisPlugin._tracer.startSpan(thisPlugin._getSpanName(this), {
-          attributes: {
-            [AttributeNames.COMPONENT]: thisPlugin.moduleName,
-            [AttributeNames.AWS_OPERATION]: this.operation,
-            [AttributeNames.AWS_SIGNATURE_VERSION]: this.service?.config
-              ?.signatureVersion,
-            [AttributeNames.AWS_REGION]: this.service?.config?.region,
-            [AttributeNames.AWS_SERVICE_API]: this.service?.api?.className,
-            [AttributeNames.AWS_SERVICE_IDENTIFIER]: this.service
-              ?.serviceIdentifier,
-            [AttributeNames.AWS_SERVICE_NAME]: this.service?.api?.abbreviation,
-            ...getRequestServiceAttributes(this),
-          },
-        });
-
-        if (thisPlugin._config?.preRequestHook) {
-          thisPlugin._safeExecute(
-            span,
-            () => thisPlugin._config.preRequestHook(span, this),
-            false
-          );
-        }
-
-        (this as AWS.Request<any, any>).on("complete", (response) => {
-          if (thisPlugin.activeRequests.has(this)) {
-            thisPlugin.activeRequests.delete(this);
-          }
-          if (!span) return;
-
-          if (response.error) {
-            span.setAttribute(AttributeNames.AWS_ERROR, response.error);
-          }
-
-          span.setAttributes({
-            [AttributeNames.AWS_REQUEST_ID]: response.requestId,
-            ...getResponseServiceAttributes(response),
-          });
-          span.end();
-          span = null;
-        });
+        return original.apply(this, arguments);
       }
 
-      const awsRequest = this;
+      const span = thisPlugin._startAwsSpan(awsRequest);
+      thisPlugin.servicesExtensions.requestHook(awsRequest, span);
+      thisPlugin._callPreRequestHooks(span, awsRequest);
+      thisPlugin._registerCompletedEvent(span, awsRequest);
+
+      const callbackWithContext = thisPlugin._tracer.bind(callback, span);
       return thisPlugin._tracer.withSpan(span, () => {
-        return original.apply(awsRequest, arguments);
+        return original.call(awsRequest, callbackWithContext);
       });
     };
-  };
+  }
+
+  private _getRequestPromisePatch(original: () => Promise<any>) {
+    const thisPlugin = this;
+    return function (): Promise<any> {
+      const awsRequest: AWS.Request<any, any> = this;
+      /* 
+        if the span was already started, we don't want to start a new one 
+        when Request.promise() is called
+      */
+      if (
+        this._asm.currentState === "complete" ||
+        awsRequest[thisPlugin.REQUEST_SPAN_KEY]
+      ) {
+        return original.apply(this, arguments);
+      }
+
+      const span = thisPlugin._startAwsSpan(awsRequest);
+      const requestMetadata = thisPlugin.servicesExtensions.requestHook(
+        awsRequest,
+        span
+      );
+      thisPlugin._callPreRequestHooks(span, awsRequest);
+      thisPlugin._registerCompletedEvent(span, awsRequest);
+
+      const origPromise: Promise<any> = thisPlugin._tracer.withSpan(
+        span,
+        () => {
+          return original.apply(awsRequest, arguments);
+        }
+      );
+
+      return requestMetadata.isIncoming
+        ? thisPlugin._bindPromise(origPromise, span)
+        : origPromise;
+    };
+  }
 
   private _getSpanName = (request: any) => {
     return `aws.${request.service?.serviceIdentifier ?? "request"}.${

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

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

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

@@ -0,0 +1,26 @@
+import { Tracer, Span } from "@opentelemetry/api";
+import { ServiceExtension, RequestMetadata } from "./ServiceExtension";
+import { SqsServiceExtension } from "./sqs";
+import * as AWS from "aws-sdk";
+
+export class ServicesExtensions implements ServiceExtension {
+  services: Map<string, ServiceExtension> = new Map();
+
+  constructor(tracer: Tracer) {
+    this.services.set("sqs", new SqsServiceExtension(tracer));
+  }
+
+  requestHook(request: AWS.Request<any, any>, span: Span): RequestMetadata {
+    const serviceId = (request as any)?.service?.serviceIdentifier;
+    const serviceExtension = this.services.get(serviceId);
+    if (!serviceExtension) return { isIncoming: false };
+    return serviceExtension.requestHook(request, span);
+  }
+
+  responseHook(response: AWS.Response<any, any>, span: Span) {
+    const serviceId = (response as any)?.request?.service?.serviceIdentifier;
+    const serviceExtension = this.services.get(serviceId);
+    if (!serviceExtension) return;
+    serviceExtension.responseHook(response, span);
+  }
+}

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

@@ -1,4 +1 @@
-export {
-  getRequestServiceAttributes,
-  getResponseServiceAttributes,
-} from "./service-attributes";
+export { ServicesExtensions } from "./ServicesExtensions";