diff --git a/dev-packages/e2e-tests/test-applications/aws-serverless-layer/src/lambda-functions-layer/Callback/index.js b/dev-packages/e2e-tests/test-applications/aws-serverless-layer/src/lambda-functions-layer/Callback/index.js index c8d4342954e9..42f56e655646 100644 --- a/dev-packages/e2e-tests/test-applications/aws-serverless-layer/src/lambda-functions-layer/Callback/index.js +++ b/dev-packages/e2e-tests/test-applications/aws-serverless-layer/src/lambda-functions-layer/Callback/index.js @@ -1,9 +1,19 @@ // Callback-style handler `(event, context, callback)` to exercise the wrapped-callback span path. +// `event.callbackDelayMs` makes the handler return before it calls back, like most real callback-style handlers. exports.handler = (event, context, callback) => { - if (event.shouldError) { - callback(new Error('callback error')); + const respond = () => { + if (event.shouldError) { + callback(new Error('callback error')); + return; + } + + callback(null, { ok: true }); + }; + + if (event.callbackDelayMs) { + setTimeout(respond, event.callbackDelayMs); return; } - callback(null, { ok: true }); + respond(); }; diff --git a/dev-packages/e2e-tests/test-applications/aws-serverless-layer/tests/layer.test.ts b/dev-packages/e2e-tests/test-applications/aws-serverless-layer/tests/layer.test.ts index 862c956d7883..538b533967e0 100644 --- a/dev-packages/e2e-tests/test-applications/aws-serverless-layer/tests/layer.test.ts +++ b/dev-packages/e2e-tests/test-applications/aws-serverless-layer/tests/layer.test.ts @@ -322,6 +322,48 @@ test.describe('Lambda layer', () => { ); }); + test('ends the span when a callback-style handler calls back after it returned', async ({ lambdaClient }) => { + const transactionEventPromise = waitForTransaction('aws-serverless-layer', transactionEvent => { + return transactionEvent?.transaction === 'LayerCallback'; + }); + + await lambdaClient.send( + new InvokeCommand({ + FunctionName: 'LayerCallback', + Payload: JSON.stringify({ callbackDelayMs: 200 }), + }), + ); + + const transactionEvent = await transactionEventPromise; + + expect(transactionEvent.contexts?.trace?.status).toBe('ok'); + expect(transactionEvent.timestamp! - transactionEvent.start_timestamp!).toBeGreaterThanOrEqual(0.2); + }); + + test('sets error status when a callback-style handler calls back with an error after it returned', async ({ + lambdaClient, + }) => { + const errorEventPromise = waitForError('aws-serverless-layer', errorEvent => { + return errorEvent?.exception?.values?.[0]?.value === 'callback error'; + }); + const transactionEventPromise = waitForTransaction('aws-serverless-layer', transactionEvent => { + return transactionEvent?.transaction === 'LayerCallback'; + }); + + await lambdaClient.send( + new InvokeCommand({ + FunctionName: 'LayerCallback', + Payload: JSON.stringify({ shouldError: true, callbackDelayMs: 200 }), + }), + ); + + const [errorEvent, transactionEvent] = await Promise.all([errorEventPromise, transactionEventPromise]); + + expect(transactionEvent.contexts?.trace?.status).toBe('internal_error'); + expect(transactionEvent.timestamp! - transactionEvent.start_timestamp!).toBeGreaterThanOrEqual(0.2); + expect(errorEvent.contexts?.trace?.trace_id).toBe(transactionEvent.contexts?.trace?.trace_id); + }); + test('continues a trace from incoming headers', async ({ lambdaClient }) => { const traceId = 'd4cda95b652f4a1592b449d5929fda1b'; const parentSpanId = '6e0c63257de34c92'; diff --git a/packages/aws-serverless/src/instrumentHandler.ts b/packages/aws-serverless/src/instrumentHandler.ts index b652def8b0a9..c1b4fb2bf5b2 100644 --- a/packages/aws-serverless/src/instrumentHandler.ts +++ b/packages/aws-serverless/src/instrumentHandler.ts @@ -134,6 +134,11 @@ function getPatchedHandler(original: Handler | StreamifyHandler, lambdaStartTime throw error; } + // Callback-style handlers usually return before invoking the callback, so the wrapped callback ends the span. + if (original.length > 2 && typeof maybePromise?.then !== 'function') { + return maybePromise; + } + return handlePromiseResult(span, maybePromise); }), ); diff --git a/packages/aws-serverless/test/instrumentHandler.test.ts b/packages/aws-serverless/test/instrumentHandler.test.ts new file mode 100644 index 000000000000..5bea12c73bb4 --- /dev/null +++ b/packages/aws-serverless/test/instrumentHandler.test.ts @@ -0,0 +1,84 @@ +import type { Context, Handler } from 'aws-lambda'; +import { beforeEach, describe, expect, test, vi } from 'vitest'; +import { instrumentHandler } from '../src/instrumentHandler'; + +const mockSpan = { + setStatus: vi.fn(), + end: vi.fn(), +}; + +vi.mock('@sentry/core', async () => { + const original = await vi.importActual('@sentry/core'); + return { + ...original, + continueTrace: (_traceData: unknown, callback: () => unknown) => callback(), + startSpanManual: (_options: unknown, callback: (span: unknown) => unknown) => callback(mockSpan), + }; +}); + +vi.mock('@sentry/node', async () => { + const original = await vi.importActual('@sentry/node'); + return { + ...original, + flush: () => Promise.resolve(true), + captureException: vi.fn(), + }; +}); + +const fakeContext = { + callbackWaitsForEmptyEventLoop: false, + functionName: 'functionName', + functionVersion: 'functionVersion', + invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:functionName', + memoryLimitInMB: '128', + awsRequestId: 'awsRequestId', + logGroupName: 'logGroupName', + logStreamName: 'logStreamName', + getRemainingTimeInMillis: () => 30_000, + done: () => {}, + fail: () => {}, + succeed: () => {}, +} as Context; + +describe('instrumentHandler', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + test('ends the span when a callback-style handler calls back asynchronously', async () => { + const handler: Handler = (_event, _context, callback) => { + setTimeout(() => callback(null, 42), 10); + }; + + const result = instrumentHandler(handler)({}, fakeContext, () => {}); + expect(mockSpan.end).not.toHaveBeenCalled(); + + await expect(result).resolves.toBe(42); + expect(mockSpan.end).toHaveBeenCalledTimes(1); + expect(mockSpan.setStatus).not.toHaveBeenCalled(); + }); + + test('sets error status when a callback-style handler calls back with an error asynchronously', async () => { + const handler: Handler = (_event, _context, callback) => { + setTimeout(() => callback(new Error('boom')), 10); + }; + + const result = instrumentHandler(handler)({}, fakeContext, () => {}); + expect(mockSpan.end).not.toHaveBeenCalled(); + + await expect(result).rejects.toThrow('boom'); + expect(mockSpan.setStatus).toHaveBeenCalledWith({ code: 2, message: 'boom' }); + expect(mockSpan.end).toHaveBeenCalledTimes(1); + }); + + test('ends the span when an async handler resolves', async () => { + const handler: Handler = async () => { + await new Promise(resolve => setTimeout(resolve, 10)); + return 42; + }; + + await expect(instrumentHandler(handler)({}, fakeContext, () => {})).resolves.toBe(42); + expect(mockSpan.end).toHaveBeenCalledTimes(1); + expect(mockSpan.setStatus).not.toHaveBeenCalled(); + }); +});