feat(lambda): add explicit callback completion mode

README.md

@@ -31,7 +31,30 @@ Besides the environment variables supported by dd-trace-js, the datadog-lambda-j
 | DD_CAPTURE_LAMBDA_PAYLOAD | [Captures incoming and outgoing AWS Lambda payloads][1] in the Datadog APM spans for Lambda invocations. | `false` |
 | DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH | Determines the level of detail captured from AWS Lambda payloads, which are then assigned as tags for the `aws.lambda` span. It specifies the nesting depth of the JSON payload structure to process. Once the specified maximum depth is reached, the tag's value is set to the stringified value of any nested elements beyond this level.  <br> For example, given the input payload: <pre>{<br>  "lv1" : {<br>    "lv2": {<br>      "lv3": "val"<br>    }<br>  }<br>}</pre> If the depth is set to `2`, the resulting tag's key is set to `function.request.lv1.lv2` and the value is `{\"lv3\": \"val\"}`. <br> If the depth is set to `0`, the resulting tag's key is set to `function.request` and value is `{\"lv1\":{\"lv2\":{\"lv3\": \"val\"}}}` | `10` |
 | DD_DURABLE_CROSS_INVOCATION_TRACING_ENABLED | For AWS Durable functions, the tracer creates extra checkpoints named `_datadog_{N}` to propagate trace context across function invocations, keeping spans from multiple invocations in one intact trace for each durable execution. | `true` |
-
+| DD_LAMBDA_HANDLER_COMPLETION | Set to `callback` for buffered handlers that complete through a callback or `context.done/succeed/fail` and may return an incidental synchronous value. `auto` preserves existing completion selection. | `auto` |
+
+### Callback completion with rest or default parameters
+
+For a buffered handler that completes through its callback, set `DD_LAMBDA_HANDLER_COMPLETION=callback`
+to ignore ordinary synchronous return values. This supports rest/default-parameter handlers whose
+`length` is less than three, without changing the handler or adding another wrapper:
+
+```js
+exports.handler = function (...[event, context, callback]) {
+  startWork(event, (error, response) => callback(error, response));
+  return { incidental: true };
+};
+```
+
+With `callback`, the library waits for the callback or a legacy context completion method. If the handler
+returns a Promise, that Promise and the callback still race; whichever completes first wins. A handler
+that never signals completion remains pending until Lambda times out. Response-streaming handlers are
+unaffected.
+
+The default `auto` behavior uses declared arity, the return value, and the existing server/emitter
+heuristics to select completion. In particular, a rest/default-parameter handler returning a plain object
+can complete with that object before its asynchronous callback. Handlers that intentionally return
+synchronous responses should keep `auto`. Unknown values also fall back to `auto`.
 
 ## Tracer compatibility
 

src/index.spec.ts

@@ -1,7 +1,7 @@
 import http from "http";
 import nock from "nock";
 
