From fb348f1ae4f22eb726ba9592be41a0e0b00b96a8 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:36:17 -0500 Subject: [PATCH 1/3] fix(lambda): prefer registered ESM loader entrypoint --- src/handler.mjs | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/src/handler.mjs b/src/handler.mjs index ae586010..252fddaf 100644 --- a/src/handler.mjs +++ b/src/handler.mjs @@ -30,6 +30,31 @@ function esmLoaderAlreadyRegistered() { return sources.some((source) => /dd-trace[\\/](?:[^\s]*\.mjs|register\.js)/.test(source)); } +function esmLoaderEnabled() { + return getEnvValue("DD_TRACE_ESM_LOADER_ENABLED", "true").toLowerCase() !== "false"; +} + +function registerDdTraceEsmLoaderHook() { + const require = Module.createRequire(import.meta.url); + const ddTraceResolvePaths = ["/var/task/node_modules", ...(require.resolve.paths("dd-trace") || [])]; + + try { + const ddTraceRegister = require.resolve("dd-trace/register.js", { paths: ddTraceResolvePaths }); + require(ddTraceRegister); + logDebug("registered dd-trace ESM loader hook for ESM instrumentation"); + return; + } catch (registerError) { + try { + const ddTraceEntry = require.resolve("dd-trace", { paths: ddTraceResolvePaths }); + const ddTraceRoot = pathToFileURL(dirname(ddTraceEntry) + "/").href; + Module.register("./loader-hook.mjs", ddTraceRoot); + logDebug("registered dd-trace ESM loader hook for ESM instrumentation through direct fallback"); + } catch (fallbackError) { + logDebug("failed to register dd-trace ESM loader hook", { registerError, fallbackError }); + } + } +} + if (getEnvValue("DD_TRACE_ENABLED", "true").toLowerCase() === "true") { const tracer = initTracer(); @@ -40,18 +65,8 @@ if (getEnvValue("DD_TRACE_ENABLED", "true").toLowerCase() === "true") { // 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 }); - } + if (tracer && esmLoaderEnabled() && typeof Module.register === "function" && !esmLoaderAlreadyRegistered()) { + registerDdTraceEsmLoaderHook(); } } From 5af0fdaacc25ec99192fe84a4603d5a2d2572bd3 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:36:32 -0500 Subject: [PATCH 2/3] fix(lambda): clear pending cold start module state --- src/trace/listener.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/trace/listener.ts b/src/trace/listener.ts index 1483a6ad..a85821ae 100644 --- a/src/trace/listener.ts +++ b/src/trace/listener.ts @@ -228,9 +228,9 @@ export class TraceListener { const coldStartTracer = new ColdStartTracer(coldStartConfig); coldStartTracer.trace(coldStartNodes); } - // Always clear the tree to prevent memory leaks, even if we skip span creation - clearTraceTree(); } + // Always clear the tree to prevent memory leaks, even when no balanced root was formed. + clearTraceTree(); if (this.config.appsecEnabled) { processAppsecResponse(this.tracerWrapper.currentSpan, result); } From 33e32f9c4c2b7dd00e3c5ad0fac3ce1e7a892fa6 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:36:46 -0500 Subject: [PATCH 3/3] fix(lambda): trace ESM handler imports during cold start --- src/runtime/index.ts | 2 +- src/runtime/require-tracer.spec.ts | 114 ++++++++++++++++++++++++++-- src/runtime/require-tracer.ts | 98 ++++++++++++++++++++++-- src/runtime/user-function.spec.ts | 43 +++++++++++ src/runtime/user-function.ts | 19 ++++- src/trace/cold-start-tracer.spec.ts | 61 +++++++++++++++ src/trace/cold-start-tracer.ts | 17 +++-- 7 files changed, 333 insertions(+), 21 deletions(-) create mode 100644 src/runtime/user-function.spec.ts diff --git a/src/runtime/index.ts b/src/runtime/index.ts index 837b0ec2..c82b230a 100644 --- a/src/runtime/index.ts +++ b/src/runtime/index.ts @@ -1,2 +1,2 @@ export { load } from "./user-function"; -export { subscribeToDC, getTraceTree, clearTraceTree, RequireNode } from "./require-tracer" +export { subscribeToDC, getTraceTree, clearTraceTree, recordModuleLoad, currentTime, RequireNode } from "./require-tracer" diff --git a/src/runtime/require-tracer.spec.ts b/src/runtime/require-tracer.spec.ts index d0bcda88..d3cad5cf 100644 --- a/src/runtime/require-tracer.spec.ts +++ b/src/runtime/require-tracer.spec.ts @@ -1,12 +1,19 @@ -import { subscribeToDC, getTraceTree, RequireNode } from "./require-tracer"; +import { subscribeToDC, getTraceTree, clearTraceTree, recordModuleLoad, RequireNode } from "./require-tracer"; const dc = require('dc-polyfill') describe('require-tracer', () => { - it('generates a trace tree', () => { + const moduleLoadStartChannel = dc.channel('dd-trace:moduleLoadStart') + const moduleLoadEndChannel = dc.channel('dd-trace:moduleLoadEnd') + + beforeEach(() => { subscribeToDC() - const moduleLoadStartChannel = dc.channel('dd-trace:moduleLoadStart') - const moduleLoadEndChannel = dc.channel('dd-trace:moduleLoadEnd') + }) + afterEach(() => { + clearTraceTree() + }) + + it('generates a trace tree', () => { // require('myLibrary') moduleLoadStartChannel.publish({ request: 'myLibrary' @@ -18,7 +25,7 @@ describe('require-tracer', () => { moduleLoadEndChannel.publish() moduleLoadEndChannel.publish() const res = getTraceTree() - expect(res).toBeDefined + expect(res).toBeDefined() expect(res[0].id).toBe('myLibrary') const resChildren = res[0].children as RequireNode[] expect(resChildren).toHaveLength(1) @@ -26,4 +33,101 @@ describe('require-tracer', () => { expect(resChild.id).toBe('myChildLibrary') }) + it('ignores unmatched module load end events', () => { + moduleLoadEndChannel.publish({ + request: 'missing' + }) + + expect(getTraceTree()).toStrictEqual([]) + }) + + it('clears pending module load state', () => { + moduleLoadStartChannel.publish({ + request: 'unfinished' + }) + + clearTraceTree() + + moduleLoadStartChannel.publish({ + request: 'finished' + }) + moduleLoadEndChannel.publish({ + request: 'finished' + }) + + const res = getTraceTree() + expect(res).toHaveLength(1) + expect(res[0].id).toBe('finished') + }) + + it('handles ESM loader payloads without a request', () => { + const filename = 'file:///var/task/index.mjs' + + moduleLoadStartChannel.publish({ + filename + }) + moduleLoadEndChannel.publish({ + filename + }) + + const res = getTraceTree() + expect(res).toHaveLength(1) + expect(res[0].id).toBe(filename) + expect(res[0].filename).toBe(filename) + expect(res[0].endTime).toBeGreaterThanOrEqual(res[0].startTime) + }) + + it('records measured ESM imports', () => { + recordModuleLoad({ + id: '/var/task/index.mjs', + filename: '/var/task/index.mjs', + startTime: 100, + endTime: 150, + kind: 'import' + }) + + const res = getTraceTree() + expect(res).toHaveLength(1) + expect(res[0].id).toBe('/var/task/index.mjs') + expect(res[0].filename).toBe('/var/task/index.mjs') + expect(res[0].startTime).toBe(100) + expect(res[0].endTime).toBe(150) + expect(res[0].kind).toBe('import') + }) + + it('does not record measured imports with invalid durations', () => { + recordModuleLoad({ + id: '/var/task/index.mjs', + filename: '/var/task/index.mjs', + startTime: 150, + endTime: 100, + kind: 'import' + }) + + expect(getTraceTree()).toStrictEqual([]) + }) + + it('parents roots loaded during a measured ESM import', () => { + recordModuleLoad({ + id: '@aws-sdk/core/client', + filename: '/var/task/node_modules/@aws-sdk/core/dist-cjs/client.js', + startTime: 120, + endTime: 130 + }) + recordModuleLoad({ + id: '/var/task/index.mjs', + filename: '/var/task/index.mjs', + startTime: 100, + endTime: 150, + kind: 'import', + absorbChildren: true + }) + + const res = getTraceTree() + expect(res).toHaveLength(1) + expect(res[0].id).toBe('/var/task/index.mjs') + expect(res[0].kind).toBe('import') + expect(res[0].children).toHaveLength(1) + expect(res[0].children[0].id).toBe('@aws-sdk/core/client') + }) }) diff --git a/src/runtime/require-tracer.ts b/src/runtime/require-tracer.ts index 2946f408..c3d399be 100644 --- a/src/runtime/require-tracer.ts +++ b/src/runtime/require-tracer.ts @@ -1,18 +1,33 @@ +import { performance } from "perf_hooks"; + const dc = require('dc-polyfill') +export type ModuleLoadKind = 'require' | 'import' + +export interface ModuleLoadTrace { + id: string + filename?: string + startTime: number + endTime: number + kind?: ModuleLoadKind + absorbChildren?: boolean +} + export class RequireNode { public id: string public filename: string public startTime: number public endTime: number public children: RequireNode[] + public kind: ModuleLoadKind - constructor(id: string, filename: string, startTime: number) { + constructor(id: string, filename: string, startTime: number, kind: ModuleLoadKind = 'require') { this.id = id this.filename = filename this.startTime = startTime this.endTime = startTime this.children = [] + this.kind = kind } public set setEnd(endTime: number) { @@ -23,12 +38,15 @@ export class RequireNode { const moduleLoadStartChannel = dc.channel('dd-trace:moduleLoadStart') const moduleLoadEndChannel = dc.channel('dd-trace:moduleLoadEnd') let rootNodes: RequireNode[] = [] +let isSubscribed = false const requireStack: RequireNode[] = [] const pushNode = (data: any) => { - const startTime = Date.now() + const startTime = now() + const id = getModuleId(data) + const filename = getModuleFilename(data, id) - const reqNode = new RequireNode(data.request, data.filename, startTime) + const reqNode = new RequireNode(id, filename, startTime) const maybeParent = requireStack[requireStack.length - 1] if (maybeParent) { @@ -37,10 +55,13 @@ const pushNode = (data: any) => { requireStack.push(reqNode) } -const popNode = () => { - const endTime = Date.now() - const reqNode = requireStack.pop() - if (reqNode){ +const popNode = (data?: any) => { + const reqNodeIndex = findMatchingNodeIndex(data) + if (reqNodeIndex < 0) return + + const endTime = now() + const reqNode = requireStack.splice(reqNodeIndex, 1)[0] + if (reqNode) { reqNode.endTime = endTime } if (requireStack.length <= 0 && reqNode) { @@ -48,7 +69,69 @@ const popNode = () => { } } +const findMatchingNodeIndex = (data?: any): number => { + const stackTopIndex = requireStack.length - 1 + if (!hasModuleIdentity(data) || stackTopIndex < 0) return stackTopIndex + + const id = getModuleId(data) + const filename = getModuleFilename(data, id) + for (let index = stackTopIndex; index >= 0; index--) { + const reqNode = requireStack[index] + if (reqNode.id === id && reqNode.filename === filename) { + return index + } + } + + return -1 +} + +const hasModuleIdentity = (data: any): boolean => { + return !!(data?.request || data?.filename) +} + +const getModuleId = (data: any): string => { + return String(data?.request || data?.filename || "") +} + +const getModuleFilename = (data: any, id: string): string => { + return String(data?.filename || data?.request || id) +} + +const now = (): number => { + return performance.timeOrigin + performance.now() +} + +const isWithin = (node: RequireNode, parent: RequireNode): boolean => { + return node.startTime >= parent.startTime && node.endTime <= parent.endTime +} + +export const recordModuleLoad = (moduleLoad: ModuleLoadTrace): void => { + if (moduleLoad.endTime < moduleLoad.startTime) return + + const reqNode = new RequireNode( + moduleLoad.id, + moduleLoad.filename || moduleLoad.id, + moduleLoad.startTime, + moduleLoad.kind + ) + reqNode.endTime = moduleLoad.endTime + + if (moduleLoad.absorbChildren) { + const children = rootNodes.filter((rootNode) => isWithin(rootNode, reqNode)) + rootNodes = rootNodes.filter((rootNode) => !isWithin(rootNode, reqNode)) + reqNode.children.push(...children) + } + + rootNodes.push(reqNode) +} + +export const currentTime = (): number => { + return now() +} + export const subscribeToDC = () => { + if (isSubscribed) return + isSubscribed = true moduleLoadStartChannel.subscribe(pushNode) moduleLoadEndChannel.subscribe(popNode) } @@ -59,4 +142,5 @@ export const getTraceTree = (): RequireNode[] => { export const clearTraceTree = () => { rootNodes = [] + requireStack.length = 0 } diff --git a/src/runtime/user-function.spec.ts b/src/runtime/user-function.spec.ts new file mode 100644 index 00000000..64501f16 --- /dev/null +++ b/src/runtime/user-function.spec.ts @@ -0,0 +1,43 @@ +import fs from "fs"; + +import { clearTraceTree, getTraceTree } from "./require-tracer"; +import { load } from "./user-function"; + +const mockImport = jest.fn(); + +jest.mock("./module_importer", () => ({ + import: (...args: any[]) => mockImport(...args), +})); + +describe("user-function", () => { + let existsSyncSpy: jest.SpyInstance; + + beforeEach(() => { + clearTraceTree(); + mockImport.mockReset(); + existsSyncSpy = jest.spyOn(fs, "existsSync").mockImplementation((filePath) => { + return String(filePath) === "/var/task/index.mjs"; + }); + }); + + afterEach(() => { + existsSyncSpy.mockRestore(); + clearTraceTree(); + }); + + it("records ESM handler imports with real measured duration", async () => { + const handler = jest.fn(); + mockImport.mockResolvedValue({ handler }); + + const loadedHandler = await load("/var/task", "index.handler"); + + expect(loadedHandler).toBe(handler); + expect(mockImport).toHaveBeenCalledWith("/var/task/index.mjs"); + const traceTree = getTraceTree(); + expect(traceTree).toHaveLength(1); + expect(traceTree[0].id).toBe("/var/task/index.mjs"); + expect(traceTree[0].filename).toBe("/var/task/index.mjs"); + expect(traceTree[0].kind).toBe("import"); + expect(traceTree[0].endTime).toBeGreaterThanOrEqual(traceTree[0].startTime); + }); +}); diff --git a/src/runtime/user-function.ts b/src/runtime/user-function.ts index 9c0df1f8..7ba8de39 100644 --- a/src/runtime/user-function.ts +++ b/src/runtime/user-function.ts @@ -19,6 +19,7 @@ import { ImportModuleError, UserCodeSyntaxError, } from "./errors"; +import { currentTime, recordModuleLoad } from "./require-tracer"; const module_importer = require("./module_importer"); const FUNCTION_EXPR = /^([^.]*)\.(.*)$/; @@ -66,10 +67,22 @@ function _tryRequireFile(file: string, extension?: string): any { } async function _tryAwaitImport(file: string, extension: string): Promise { - const path = file + (extension || ""); + const modulePath = file + (extension || ""); - if (fs.existsSync(path)) { - return await module_importer.import(path); + if (fs.existsSync(modulePath)) { + const startTime = currentTime(); + try { + return await module_importer.import(modulePath); + } finally { + recordModuleLoad({ + id: modulePath, + filename: modulePath, + startTime, + endTime: currentTime(), + kind: "import", + absorbChildren: true, + }); + } } } diff --git a/src/trace/cold-start-tracer.spec.ts b/src/trace/cold-start-tracer.spec.ts index 78827d86..7c3137ec 100644 --- a/src/trace/cold-start-tracer.spec.ts +++ b/src/trace/cold-start-tracer.spec.ts @@ -260,4 +260,65 @@ describe("ColdStartTracer", () => { service: "aws.lambda", }); }); + + it("traces measured ESM imports and extends the load span to handler ready", () => { + const requireNodes: RequireNode[] = [ + { + id: "./hooks", + filename: "/opt/nodejs/node_modules/dd-trace/packages/dd-trace/src/hooks.js", + startTime: 20, + endTime: 30, + children: [], + } as any as RequireNode, + { + id: "/var/task/index.mjs", + filename: "/var/task/index.mjs", + startTime: 100, + endTime: 220, + kind: "import", + children: [ + { + id: "@aws-sdk/core/client", + filename: "/var/task/node_modules/@aws-sdk/core/dist-cjs/client.js", + startTime: 150, + endTime: 200, + children: [], + }, + ], + } as any as RequireNode, + ]; + const coldStartConfig: ColdStartTracerConfig = { + tracerWrapper: new TracerWrapper(), + parentSpan: { + span: {}, + name: "my-lambda-span", + } as any as SpanWrapper, + lambdaFunctionName: "my-function-name", + currentSpanStartTime: 500, + minDuration: 1, + ignoreLibs: "", + isColdStart: true, + }; + const coldStartTracer = new ColdStartTracer(coldStartConfig); + coldStartTracer.trace(requireNodes); + + expect(startSpanSpy).toHaveBeenCalledTimes(4); + expect(startSpanSpy.mock.calls[0][0]).toEqual("aws.lambda.load"); + expect((startSpanSpy.mock.calls[0][1] as any).startTime).toEqual(20); + expect(mockFinishSpan.mock.calls[0][0]).toEqual(220); + + const importSpan = startSpanSpy.mock.calls[2]; + expect(importSpan[0]).toEqual("aws.lambda.import"); + expect(importSpan[1].tags).toEqual({ + operation_name: "aws.lambda.import", + "resource.name": "/var/task/index.mjs", + resource_names: "/var/task/index.mjs", + service: "aws.lambda", + filename: "/var/task/index.mjs", + }); + + const importChildSpan = startSpanSpy.mock.calls[3]; + expect(importChildSpan[0]).toEqual("aws.lambda.require"); + expect((importChildSpan[1] as any).childOf.spanName).toEqual("aws.lambda.import"); + }); }); diff --git a/src/trace/cold-start-tracer.ts b/src/trace/cold-start-tracer.ts index 7f843408..584dff32 100644 --- a/src/trace/cold-start-tracer.ts +++ b/src/trace/cold-start-tracer.ts @@ -32,8 +32,10 @@ export class ColdStartTracer { } trace(rootNodes: RequireNode[]) { - const coldStartSpanStartTime = rootNodes[0]?.startTime; - const coldStartSpanEndTime = Math.min(rootNodes[rootNodes.length - 1]?.endTime, this.currentSpanStartTime); + if (rootNodes.length <= 0) return; + + const coldStartSpanStartTime = Math.min(...rootNodes.map((node) => node.startTime)); + const coldStartSpanEndTime = Math.min(Math.max(...rootNodes.map((node) => node.endTime)), this.currentSpanStartTime); let targetParentSpan: SpanWrapper | undefined; if (this.isColdStart) { const coldStartSpan = this.createColdStartSpan(coldStartSpanStartTime, coldStartSpanEndTime, this.parentSpan); @@ -64,7 +66,12 @@ export class ColdStartTracer { return newSpan; } - private coldStartSpanOperationName(filename: string): string { + private coldStartSpanOperationName(reqNode: RequireNode): string { + const filename = reqNode.filename; + if (reqNode.kind === "import" || filename.startsWith("file://") || filename.endsWith(".mjs")) { + return "aws.lambda.import"; + } + if (filename.startsWith("/opt/")) { return "aws.lambda.require_layer"; } else if (filename.startsWith("/var/runtime/")) { @@ -88,7 +95,7 @@ export class ColdStartTracer { const options: SpanOptions = { tags: { service: "aws.lambda", - operation_name: this.coldStartSpanOperationName(reqNode.filename), + operation_name: this.coldStartSpanOperationName(reqNode), resource_names: reqNode.id, "resource.name": reqNode.id, filename: reqNode.filename, @@ -99,7 +106,7 @@ export class ColdStartTracer { options.childOf = parentSpan.span; } const newSpan = new SpanWrapper( - this.tracerWrapper.startSpan(this.coldStartSpanOperationName(reqNode.filename), options), + this.tracerWrapper.startSpan(this.coldStartSpanOperationName(reqNode), options), {}, ); if (reqNode.endTime - reqNode.startTime > this.minDuration) {