feat(aws-sdk): add exception hook (#2398)

Author: project0

packages/instrumentation-aws-sdk/README.md

@@ -52,6 +52,7 @@ aws-sdk instrumentation has few options available to choose from. You can set th
 | ----------------------------------------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | `preRequestHook`                          | `AwsSdkRequestCustomAttributeFunction`    | Hook called before request send, which allow to add custom attributes to span.                                                                                                  |
 | `responseHook`                            | `AwsSdkResponseCustomAttributeFunction`   | Hook for adding custom attributes when response is received from aws.                                                                                                           |
+| `exceptionHook`                           | `AwsSdkExceptionCustomAttributeFunction`  | Hook for adding custom attributes when exception is received from aws.                                                                                                          |
 | `sqsProcessHook`                          | `AwsSdkSqsProcessCustomAttributeFunction` | Hook called after starting sqs `process` span (for each sqs received message), which allow to add custom attributes to it.                                                      |
 | `suppressInternalInstrumentation`         | `boolean`                                 | Most aws operation use http requests under the hood. Set this to `true` to hide all underlying http spans.                                                                      |
 | `sqsExtractContextPropagationFromPayload` | `boolean`                                 | Will parse and extract context propagation headers from SQS Payload, false by default. [When should it be used?](./doc/sns.md#integration-with-sqs)                             |

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

@@ -233,6 +233,30 @@ export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentatio
     );
   }
 
+  private _callUserExceptionResponseHook(
+    span: Span,
+    request: NormalizedRequest,
+    err: any
+  ) {
+    const { exceptionHook } = this.getConfig();
+    if (!exceptionHook) return;
+    const requestInfo: AwsSdkRequestHookInformation = {
+      request,
+    };
+
+    safeExecuteInTheMiddle(
+      () => exceptionHook(span, requestInfo, err),
+      (e: Error | undefined) => {
+        if (e)
+          diag.error(
+            `${AwsInstrumentation.component} instrumentation: exceptionHook error`,
+            e
+          );
+      },
+      true
+    );
+  }
+
   private _getV3ConstructStackPatch(
     moduleVersion: string | undefined,
     original: (...args: unknown[]) => MiddlewareStack<any, any>
@@ -443,6 +467,11 @@ export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentatio
                     message: err.message,
                   });
                   span.recordException(err);
+                  self._callUserExceptionResponseHook(
+                    span,
+                    normalizedRequest,
+                    err
+                  );
                   throw err;
                 })
                 .finally(() => {

packages/instrumentation-aws-sdk/src/types.ts

@@ -48,6 +48,7 @@ export interface AwsSdkResponseHookInformation {
   moduleVersion?: string;
   response: NormalizedResponse;
 }
+
 /**
  * span can be used to add custom attributes, or for any other need.
  * response is the object that is returned to the user calling the aws-sdk operation.
@@ -57,6 +58,14 @@ export interface AwsSdkResponseCustomAttributeFunction {
   (span: Span, responseInfo: AwsSdkResponseHookInformation): void;
 }
 
+/**
+ * span can be used to modify the status.
+ * As we have no response object, the request can be used to determine the failed aws-sdk operation.
+ */
+export interface AwsSdkExceptionCustomAttributeFunction {
+  (span: Span, requestInfo: AwsSdkRequestHookInformation, err: any): void;
+}
+
 export interface AwsSdkSqsProcessHookInformation {
   message: SQS.Message;
 }
@@ -76,6 +85,12 @@ export interface AwsSdkInstrumentationConfig extends InstrumentationConfig {
   /** hook for adding custom attributes when response is received from aws */
   responseHook?: AwsSdkResponseCustomAttributeFunction;
 
+  /**
+   * Hook for adding custom attributes when exception is received from aws.
+   * This hook is only available with aws sdk v3
+   */
+  exceptionHook?: AwsSdkExceptionCustomAttributeFunction;
+
   /** hook for adding custom attribute when an sqs process span is started */
   sqsProcessHook?: AwsSdkSqsProcessCustomAttributeFunction;
 

packages/instrumentation-aws-sdk/test/aws-sdk-v3-s3.test.ts

@@ -270,5 +270,48 @@ describe('instrumentation-aws-sdk-v3 (client-s3)', () => {
         );
       });
     });
+
+    it('handle aws sdk exception', async () => {
+      instrumentation.disable();
+      instrumentation.setConfig({
+        exceptionHook: (
+          span: Span,
+          requestInfo: AwsSdkRequestHookInformation,
+          err: any
+        ) => {
+          span.setAttribute(
+            'attribute.from.exception.hook',
+            requestInfo.request.commandInput.Bucket
+          );
+          span.setAttribute('error.from.exception.hook', err.name);
+        },
+        suppressInternalInstrumentation: true,
+      });
+      instrumentation.enable();
+
+      nock(`https://ot-demo-test.s3.${region}.amazonaws.com/`)
+        .put('/aws-ot-s3-test-object.txt?x-id=PutObject')
+        .reply(
+          404,
+          fs.readFileSync('./test/mock-responses/s3-put-object.xml', 'utf8')
+        );
+
+      const params = {
+        Bucket: 'ot-demo-test',
+        Key: 'aws-ot-s3-test-object.txt',
+      };
+      try {
+        await s3Client.putObject(params);
+      } catch {
+        expect(getTestSpans().length).toBe(1);
+        const [span] = getTestSpans();
+        expect(span.attributes['attribute.from.exception.hook']).toEqual(
+          params.Bucket
+        );
+        expect(span.attributes['error.from.exception.hook']).toEqual(
+          'NotFound'
+        );
+      }
+    });
   });
 });