-import { Context, Handler } from "aws-lambda";
+import { Callback, Context, Handler } from "aws-lambda";
 import {
   datadog,
   getTraceHeaders,
@@ -115,6 +115,51 @@ describe("datadog", () => {
     process.env = oldEnv;
   });
 
+  describe("explicit callback completion", () => {
+    it.each(["rest", "default"])("uses the environment setting for %s parameters", async (kind) => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      let complete: Callback<unknown> = () => {
+        throw new Error("handler has not started");
+      };
+      let entered: () => void = () => {};
+      const started = new Promise<void>((resolve) => {
+        entered = resolve;
+      });
+      const work = (callback: Callback<unknown>) => {
+        complete = callback;
+        entered();
+        return { incidental: true };
+      };
+      const candidate =
+        kind === "rest"
+          ? (...args: Parameters<Handler>) => work(args[2])
+          : (event: unknown, context: Context, callback: Callback<unknown> = () => {}) => work(callback);
+      const wrapped = datadog(candidate, { enhancedMetrics: false, autoPatchHTTP: false });
+      const result = wrapped({}, { ...mockContext });
+      await started;
+      const response = { statusCode: 200, body: "callback response" };
+      complete(null, response);
+      await expect(result).resolves.toBe(response);
+    });
+
+    it("restores automatic synchronous completion after the environment setting is removed", async () => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      const callbackResponse = { callback: true };
+      const callbackHandler = datadog(
+        (_event: unknown, context: Context) => {
+          context.succeed(callbackResponse);
+          return { incidental: true };
+        },
+        { enhancedMetrics: false, autoPatchHTTP: false },
+      );
+      await expect(callbackHandler({}, { ...mockContext })).resolves.toBe(callbackResponse);
+      delete process.env.DD_LAMBDA_HANDLER_COMPLETION;
+      const response = { synchronous: true };
+      const synchronousHandler = datadog(() => response, { enhancedMetrics: false, autoPatchHTTP: false });
+      await expect(synchronousHandler({}, { ...mockContext })).resolves.toBe(response);
+    });
+  });
+
   it("patches http request when autoPatch enabled", async () => {
     mockTraceHeaders = {
       "x-datadog-parent-id": "9101112",

src/utils/handler.spec.ts

@@ -1,8 +1,9 @@
-import { Context, Handler } from "aws-lambda";
+import { Callback, Context, Handler } from "aws-lambda";
 
 import { didFunctionColdStart } from "./cold-start";
 import { promisifiedHandler } from "./handler";
 import { LogLevel, setLogLevel } from "./log";
+import { HANDLER_STREAMING, STREAM_RESPONSE } from "../constants";
 
 const mockContext = {
   invokedFunctionArn: "arn:aws:lambda:us-east-1:123497598159:function:my-test-lambda",
@@ -330,3 +331,165 @@ describe("promisifiedHandler", () => {
     expect(result).toEqual({ statusCode: 200, body: "Plain response", headers: { "Content-Type": "text/plain" } });
   });
 });
+
+describe("buffered handler completion mode", () => {
+  let oldCompletion: string | undefined;
+  const response = { statusCode: 200, body: "callback response" };
+  const incidental = { incidental: true };
+
+  function createHandler(kind: string, work: (callback: Callback<unknown>, context: Context) => unknown) {
+    if (kind === "rest") {
+      return (...args: Parameters<Handler>) => work(args[2], args[1]);
+    }
+    return (event: unknown, context: Context, callback: Callback<unknown> = () => {}) => work(callback, context);
+  }
+
+  function wrap(handler: unknown) {
+    return promisifiedHandler(handler) as (event: unknown, context: Context) => Promise<unknown>;
+  }
+
+  beforeEach(() => {
+    oldCompletion = process.env.DD_LAMBDA_HANDLER_COMPLETION;
+    delete process.env.DD_LAMBDA_HANDLER_COMPLETION;
+    jest.useFakeTimers();
+  });
+
+  afterEach(() => {
+    jest.useRealTimers();
+    if (oldCompletion === undefined) {
+      delete process.env.DD_LAMBDA_HANDLER_COMPLETION;
+    } else {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = oldCompletion;
+    }
+  });
+
+  describe.each(["rest", "default"])("%s parameters", (kind) => {
+    it("preserves the existing incidental-return selection in auto mode", async () => {
+      const handler = createHandler(kind, (callback) => {
+        setTimeout(() => callback(null, response), 10);
+        return incidental;
+      });
+      expect(handler.length).toBe(kind === "rest" ? 0 : 2);
+      await expect(wrap(handler)({}, { ...mockContext })).resolves.toBe(incidental);
+      jest.advanceTimersByTime(10);
+    });
+
+    it("waits for the asynchronous callback instead of returning the incidental object", async () => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      const handler = createHandler(kind, (callback) => {
+        setTimeout(() => callback(null, response), 10);
+        return incidental;
+      });
+      const settled = jest.fn();
+      const result = wrap(handler)({}, { ...mockContext });
+      const observed = result.then(settled);
+      jest.advanceTimersByTime(9);
+      await Promise.resolve();
+      expect(settled).not.toHaveBeenCalled();
+      jest.advanceTimersByTime(1);
+      await expect(result).resolves.toBe(response);
+      await observed;
+    });
+
+    it("propagates asynchronous callback errors", async () => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      const error = new Error("callback failure");
+      const handler = createHandler(kind, (callback) => {
+        setTimeout(() => callback(error), 10);
+        return incidental;
+      });
+      const result = wrap(handler)({}, { ...mockContext });
+      const observed = expect(result).rejects.toBe(error);
+      jest.advanceTimersByTime(10);
+      await observed;
+    });
+
+    it("stays pending if neither callback nor context completion fires", async () => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      const handler = createHandler(kind, () => incidental);
+      const settled = jest.fn();
+      wrap(handler)({}, { ...mockContext }).then(settled, settled);
+      jest.advanceTimersByTime(1000);
+      await Promise.resolve();
+      expect(settled).not.toHaveBeenCalled();
+    });
+
+    it.each(["done", "succeed", "fail"])("honors context.%s with an incidental return", async (method) => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      const error = new Error("context failure");
+      const handler = createHandler(kind, (_callback, context) => {
+        setTimeout(() => {
+          if (method === "done") context.done(undefined, response);
+          else if (method === "succeed") context.succeed(response);
+          else context.fail(error);
+        }, 10);
+        return incidental;
+      });
+      const result = wrap(handler)({}, { ...mockContext });
+      const observed = method === "fail" ? expect(result).rejects.toBe(error) : expect(result).resolves.toBe(response);
+      jest.advanceTimersByTime(10);
+      await observed;
+    });
+
+    it.each(["callback", "promise", "rejection"])("preserves first completion from %s", async (winner) => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      const error = new Error("promise failure");
+      const handler = createHandler(kind, (callback) => {
+        setTimeout(() => callback(null, response), winner === "callback" ? 10 : 20);
+        return new Promise((resolve, reject) => {
+          setTimeout(
+            () => (winner === "rejection" ? reject(error) : resolve("promise response")),
+            winner === "callback" ? 20 : 10,
+          );
+        });
+      });
+      const result = wrap(handler)({}, { ...mockContext });
+      const observed =
+        winner === "rejection"
+          ? expect(result).rejects.toBe(error)
+          : expect(result).resolves.toEqual(winner === "callback" ? response : "promise response");
+      jest.advanceTimersByTime(10);
+      await observed;
+      jest.advanceTimersByTime(10);
+    });
+
+    it("preserves synchronous throws", () => {
+      process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+      const error = new Error("handler failure");
+      expect(() =>
+        wrap(
+          createHandler(kind, () => {
+            throw error;
+          }),
+        )({}, { ...mockContext }),
+      ).toThrow(error);
+    });
+  });
+
+  it.each([undefined, "auto", "unknown"])("preserves genuine synchronous responses for %s", async (mode) => {
+    if (mode !== undefined) process.env.DD_LAMBDA_HANDLER_COMPLETION = mode;
+    await expect(wrap(() => response)({}, { ...mockContext })).resolves.toBe(response);
+  });
+
+  it("reads the mode when wrapping and accepts uppercase callback", async () => {
+    process.env.DD_LAMBDA_HANDLER_COMPLETION = "CALLBACK";
+    const invoke = wrap((_event: unknown, context: Context) => {
+      context.succeed(response);
+      return incidental;
+    });
+    delete process.env.DD_LAMBDA_HANDLER_COMPLETION;
+    await expect(invoke({}, { ...mockContext })).resolves.toBe(response);
+    await expect(wrap(() => incidental)({}, { ...mockContext })).resolves.toBe(incidental);
+  });
+
+  it("leaves response streaming completion and arguments unchanged", async () => {
+    process.env.DD_LAMBDA_HANDLER_COMPLETION = "callback";
+    const stream = {};
+    const context = { ...mockContext };
+    const original = jest.fn().mockResolvedValue(response);
+    const handler = Object.assign(original, { [HANDLER_STREAMING]: STREAM_RESPONSE });
+    const invoke = promisifiedHandler(handler) as (event: unknown, output: unknown, ctx: Context) => Promise<unknown>;
+    await expect(invoke({}, stream, context)).resolves.toBe(response);
+    expect(original).toHaveBeenCalledWith({}, stream, context);
+  });
+});

src/utils/handler.ts

@@ -12,6 +12,8 @@ export function promisifiedHandler<TEvent, TResult>(handler: Handler<TEvent, TRe
     };
   }
 
+  const callbackCompletion = process.env.DD_LAMBDA_HANDLER_COMPLETION?.toLowerCase() === "callback";
+
   // Buffered Lambda function.
   return (event: TEvent, context: Context) => {
     // Lambda functions in node complete in one of two possible ways.
@@ -60,8 +62,8 @@ export function promisifiedHandler<TEvent, TResult>(handler: Handler<TEvent, TRe
     if (asyncProm !== undefined && typeof (asyncProm as any).then === "function") {
       // Mimics behaviour of lambda runtime, the first method of returning a result always wins.
       promise = Promise.race([callbackProm, asyncProm as Promise<TResult>]);
-    } else if (handler.length >= 3) {
-      // Handler takes a callback, wait for the callback to be called
+    } else if (callbackCompletion || handler.length >= 3) {
+      // Explicit callback completion also covers rest/default parameters, which reduce handler.length.
       promise = callbackProm;
     } else if (asyncProm === undefined) {
       // Handler returned nothing (implicit `undefined`) and doesn't take a callback parameter.