feat(aws-lambda): added request and response hooks (open-telemetry#486)

Co-authored-by: Valentin Marchaud <[email protected]>

plugins/node/opentelemetry-instrumentation-aws-lambda/README.md

@@ -29,7 +29,9 @@ provider.register();
 
 registerInstrumentations({
   instrumentations: [
-    new AwsLambdaInstrumentation()
+    new AwsLambdaInstrumentation({
+        // see under for available configuration
+    })
   ],
 });
 ```
@@ -38,6 +40,29 @@ In your Lambda function configuration, add or update the `NODE_OPTIONS` environm
 
 `NODE_OPTIONS=--require lambda-wrapper`
 
+## AWS Lambda Instrumentation Options
+
+| Options | Type  | Description |
+| --- | --- | --- |
+| `requestHook` | `RequestHook` (function) | Hook for adding custom attributes before lambda starts handling the request. Receives params: `span, { event, context }` |
+| `responseHook` | `ResponseHook` (function) | Hook for adding custom attributes before lambda returns the response. Receives params: `span, { err?, res? } ` |
+
+### Hooks Usage Example
+
+```js
+const { AwsLambdaInstrumentation } = require('@opentelemetry/instrumentation-aws-lambda');
+
+new AwsLambdaInstrumentation({
+    requestHook: (span, { event, context }) => {
+        span.setAttributes('faas.name', context.functionName);
+    },
+    responseHook: (span, { err, res }) => {
+        if (err instanceof Error) span.setAttributes('faas.error', err.message);
+        if (res) span.setAttributes('faas.res', res);
+    }
+})
+```
+
 ## Useful links
 
 - For more information on OpenTelemetry, visit: <https://opentelemetry.io/>

plugins/node/opentelemetry-instrumentation-aws-lambda/src/index.ts

@@ -15,3 +15,8 @@
  */
 
 export * from './instrumentation';
+export {
+  AwsLambdaInstrumentationConfig,
+  RequestHook,
+  ResponseHook,
+} from './types';

plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts

@@ -54,7 +54,7 @@ import {
   Handler,
 } from 'aws-lambda';
 
-import { LambdaModule } from './types';
+import { LambdaModule, AwsLambdaInstrumentationConfig } from './types';
 import { VERSION } from './version';
 
 const awsPropagator = new AWSXRayPropagator();
@@ -72,8 +72,12 @@ export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
 export class AwsLambdaInstrumentation extends InstrumentationBase {
   private _tracerProvider: TracerProvider | undefined;
 
-  constructor() {
-    super('@opentelemetry/instrumentation-aws-lambda', VERSION);
+  constructor(protected _config: AwsLambdaInstrumentationConfig = {}) {
+    super('@opentelemetry/instrumentation-aws-lambda', VERSION, _config);
+  }
+
+  setConfig(config: AwsLambdaInstrumentationConfig = {}) {
+    this._config = config;
   }
 
   init() {
@@ -167,6 +171,17 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
         parent
       );
 
+      if (plugin._config.requestHook) {
+        safeExecuteInTheMiddle(
+          () => plugin._config.requestHook!(span, { event, context }),
+          e => {
+            if (e)
+              diag.error('aws-lambda instrumentation: requestHook error', e);
+          },
+          true
+        );
+      }
+
       return otelContext.with(setSpan(otelContext.active(), span), () => {
         // Lambda seems to pass a callback even if handler is of Promise form, so we wrap all the time before calling
         // the handler and see if the result is a Promise or not. In such a case, the callback is usually ignored. If
@@ -178,20 +193,25 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
           error => {
             if (error != null) {
               // Exception thrown synchronously before resolving callback / promise.
+              plugin._applyResponseHook(span, error);
               plugin._endSpan(span, error, () => {});
             }
           }
         ) as Promise<{}> | undefined;
         if (typeof maybePromise?.then === 'function') {
           return maybePromise.then(
-            value =>
-              new Promise(resolve =>
+            value => {
+              plugin._applyResponseHook(span, null, value);
+              return new Promise(resolve =>
                 plugin._endSpan(span, undefined, () => resolve(value))
-              ),
-            (err: Error | string) =>
-              new Promise((resolve, reject) =>
+              );
+            },
+            (err: Error | string) => {
+              plugin._applyResponseHook(span, err);
+              return new Promise((resolve, reject) =>
                 plugin._endSpan(span, err, () => reject(err))
-              )
+              );
+            }
           );
         }
         return maybePromise;
@@ -208,6 +228,7 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
     const plugin = this;
     return function wrappedCallback(this: never, err, res) {
       diag.debug('executing wrapped lookup callback function');
+      plugin._applyResponseHook(span, err, res);
 
       plugin._endSpan(span, err, () => {
         diag.debug('executing original lookup callback function');
@@ -252,6 +273,23 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
     }
   }
 
+  private _applyResponseHook(
+    span: Span,
+    err?: Error | string | null,
+    res?: any
+  ) {
+    if (this._config?.responseHook) {
+      safeExecuteInTheMiddle(
+        () => this._config.responseHook!(span, { err, res }),
+        e => {
+          if (e)
+            diag.error('aws-lambda instrumentation: responseHook error', e);
+        },
+        true
+      );
+    }
+  }
+
   private static _extractAccountId(arn: string): string | undefined {
     const parts = arn.split(':');
     if (parts.length >= 5) {

plugins/node/opentelemetry-instrumentation-aws-lambda/src/types.ts

@@ -14,6 +14,26 @@
  * limitations under the License.
  */
 
-import { Handler } from 'aws-lambda';
+import { Span } from '@opentelemetry/api';
+import { InstrumentationConfig } from '@opentelemetry/instrumentation';
+import { Handler, Context } from 'aws-lambda';
 
 export type LambdaModule = Record<string, Handler>;
+
+export type RequestHook = (
+  span: Span,
+  hookInfo: { event: any; context: Context }
+) => void;
+
+export type ResponseHook = (
+  span: Span,
+  hookInfo: {
+    err?: Error | string | null;
+    res?: any;
+  }
+) => void;
+
+export interface AwsLambdaInstrumentationConfig extends InstrumentationConfig {
+  requestHook?: RequestHook;
+  responseHook?: ResponseHook;
+}

plugins/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts

@@ -21,6 +21,7 @@ import * as path from 'path';
 
 import {
   AwsLambdaInstrumentation,
+  AwsLambdaInstrumentationConfig,
   traceContextEnvironmentKey,
 } from '../../src';
 import {
@@ -31,6 +32,10 @@ import {
 import { NodeTracerProvider } from '@opentelemetry/node';
 import { Context } from 'aws-lambda';
 import * as assert from 'assert';
+import {
+  SemanticAttributes,
+  ResourceAttributes,
+} from '@opentelemetry/semantic-conventions';
 import {
   context,
   setSpanContext,
@@ -41,7 +46,6 @@ import {
 } from '@opentelemetry/api';
 import { AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
 import { HttpTraceContext } from '@opentelemetry/core';
-import { SemanticAttributes } from '@opentelemetry/semantic-conventions';
 
 const memoryExporter = new InMemorySpanExporter();
 const provider = new NodeTracerProvider();
@@ -105,10 +109,13 @@ describe('lambda handler', () => {
     awsRequestId: 'aws_request_id',
   } as Context;
 
-  const initializeHandler = (handler: string) => {
+  const initializeHandler = (
+    handler: string,
+    config: AwsLambdaInstrumentationConfig = {}
+  ) => {
     process.env._HANDLER = handler;
 
-    instrumentation = new AwsLambdaInstrumentation();
+    instrumentation = new AwsLambdaInstrumentation(config);
     instrumentation.setTracerProvider(provider);
   };
 
@@ -520,4 +527,117 @@ describe('lambda handler', () => {
       assert.strictEqual(spans.length, 0);
     });
   });
+
+  describe('hooks', () => {
+    describe('requestHook', () => {
+      it('sync - success', async () => {
+        initializeHandler('lambda-test/async.handler', {
+          requestHook: (span, { context }) => {
+            span.setAttribute(
+              ResourceAttributes.FAAS_NAME,
+              context.functionName
+            );
+          },
+        });
+
+        await lambdaRequire('lambda-test/async').handler('arg', ctx);
+        const spans = memoryExporter.getFinishedSpans();
+        const [span] = spans;
+        assert.strictEqual(spans.length, 1);
+        assert.strictEqual(
+          span.attributes[ResourceAttributes.FAAS_NAME],
+          ctx.functionName
+        );
+        assertSpanSuccess(span);
+      });
+    });
+
+    describe('responseHook', () => {
+      const RES_ATTR = 'test.res';
+      const ERR_ATTR = 'test.error';
+
+      const config: AwsLambdaInstrumentationConfig = {
+        responseHook: (span, { err, res }) => {
+          if (err)
+            span.setAttribute(
+              ERR_ATTR,
+              typeof err === 'string' ? err : err.message
+            );
+          if (res)
+            span.setAttribute(
+              RES_ATTR,
+              typeof res === 'string' ? res : JSON.stringify(res)
+            );
+        },
+      };
+      it('async - success', async () => {
+        initializeHandler('lambda-test/async.handler', config);
+
+        const res = await lambdaRequire('lambda-test/async').handler(
+          'arg',
+          ctx
+        );
+        const [span] = memoryExporter.getFinishedSpans();
+        assert.strictEqual(span.attributes[RES_ATTR], res);
+      });
+
+      it('async - error', async () => {
+        initializeHandler('lambda-test/async.error', config);
+
+        let err: Error;
+        try {
+          await lambdaRequire('lambda-test/async').error('arg', ctx);
+        } catch (e) {
+          err = e;
+        }
+        const [span] = memoryExporter.getFinishedSpans();
+        assert.strictEqual(span.attributes[ERR_ATTR], err!.message);
+      });
+
+      it('sync - success', async () => {
+        initializeHandler('lambda-test/sync.handler', config);
+
+        const result = await new Promise((resolve, _reject) => {
+          lambdaRequire('lambda-test/sync').handler(
+            'arg',
+            ctx,
+            (_err: Error, res: any) => resolve(res)
+          );
+        });
+        const [span] = memoryExporter.getFinishedSpans();
+        assert.strictEqual(span.attributes[RES_ATTR], result);
+      });
+
+      it('sync - error', async () => {
+        initializeHandler('lambda-test/sync.error', config);
+
+        let err: Error;
+        try {
+          lambdaRequire('lambda-test/sync').error('arg', ctx, () => {});
+        } catch (e) {
+          err = e;
+        }
+        const [span] = memoryExporter.getFinishedSpans();
+        assert.strictEqual(span.attributes[ERR_ATTR], err!.message);
+      });
+
+      it('sync - error with callback', async () => {
+        initializeHandler('lambda-test/sync.callbackerror', config);
+
+        let error: Error;
+        await new Promise((resolve, _reject) => {
+          lambdaRequire('lambda-test/sync').callbackerror(
+            'arg',
+            ctx,
+            (err: Error, _res: any) => {
+              error = err;
+              resolve({});
+            }
+          );
+        });
+        const [span] = memoryExporter.getFinishedSpans();
+        assert.strictEqual(span.attributes[ERR_ATTR], error!.message);
+      });
+    });
+  });
 });