feat(aws-lambda): disableAwsContextPropagation config option (#546)

Co-authored-by: Daniel Dyla <[email protected]>

Author: Nir Hadassi

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

@@ -46,6 +46,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 |
 
 ### Hooks Usage Example
 

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

@@ -151,7 +151,10 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
     ) {
       const httpHeaders =
         typeof event.headers === 'object' ? event.headers : {};
-      const parent = AwsLambdaInstrumentation._determineParent(httpHeaders);
+      const parent = AwsLambdaInstrumentation._determineParent(
+        httpHeaders,
+        plugin._config.disableAwsContextPropagation === true
+      );
 
       const name = context.functionName;
       const span = plugin.tracer.startSpan(
@@ -298,29 +301,34 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
   }
 
   private static _determineParent(
-    httpHeaders: APIGatewayProxyEventHeaders
+    httpHeaders: APIGatewayProxyEventHeaders,
+    disableAwsContextPropagation: boolean
   ): OtelContext {
     let parent: OtelContext | undefined = undefined;
-    const lambdaTraceHeader = process.env[traceContextEnvironmentKey];
-    if (lambdaTraceHeader) {
-      parent = awsPropagator.extract(
-        otelContext.active(),
-        { [AWSXRAY_TRACE_ID_HEADER]: lambdaTraceHeader },
-        headerGetter
-      );
-    }
-    if (parent) {
-      const spanContext = trace.getSpan(parent)?.spanContext();
-      if (
-        spanContext &&
-        (spanContext.traceFlags & TraceFlags.SAMPLED) === TraceFlags.SAMPLED
-      ) {
-        // Trace header provided by Lambda only sampled if a sampled context was propagated from
-        // an upstream cloud service such as S3, or the user is using X-Ray. In these cases, we
-        // need to use it as the parent.
-        return parent;
+
+    if (!disableAwsContextPropagation) {
+      const lambdaTraceHeader = process.env[traceContextEnvironmentKey];
+      if (lambdaTraceHeader) {
+        parent = awsPropagator.extract(
+          otelContext.active(),
+          { [AWSXRAY_TRACE_ID_HEADER]: lambdaTraceHeader },
+          headerGetter
+        );
+      }
+      if (parent) {
+        const spanContext = trace.getSpan(parent)?.spanContext();
+        if (
+          spanContext &&
+          (spanContext.traceFlags & TraceFlags.SAMPLED) === TraceFlags.SAMPLED
+        ) {
+          // Trace header provided by Lambda only sampled if a sampled context was propagated from
+          // an upstream cloud service such as S3, or the user is using X-Ray. In these cases, we
+          // need to use it as the parent.
+          return parent;
+        }
       }
     }
+
     // There was not a sampled trace header from Lambda so try from HTTP headers.
     const httpContext = propagation.extract(
       otelContext.active(),

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

@@ -36,4 +36,5 @@ export type ResponseHook = (
 export interface AwsLambdaInstrumentationConfig extends InstrumentationConfig {
   requestHook?: RequestHook;
   responseHook?: ResponseHook;
+  disableAwsContextPropagation?: boolean;
 }

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

@@ -526,6 +526,57 @@ describe('lambda handler', () => {
       // Parent unsampled so no spans exported.
       assert.strictEqual(spans.length, 0);
     });
+
+    it('ignores sampled lambda context if "disableAwsContextPropagation" config option is true', async () => {
+      process.env[traceContextEnvironmentKey] = sampledAwsHeader;
+      initializeHandler('lambda-test/async.handler', {
+        disableAwsContextPropagation: true,
+      });
+
+      const result = await lambdaRequire('lambda-test/async').handler(
+        'arg',
+        ctx
+      );
+      assert.strictEqual(result, 'ok');
+      const spans = memoryExporter.getFinishedSpans();
+      const [span] = spans;
+      assert.strictEqual(spans.length, 1);
+      assertSpanSuccess(span);
+      assert.notDeepStrictEqual(
+        span.spanContext().traceId,
+        sampledAwsSpanContext.traceId
+      );
+      assert.strictEqual(span.parentSpanId, undefined);
+    });
+
+    it('takes sampled http context over sampled lambda context if "disableAwsContextPropagation" config option is true', async () => {
+      process.env[traceContextEnvironmentKey] = sampledAwsHeader;
+      initializeHandler('lambda-test/async.handler', {
+        disableAwsContextPropagation: true,
+      });
+
+      const proxyEvent = {
+        headers: {
+          traceparent: sampledHttpHeader,
+        },
+      };
+
+      const result = await lambdaRequire('lambda-test/async').handler(
+        proxyEvent,
+        ctx
+      );
+
+      assert.strictEqual(result, 'ok');
+      const spans = memoryExporter.getFinishedSpans();
+      const [span] = spans;
+      assert.strictEqual(spans.length, 1);
+      assertSpanSuccess(span);
+      assert.strictEqual(
+        span.spanContext().traceId,
+        sampledHttpSpanContext.traceId
+      );
+      assert.strictEqual(span.parentSpanId, sampledHttpSpanContext.spanId);
+    });
   });
 
   describe('hooks', () => {