Skip to content

Commit

Permalink
fix(aws-serverless): Extract sentry trace data from handler context
Browse files Browse the repository at this point in the history
… 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
  • Loading branch information
andreiborza authored Aug 8, 2024
1 parent 6cbc416 commit b17ac59
Show file tree
Hide file tree
Showing 5 changed files with 217 additions and 21 deletions.
4 changes: 4 additions & 0 deletions packages/aws-serverless/rollup.npm.config.mjs
Original file line number Diff line number Diff line change
Expand Up @@ -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'),
Expand Down
44 changes: 34 additions & 10 deletions packages/aws-serverless/src/integration/awslambda.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand Down
14 changes: 4 additions & 10 deletions packages/aws-serverless/src/sdk.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand All @@ -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;

Expand Down Expand Up @@ -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,
Expand Down
74 changes: 73 additions & 1 deletion packages/aws-serverless/src/utils.ts
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -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);
}
102 changes: 102 additions & 0 deletions packages/aws-serverless/test/utils.test.ts
Original file line number Diff line number Diff line change
@@ -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]));
});
});

0 comments on commit b17ac59

Please sign in to comment.