feat: esm cold start tracing visibility

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

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"

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,12 +25,109 @@ 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)
     const resChild = resChildren.pop() as RequireNode
     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')
+  })
 })

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,18 +55,83 @@ 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) {
     rootNodes.push(reqNode)
   }
 }
 
+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 || "<unknown>")
+}
+
+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
 }

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