feat: support using lambda context in the aws lambda context extractor (

open-telemetry#860)
dynatrace-oss-contrib · Jan 31, 2022 · 5cb3266 · 5cb3266
1 parent 0d90e9d
commit 5cb3266
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 3 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/README.md b/plugins/node/opentelemetry-instrumentation-aws-lambda/README.md
@@ -51,7 +51,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` |
+| `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, context` |
 
 ### Hooks Usage Example
 

diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
@@ -156,6 +156,7 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
       const config = plugin._config;
       const parent = AwsLambdaInstrumentation._determineParent(
         event,
+        context,
         config.disableAwsContextPropagation === true,
         config.eventContextExtractor ||
           AwsLambdaInstrumentation._defaultEventContextExtractor
@@ -331,6 +332,7 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
 
   private static _determineParent(
     event: any,
+    context: Context,
     disableAwsContextPropagation: boolean,
     eventContextExtractor: EventContextExtractor
   ): OtelContext {
@@ -358,7 +360,7 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
       }
     }
     const extractedContext = safeExecuteInTheMiddle(
-      () => eventContextExtractor(event),
+      () => eventContextExtractor(event, context),
       e => {
         if (e)
           diag.error(

diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/types.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/types.ts
@@ -33,7 +33,10 @@ export type ResponseHook = (
   }
 ) => void;
 
-export type EventContextExtractor = (event: any) => OtelContext;
+export type EventContextExtractor = (
+  event: any,
+  context: Context
+) => OtelContext;
 export interface AwsLambdaInstrumentationConfig extends InstrumentationConfig {
   requestHook?: RequestHook;
   responseHook?: ResponseHook;

diff --git a/...ns/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts b/...ns/node/opentelemetry-instrumentation-aws-lambda/test/integrations/lambda-handler.test.ts
@@ -656,6 +656,64 @@ describe('lambda handler', () => {
       const [span] = spans;
       assert.strictEqual(span.parentSpanId, undefined);
     });
+
+    it('passes the lambda context object into the eventContextExtractor for scenarios where it is the otel context carrier', async () => {
+      process.env[traceContextEnvironmentKey] = sampledAwsHeader;
+      const customExtractor = (
+        event: any,
+        handlerContext: Context
+      ): OtelContext => {
+        const contextCarrier = handlerContext.clientContext?.Custom ?? {};
+        return propagation.extract(context.active(), contextCarrier);
+      };
+
+      initializeHandler('lambda-test/async.handler', {
+        disableAwsContextPropagation: true,
+        eventContextExtractor: customExtractor,
+      });
+
+      const otherEvent = {};
+      const ctxWithCustomData = {
+        functionName: 'my_function',
+        invokedFunctionArn: 'my_arn',
+        awsRequestId: 'aws_request_id',
+        clientContext: {
+          client: {
+            installationId: '',
+            appTitle: '',
+            appVersionName: '',
+            appVersionCode: '',
+            appPackageName: '',
+          },
+          Custom: {
+            traceparent: sampledGenericSpan,
+          },
+          env: {
+            platformVersion: '',
+            platform: '',
+            make: '',
+            model: '',
+            locale: '',
+          },
+        },
+      } as Context;
+
+      const result = await lambdaRequire('lambda-test/async').handler(
+        otherEvent,
+        ctxWithCustomData
+      );
+
+      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', () => {