fix(aws-serverless): Extract sentry trace data from handler context over event (#13266)

Currently, the AWS otel integration (and our `wrapHandler` fallback) try
to extract sentry trace data from the `event` object passed to a Lambda
call. The aws-sdk integration, however, places tracing data onto
`context.clientContext.Custom`.

This PR adds a custom `eventContextExtractor` that attempts extracting
sentry trace data from the `context`, with a fallback to `event` to
enable distributed tracing among Lambda invocations.

Traces are now connected. Here an example:

`Lambda-A` calling `Lambda-B`:
```
import { LambdaClient, InvokeCommand } from "@aws-sdk/client-lambda";
import * as Sentry from "@sentry/aws-serverless";

export const handler = Sentry.wrapHandler(async (event, context) => {
  const client = new LambdaClient();
  const command = new InvokeCommand({
    FunctionName: `Lambda-B`,
    InvocationType: "RequestResponse",
    Payload: new Uint16Array(),
  })
  return client.send(command);

});
```

`Lambda-B`:
```
import * as Sentry from "@sentry/aws-serverless";
Sentry.addIntegration(Sentry.postgresIntegration())

export const handler = Sentry.wrapHandler(async (event) => {
  const queryString = "select count(*) from myTable;";
  return await Sentry.startSpan({
    name: queryString,
    op: "db.sql.execute"
  }, async (span) => {
    console.log('executing query', queryString);
  })
})
```

![CleanShot 2024-08-07 at 16 34
51@2x](https://github.com/user-attachments/assets/43f5dd9e-e5af-4667-9551-05fac90f03a6)

Closes: #13146

Author: andreiborza

packages/aws-serverless/rollup.npm.config.mjs

@@ -9,6 +9,10 @@ export default [
       entrypoints: ['src/index.ts', 'src/awslambda-auto.ts'],
       // packages with bundles have a different build directory structure
       hasBundles: true,
+      packageSpecificConfig: {
+        // Used for our custom eventContextExtractor
+        external: ['@opentelemetry/api'],
+      },
     }),
   ),
   ...makeOtelLoaders('./build', 'sentry-node'),

packages/aws-serverless/src/integration/awslambda.ts

@@ -1,20 +1,44 @@
 import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lambda';
 import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, defineIntegration } from '@sentry/core';
-import { addOpenTelemetryInstrumentation } from '@sentry/node';
+import { generateInstrumentOnce } from '@sentry/node';
 import type { IntegrationFn } from '@sentry/types';
+import { eventContextExtractor } from '../utils';
 
-const _awsLambdaIntegration = (() => {
+interface AwsLambdaOptions {
+  /**
+   * Disables the AWS context propagation and instead uses
+   * Sentry's context. Defaults to `true`, in order for
+   * Sentry trace propagation to take precedence, but can
+   * be disabled if you want AWS propagation to take take
+   * precedence.
+   */
+  disableAwsContextPropagation?: boolean;
+}
+
+export const instrumentAwsLambda = generateInstrumentOnce<AwsLambdaOptions>(
+  'AwsLambda',
+  (_options: AwsLambdaOptions = {}) => {
+    const options = {
+      disableAwsContextPropagation: true,
+      ..._options,
+    };
+
+    return new AwsLambdaInstrumentation({
+      ...options,
+      eventContextExtractor,
+      requestHook(span) {
+        span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.otel.aws-lambda');
+        span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'function.aws.lambda');
+      },
+    });
+  },
+);
+
+const _awsLambdaIntegration = ((options: AwsLambdaOptions = {}) => {
   return {
     name: 'AwsLambda',
     setupOnce() {
-      addOpenTelemetryInstrumentation(
-        new AwsLambdaInstrumentation({
-          requestHook(span) {
-            span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, 'auto.otel.aws-lambda');
-            span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'function.aws.lambda');
-          },
-        }),
-      );
+      instrumentAwsLambda(options);
     },
   };
 }) satisfies IntegrationFn;

packages/aws-serverless/src/sdk.ts

@@ -16,7 +16,7 @@ import {
   withScope,
 } from '@sentry/node';
 import type { Integration, Options, Scope, SdkMetadata, Span } from '@sentry/types';
