perf(runtime): use dd-trace loader registration

.github/workflows/build.yml

@@ -69,9 +69,6 @@ jobs:
       - name: Install dependencies
         run: yarn install
 
-      - name: Build
-        run: yarn build
-
       - name: Run tests
         run: yarn test
 

package.json

@@ -12,8 +12,8 @@
   "license": "Apache-2.0",
   "scripts": {
     "build": "tsc",
-    "test": "jest",
-    "test:watch": "jest --watch",
+    "test": "node scripts/prepare_tests.js && jest",
+    "test:watch": "node scripts/prepare_tests.js && jest --watch",
     "lint": "tslint --project tsconfig.json",
     "check-formatting": "prettier --check src/**",
     "format": "prettier --write src/**",

scripts/prepare_tests.js

@@ -0,0 +1,25 @@
+const { spawnSync } = require("node:child_process");
+const path = require("node:path");
+
+const yarn = process.platform === "win32" ? "yarn.cmd" : "yarn";
+
+/**
+ * @param {string[]} arguments_ Yarn arguments.
+ * @param {string} [cwd] Working directory.
+ */
+function runYarn(arguments_, cwd) {
+  const result = spawnSync(yarn, arguments_, { cwd, stdio: "inherit" });
+  if (result.error !== undefined) {
+    throw result.error;
+  }
+  if (result.status !== 0) {
+    process.exit(result.status ?? 1);
+  }
+}
+
+runYarn(["build"]);
+
+const nodeMajor = Number(process.versions.node.split(".")[0]);
+if (nodeMajor >= 22) {
+  runYarn(["install", "--frozen-lockfile"], path.join(__dirname, "..", "src", "runtime", "fixtures"));
+}

src/handler.mjs

@@ -1,6 +1,3 @@
-import * as Module from "module";
-import { dirname } from "path";
-import { pathToFileURL } from "url";
 import {
   datadog,
   datadogHandlerEnvVar,
@@ -18,41 +15,8 @@ if (process.env.DD_TRACE_DISABLED_PLUGINS === undefined) {
   logDebug("disabled the dd-trace plugin 'fs'");
 }
 
-// True when dd-trace's ESM loader hook is already active via NODE_OPTIONS or
-// execArgv. Covers every preload entry point that registers the hook:
-//   --import dd-trace/initialize.mjs  (documented ESM setup)
-//   --loader dd-trace/initialize.mjs  (legacy)
-//   --require dd-trace/register.js    (documented CJS setup; also registers the hook)
-// module.register() does not dedupe — a second registration chains another hooks
-// worker and every module would be rewritten twice.
-function esmLoaderAlreadyRegistered() {
-  const sources = [process.env.NODE_OPTIONS || "", ...process.execArgv];
-  return sources.some((source) => /dd-trace[\\/](?:[^\s]*\.mjs|register\.js)/.test(source));
-}
-
 if (getEnvValue("DD_TRACE_ENABLED", "true").toLowerCase() === "true") {
-  const tracer = initTracer();
-
-  // Register dd-trace's ESM loader hooks programmatically so that ESM imports
-  // (e.g. @aws/durable-execution-sdk-js) are rewritten for instrumentation.
-  // Normally this happens via --import dd-trace/initialize.mjs, but the AWS
-  // durable runtime ignores NODE_OPTIONS so the loader is never registered.
-  // This mirrors what dd-trace/initialize.mjs does at lines 77-84, including
-  // only registering when the tracer initialized — a bail-out leaves the hooks
-  // worker with nothing to instrument and can keep the process from exiting.
-  if (tracer && typeof Module.register === "function" && !esmLoaderAlreadyRegistered()) {
-    try {
-      const require = Module.createRequire(import.meta.url);
-      const ddTraceEntry = require.resolve("dd-trace", {
-        paths: ["/var/task/node_modules", ...(require.resolve.paths("dd-trace") || [])],
-      });
-      const ddTraceRoot = pathToFileURL(dirname(ddTraceEntry) + "/").href;
-      Module.register("./loader-hook.mjs", ddTraceRoot);
-      logDebug("registered dd-trace ESM loader hook for ESM instrumentation");
-    } catch (error) {
-      logDebug("failed to register dd-trace ESM loader hook", { error });
-    }
-  }
+  initTracer();
 }
 
 const taskRootEnv = getEnvValue(lambdaTaskRootEnvVar, "");

src/handler.spec.ts

@@ -0,0 +1,211 @@
+import { spawn } from "node:child_process";
+import { once } from "node:events";
+import * as http from "node:http";
+import * as Module from "node:module";
+import * as path from "node:path";
+import { pathToFileURL } from "node:url";
+import * as zlib from "node:zlib";
+
+const fixtureDirectory = path.join(__dirname, "runtime", "fixtures");
+const handlerEntry = pathToFileURL(path.join(__dirname, "..", "dist", "handler.mjs")).href;
+const runnerPath = path.join(fixtureDirectory, "published-handler-runner.mjs");
+const tracePayloads: Buffer[] = [];
+
+let agent: http.Server;
+let resolveTrace: (() => void) | undefined;
+let traceReceived: Promise<void>;
+
+/**
+ * @param {Buffer} body Encoded trace body.
+ * @param {string | string[] | undefined} encoding Content encoding from the trace request.
+ */
+function recordTrace(body: Buffer, encoding: string | string[] | undefined): void {
+  let payload = body;
+  if (encoding === "gzip") {
+    payload = zlib.gunzipSync(body);
+  } else if (encoding === "deflate") {
+    payload = zlib.inflateSync(body);
+  }
+
+  tracePayloads.push(payload);
+  resolveTrace?.();
+}
+
+/**
+ * @param {http.IncomingMessage} request Agent request.
+ * @param {http.ServerResponse} response Agent response.
+ */
+function respond(request: http.IncomingMessage, response: http.ServerResponse): void {
+  const chunks: Buffer[] = [];
+
+  /** @param {Buffer} chunk Request body chunk. */
+  function recordChunk(chunk: Buffer): void {
+    chunks.push(chunk);
+  }
+
+  request.on("data", recordChunk);
+  request.once("end", () => {
+    response.setHeader("content-type", "application/json");
+
+    if (request.url === "/info") {
+      response.end(JSON.stringify({ endpoints: ["/v0.4/traces", "/v0.5/traces"] }));
+      return;
+    }
+
+    if (request.url?.startsWith("/v0.4/traces") || request.url?.startsWith("/v0.5/traces")) {
+      recordTrace(Buffer.concat(chunks), request.headers["content-encoding"]);
+    }
+
+    response.end(JSON.stringify({ rate_by_service: {} }));
+  });
+}
+
+/** @param {() => void} resolve Trace promise resolver. */
+function captureTraceResolver(resolve: () => void): void {
+  resolveTrace = resolve;
+}
+
+/** @param {Promise<[unknown]>} messagePromise Child IPC message signal. */
+async function readChildMessage(messagePromise: Promise<[unknown]>): Promise<unknown> {
+  const [message] = await messagePromise;
+  return message;
+}
+
+/** @param {boolean} traceEnabled Whether tracing is enabled for the child process. */
+async function runPublishedHandler(traceEnabled: boolean): Promise<unknown> {
+  const address = agent.address();
+  if (address === null || typeof address === "string") {
+    throw new Error("Mock agent is not listening on a TCP port");
+  }
+
+  const env: NodeJS.ProcessEnv = {
+    ...process.env,
+    AWS_ACCESS_KEY_ID: "test-access-key",
+    AWS_EXECUTION_ENV: `AWS_Lambda_nodejs${process.versions.node.split(".")[0]}.x`,
+    AWS_LAMBDA_FUNCTION_MEMORY_SIZE: "1024",
+    AWS_LAMBDA_FUNCTION_VERSION: "$LATEST",
+    AWS_REGION: "us-east-1",
+    AWS_SECRET_ACCESS_KEY: "test-secret-key",
+    DD_COLD_START_TRACING: "false",
+    DD_INSTRUMENTATION_TELEMETRY_ENABLED: "false",
+    DD_LAMBDA_HANDLER: "esm-instrumentation-handler.handle",
+    DD_PROFILING_ENABLED: "false",
+    DD_REMOTE_CONFIGURATION_ENABLED: "false",
+    DD_RUNTIME_METRICS_ENABLED: "false",
+    DD_TELEMETRY_ENABLED: "false",
+    DD_TRACE_AGENT_URL: `http://127.0.0.1:${address.port}`,
+    DD_TRACE_ENABLED: String(traceEnabled),
+    DD_TRACE_STARTUP_LOGS: "false",
+    LAMBDA_TASK_ROOT: fixtureDirectory,
+    HANDLER_TEST_ENTRY: handlerEntry,
+  };
+  delete env.NODE_OPTIONS;
+  delete env.OTEL_LOGS_EXPORTER;
+  delete env.OTEL_METRICS_EXPORTER;
+  delete env.OTEL_TRACES_EXPORTER;
+
+  const child = spawn(process.execPath, [runnerPath], {
+    env,
+    stdio: ["ignore", "pipe", "pipe", "ipc"],
+  });
+  const closePromise = once(child, "close") as Promise<[number | null, NodeJS.Signals | null]>;
+  const messagePromise = once(child, "message") as Promise<[unknown]>;
+  let standardOutput = "";
+  let standardError = "";
+
+  child.stdout?.setEncoding("utf8");
+  child.stderr?.setEncoding("utf8");
+
+  /** @param {string} chunk Standard output chunk. */
+  function recordStandardOutput(chunk: string): void {
+    standardOutput += chunk;
+  }
+
+  /** @param {string} chunk Standard error chunk. */
+  function recordStandardError(chunk: string): void {
+    standardError += chunk;
+  }
+
+  async function failOnEarlyClose(): Promise<never> {
+    const [exitCode, signal] = await closePromise;
+    throw new Error(
+      `Published handler exited before completing the test: code=${exitCode}, signal=${signal}` +
+        `\nstdout:\n${standardOutput}\nstderr:\n${standardError}`,
+    );
+  }
+
+  child.stdout?.on("data", recordStandardOutput);
+  child.stderr?.on("data", recordStandardError);
+  const watchdog = setTimeout(() => child.kill("SIGKILL"), 15_000);
+  watchdog.unref();
+
+  try {
+    const message = await Promise.race([readChildMessage(messagePromise), failOnEarlyClose()]);
+
+    if (traceEnabled) {
+      await Promise.race([traceReceived, failOnEarlyClose()]);
+    }
+
+    child.send("exit");
+    const [exitCode, signal] = await closePromise;
+    if (exitCode !== 0 || signal !== null) {
+      throw new Error(`Published handler failed: code=${exitCode}, signal=${signal}\n${standardError}`);
+    }
+
+    return message;
+  } finally {
+    clearTimeout(watchdog);
+    if (child.exitCode === null && child.signalCode === null) {
+      child.kill("SIGKILL");
+    }
+  }
+}
+
+const moduleWithRegister = Module as typeof Module & { register?: unknown };
+const nodeMajor = Number(process.versions.node.split(".")[0]);
+const supportsDurableFixture = nodeMajor >= 22 && typeof moduleWithRegister.register === "function";
+const describeWithESMLoader = supportsDurableFixture ? describe : describe.skip;
+
+describeWithESMLoader("published ESM handler", () => {
+  jest.setTimeout(30_000);
+
+  beforeAll(async () => {
+    agent = http.createServer(respond);
+    const listening = once(agent, "listening");
+    agent.listen(0, "127.0.0.1");
+    await listening;
+  });
+
+  afterAll(async () => {
+    const closed = once(agent, "close");
+    agent.close();
+    await closed;
+  });
+
+  beforeEach(() => {
+    tracePayloads.length = 0;
+    traceReceived = new Promise(captureTraceResolver);
+  });
+
+  it("instruments an ESM durable handler through the published handler", async () => {
+    const message = await runPublishedHandler(true);
+    const payload = Buffer.concat(tracePayloads);
+
+    expect(message).toEqual({
+      registerLoaded: true,
+      result: expect.objectContaining({ Status: "SUCCEEDED" }),
+    });
+    expect(payload.includes(Buffer.from("aws.lambda"))).toBe(true);
+    expect(payload.includes(Buffer.from("aws.durable.execute"))).toBe(true);
+  });
+
+  it("does not register tracing when tracing is disabled", async () => {
+    const message = await runPublishedHandler(false);
+
+    expect(message).toEqual({
+      registerLoaded: false,
+      result: expect.objectContaining({ Status: "SUCCEEDED" }),
+    });
+    expect(tracePayloads).toHaveLength(0);
+  });
+});

src/runtime/fixtures/esm-instrumentation-handler.mjs

@@ -0,0 +1,19 @@
+import { withDurableExecution } from "@aws/durable-execution-sdk-js";
+
+const durableExecutionClient = {
+  getExecutionState() {
+    return Promise.resolve({ Operations: [] });
+  },
+  checkpoint() {
+    return Promise.resolve({
+      CheckpointToken: "next-checkpoint-token",
+      NewExecutionState: { Operations: [] },
+    });
+  },
+};
+
+async function customerHandler() {
+  return { statusCode: 200 };
+}
+
+export const handle = withDurableExecution(customerHandler, { durableExecutionClient });

src/runtime/fixtures/package.json

@@ -0,0 +1,8 @@
+{
+  "name": "datadog-lambda-js-durable-handler-fixture",
+  "private": true,
+  "type": "module",
+  "dependencies": {
+    "@aws/durable-execution-sdk-js": "2.3.0"
+  }
+}

src/runtime/fixtures/published-handler-runner.mjs

@@ -0,0 +1,43 @@
+import { once } from "node:events";
+import { createRequire } from "node:module";
+
+const handlerEntry = process.env.HANDLER_TEST_ENTRY;
+if (handlerEntry === undefined) {
+  throw new Error("HANDLER_TEST_ENTRY is required");
+}
+if (process.send === undefined) {
+  throw new Error("The published handler runner requires an IPC channel");
+}
+
+const { handler } = await import(handlerEntry);
+const require = createRequire(handlerEntry);
+const registerPath = require.resolve("dd-trace/register.js");
+const result = await handler(
+  {
+    DurableExecutionArn:
+      "arn:aws:lambda:us-east-1:123456789012:function:loader-instrumentation-test:$LATEST/durable-execution/test/test-id",
+    CheckpointToken: "checkpoint-token",
+    InitialExecutionState: { Operations: [] },
+  },
+  {
+    callbackWaitsForEmptyEventLoop: false,
+    functionName: "loader-instrumentation-test",
+    functionVersion: "$LATEST",
+    invokedFunctionArn: "arn:aws:lambda:us-east-1:123456789012:function:loader-instrumentation-test",
+    memoryLimitInMB: "1024",
+    awsRequestId: "test-request",
+    logGroupName: "/aws/lambda/loader-instrumentation-test",
+    logStreamName: "2026/08/28/[$LATEST]test",
+    getRemainingTimeInMillis() {
+      return 30_000;
+    },
+  },
+);
+
+const exitSignal = once(process, "message");
+process.send({
+  registerLoaded: require.cache[registerPath] !== undefined,
+  result,
+});
+await exitSignal;
+process.disconnect();