fix(aws-serverless): End the function.aws span when a callback-style handler calls back

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();
 };

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';

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);
       }),
     );

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();
+  });
+});