-import { isString, logger } from '@sentry/utils';
+import { logger } from '@sentry/utils';
 import type { Context, Handler } from 'aws-lambda';
 import { performance } from 'perf_hooks';
 
@@ -25,7 +25,7 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SEMANTIC_ATTRIBUTE_SENTRY_SOURCE } fr
 import { DEBUG_BUILD } from './debug-build';
 import { awsIntegration } from './integration/aws';
 import { awsLambdaIntegration } from './integration/awslambda';
-import { markEventUnhandled } from './utils';
+import { getAwsTraceData, markEventUnhandled } from './utils';
 
 const { isPromise } = types;
 
@@ -334,15 +334,9 @@ export function wrapHandler<TEvent, TResult>(
     // Otherwise, we create two root spans (one from otel, one from our wrapper).
     // If Otel instrumentation didn't work or was filtered by users, we still want to trace the handler.
     if (options.startTrace && !isWrappedByOtel(handler)) {
-      const eventWithHeaders = event as { headers?: { [key: string]: string } };
+      const traceData = getAwsTraceData(event as { headers?: Record<string, string> }, context);
 
-      const sentryTrace =
-        eventWithHeaders.headers && isString(eventWithHeaders.headers['sentry-trace'])
-          ? eventWithHeaders.headers['sentry-trace']
-          : undefined;
-      const baggage = eventWithHeaders.headers?.baggage;
-
-      return continueTrace({ sentryTrace, baggage }, () => {
+      return continueTrace({ sentryTrace: traceData['sentry-trace'], baggage: traceData.baggage }, () => {
         return startSpanManual(
           {
             name: context.functionName,

packages/aws-serverless/src/utils.ts

@@ -1,5 +1,29 @@
+import type { TextMapGetter } from '@opentelemetry/api';
+import type { Context as OtelContext } from '@opentelemetry/api';
+import { context as otelContext, propagation } from '@opentelemetry/api';
 import type { Scope } from '@sentry/types';
-import { addExceptionMechanism } from '@sentry/utils';
+import { addExceptionMechanism, isString } from '@sentry/utils';
+import type { Handler } from 'aws-lambda';
+import type { APIGatewayProxyEventHeaders } from 'aws-lambda';
+
+type HandlerEvent = Parameters<Handler<{ headers?: Record<string, string> }>>[0];
+type HandlerContext = Parameters<Handler>[1];
+
+type TraceData = {
+  'sentry-trace'?: string;
+  baggage?: string;
+};
+
+// vendored from
+// https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts#L65-L72
+const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
+  keys(carrier): string[] {
+    return Object.keys(carrier);
+  },
+  get(carrier, key: string) {
+    return carrier[key];
+  },
+};
 
 /**
  * Marks an event as unhandled by adding a span processor to the passed scope.
@@ -12,3 +36,51 @@ export function markEventUnhandled(scope: Scope): Scope {
 
   return scope;
 }
+
+/**
+ * Extracts sentry trace data from the handler `context` if available and falls
+ * back to the `event`.
+ *
+ * When instrumenting the Lambda function with Sentry, the sentry trace data
+ * is placed on `context.clientContext.Custom`. Users are free to modify context
+ * tho and provide this data via `event` or `context`.
+ */
+export function getAwsTraceData(event: HandlerEvent, context?: HandlerContext): TraceData {
+  const headers = event.headers || {};
+
+  const traceData: TraceData = {
+    'sentry-trace': headers['sentry-trace'],
+    baggage: headers.baggage,
+  };
+
+  if (context && context.clientContext && context.clientContext.Custom) {
+    const customContext: Record<string, unknown> = context.clientContext.Custom;
+    const sentryTrace = isString(customContext['sentry-trace']) ? customContext['sentry-trace'] : undefined;
+
+    if (sentryTrace) {
+      traceData['sentry-trace'] = sentryTrace;
+      traceData.baggage = isString(customContext.baggage) ? customContext.baggage : undefined;
+    }
+  }
+
+  return traceData;
+}
+
+/**
+ * A custom event context extractor for the aws integration. It takes sentry trace data
+ * from the context rather than the event, with the event being a fallback.
+ *
+ * Is only used when the handler was successfully wrapped by otel and the integration option
+ * `disableAwsContextPropagation` is `true`.
+ */
+export function eventContextExtractor(event: HandlerEvent, context?: HandlerContext): OtelContext {
+  // The default context extractor tries to get sampled trace headers from HTTP headers
+  // The otel aws integration packs these onto the context, so we try to extract them from
+  // there instead.
+  const httpHeaders = {
+    ...(event.headers || {}),
+    ...getAwsTraceData(event, context),
+  };
+
+  return propagation.extract(otelContext.active(), httpHeaders, headerGetter);
+}

packages/aws-serverless/test/utils.test.ts

@@ -0,0 +1,102 @@
+import { eventContextExtractor, getAwsTraceData } from '../src/utils';
+
+const mockExtractContext = jest.fn();
+jest.mock('@opentelemetry/api', () => {
+  const actualApi = jest.requireActual('@opentelemetry/api');
+  return {
+    ...actualApi,
+    propagation: {
+      extract: (...args: unknown[]) => mockExtractContext(args),
+    },
+  };
+});
+
+const mockContext = {
+  clientContext: {
+    Custom: {
+      'sentry-trace': '12345678901234567890123456789012-1234567890123456-1',
+      baggage: 'sentry-environment=production',
+    },
+  },
+};
+const mockEvent = {
+  headers: {
+    'sentry-trace': '12345678901234567890123456789012-1234567890123456-2',
+    baggage: 'sentry-environment=staging',
+  },
+};
+
+describe('getTraceData', () => {
+  test('gets sentry trace data from the context', () => {
+    // @ts-expect-error, a partial context object is fine here
+    const traceData = getAwsTraceData({}, mockContext);
+
+    expect(traceData['sentry-trace']).toEqual('12345678901234567890123456789012-1234567890123456-1');
+    expect(traceData.baggage).toEqual('sentry-environment=production');
+  });
+
+  test('gets sentry trace data from the context even if event has data', () => {
+    // @ts-expect-error, a partial context object is fine here
+    const traceData = getAwsTraceData(mockEvent, mockContext);
+
+    expect(traceData['sentry-trace']).toEqual('12345678901234567890123456789012-1234567890123456-1');
+    expect(traceData.baggage).toEqual('sentry-environment=production');
+  });
+
+  test('gets sentry trace data from the event if no context is passed', () => {
+    const traceData = getAwsTraceData(mockEvent);
+
+    expect(traceData['sentry-trace']).toEqual('12345678901234567890123456789012-1234567890123456-2');
+    expect(traceData.baggage).toEqual('sentry-environment=staging');
+  });
+
+  test('gets sentry trace data from the event if the context sentry trace is undefined', () => {
+    const traceData = getAwsTraceData(mockEvent, {
+      // @ts-expect-error, a partial context object is fine here
+      clientContext: { Custom: { 'sentry-trace': undefined, baggage: '' } },
+    });
+
+    expect(traceData['sentry-trace']).toEqual('12345678901234567890123456789012-1234567890123456-2');
+    expect(traceData.baggage).toEqual('sentry-environment=staging');
+  });
+});
+
+describe('eventContextExtractor', () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  test('passes sentry trace data to the propagation extractor', () => {
+    // @ts-expect-error, a partial context object is fine here
+    eventContextExtractor(mockEvent, mockContext);
+
+    // @ts-expect-error, a partial context object is fine here
+    const expectedTraceData = getAwsTraceData(mockEvent, mockContext);
+
+    expect(mockExtractContext).toHaveBeenCalledTimes(1);
+    expect(mockExtractContext).toHaveBeenCalledWith(expect.arrayContaining([expectedTraceData]));
+  });
+
+  test('passes along non-sentry trace headers along', () => {
+    eventContextExtractor(
+      {
+        ...mockEvent,
+        headers: {
+          ...mockEvent.headers,
+          'X-Custom-Header': 'Foo',
+        },
+      },
+      // @ts-expect-error, a partial context object is fine here
+      mockContext,
+    );
+
+    const expectedHeaders = {
+      'X-Custom-Header': 'Foo',
+      // @ts-expect-error, a partial context object is fine here
+      ...getAwsTraceData(mockEvent, mockContext),
+    };
+
+    expect(mockExtractContext).toHaveBeenCalledTimes(1);
+    expect(mockExtractContext).toHaveBeenCalledWith(expect.arrayContaining([expectedHeaders]));
+  });
+});