feat(aws-lambda): added eventContextExtractor config option

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

@@ -47,6 +47,7 @@ In your Lambda function configuration, add or update the `NODE_OPTIONS` environm
 | `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? } ` |
 | `disableAwsContextPropagation` | `boolean` | By default, this instrumentation will try to read the context from the `_X_AMZN_TRACE_ID` environment variable set by Lambda, set this to `true` to disable this behavior |
+| `eventContextExtractor` | `EventContextExtractor` (function) | Function for providing custom context extractor in order to support different event types that are handled by AWS Lambda (e.g., SQS, CloudWatch, Kinesis, API Gateway). Applied only when `disableAwsContextPropagation` is set to `true`. Receives params: `event` |
 
 ### Hooks Usage Example
 

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

@@ -53,7 +53,11 @@ import {
   Handler,
 } from 'aws-lambda';
 
-import { LambdaModule, AwsLambdaInstrumentationConfig } from './types';
+import {
+  LambdaModule,
+  AwsLambdaInstrumentationConfig,
+  EventContextExtractor,
+} from './types';
 import { VERSION } from './version';
 
 const awsPropagator = new AWSXRayPropagator();
@@ -149,11 +153,12 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
       context: Context,
       callback: Callback
     ) {
-      const httpHeaders =
-        typeof event.headers === 'object' ? event.headers : {};
+      const config = plugin._config;
       const parent = AwsLambdaInstrumentation._determineParent(
-        httpHeaders,
-        plugin._config.disableAwsContextPropagation === true
+        event,
+        config.disableAwsContextPropagation === true,
+        config.eventContextExtractor ||
+          AwsLambdaInstrumentation._defaultEventContextExtractor
       );
 
       const name = context.functionName;
@@ -173,9 +178,9 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
         parent
       );
 
-      if (plugin._config.requestHook) {
+      if (config.requestHook) {
         safeExecuteInTheMiddle(
-          () => plugin._config.requestHook!(span, { event, context }),
+          () => config.requestHook!(span, { event, context }),
           e => {
             if (e)
               diag.error('aws-lambda instrumentation: requestHook error', e);
@@ -300,12 +305,18 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
     return undefined;
   }
 
+  private static _defaultEventContextExtractor(event: any): OtelContext {
+    // The default extractor tries to get sampled trace header from HTTP headers.
+    const httpHeaders = event.headers || {};
+    return propagation.extract(otelContext.active(), httpHeaders, headerGetter);
+  }
+
   private static _determineParent(
-    httpHeaders: APIGatewayProxyEventHeaders,
-    disableAwsContextPropagation: boolean
+    event: any,
+    disableAwsContextPropagation: boolean,
+    eventContextExtractor: EventContextExtractor
   ): OtelContext {
     let parent: OtelContext | undefined = undefined;
-
     if (!disableAwsContextPropagation) {
       const lambdaTraceHeader = process.env[traceContextEnvironmentKey];
       if (lambdaTraceHeader) {
@@ -328,15 +339,19 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
         }
       }
     }
-
-    // There was not a sampled trace header from Lambda so try from HTTP headers.
-    const httpContext = propagation.extract(
-      otelContext.active(),
-      httpHeaders,
-      headerGetter
+    const extractedContext = safeExecuteInTheMiddle(
+      () => eventContextExtractor(event),
+      e => {
+        if (e)
+          diag.error(
+            'aws-lambda instrumentation: eventContextExtractor error',
+            e
+          );
+      },
+      true
     );
-    if (trace.getSpan(httpContext)?.spanContext()) {
-      return httpContext;
+    if (trace.getSpan(extractedContext)?.spanContext()) {
+      return extractedContext;
     }
     if (!parent) {
       // No context in Lambda environment or HTTP headers.

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

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { Span } from '@opentelemetry/api';
+import { Span, Context as OtelContext } from '@opentelemetry/api';
 import { InstrumentationConfig } from '@opentelemetry/instrumentation';
 import { Handler, Context } from 'aws-lambda';
 
@@ -33,8 +33,10 @@ export type ResponseHook = (
   }
 ) => void;
 
+export type EventContextExtractor = (event: any) => OtelContext;
 export interface AwsLambdaInstrumentationConfig extends InstrumentationConfig {
   requestHook?: RequestHook;
   responseHook?: ResponseHook;
   disableAwsContextPropagation?: boolean;
+  eventContextExtractor?: EventContextExtractor;
 }

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

@@ -37,7 +37,9 @@ import {
   ResourceAttributes,
 } from '@opentelemetry/semantic-conventions';
 import {
+  Context as OtelContext,
   context,
+  propagation,
   trace,
   SpanContext,
   SpanKind,
@@ -166,6 +168,17 @@ describe('lambda handler', () => {
     new HttpTraceContextPropagator()
   );
 
+  const sampledGenericSpanContext: SpanContext = {
+    traceId: '8a3c60f7d188f8fa79d48a391a778faa',
+    spanId: '0000000000000460',
+    traceFlags: 1,
+    isRemote: true,
+  };
+  const sampledGenericSpan = serializeSpanContext(
+    sampledGenericSpanContext,
+    new HttpTraceContextPropagator()
+  );
+
   beforeEach(() => {
     oldEnv = { ...process.env };
     process.env.LAMBDA_TASK_ROOT = path.resolve(__dirname, '..');
@@ -577,6 +590,40 @@ describe('lambda handler', () => {
       );
       assert.strictEqual(span.parentSpanId, sampledHttpSpanContext.spanId);
     });
+
+    it('takes sampled custom context over sampled lambda context if "eventContextExtractor" is defined', async () => {
+      process.env[traceContextEnvironmentKey] = sampledAwsHeader;
+      const customExtractor = (event: any): OtelContext => {
+        return propagation.extract(context.active(), event.contextCarrier);
+      };
+
+      initializeHandler('lambda-test/async.handler', {
+        disableAwsContextPropagation: true,
+        eventContextExtractor: customExtractor,
+      });
+
+      const otherEvent = {
+        contextCarrier: {
+          traceparent: sampledGenericSpan,
+        },
+      };
+
+      const result = await lambdaRequire('lambda-test/async').handler(
+        otherEvent,
+        ctx
+      );
+
+      assert.strictEqual(result, 'ok');
+      const spans = memoryExporter.getFinishedSpans();
+      const [span] = spans;
+      assert.strictEqual(spans.length, 1);
+      assertSpanSuccess(span);
+      assert.strictEqual(
+        span.spanContext().traceId,
+        sampledGenericSpanContext.traceId
+      );
+      assert.strictEqual(span.parentSpanId, sampledGenericSpanContext.spanId);
+    });
   });
 
   describe('hooks', () => {