perf(rewriter): defer transformer load to first rewrite

packages/datadog-instrumentations/src/helpers/rewriter/index.js

@@ -4,57 +4,32 @@ const { readFileSync } = require('fs')
 const { join } = require('path')
 const { pathToFileURL } = require('url')
 const log = require('../../../../dd-trace/src/log')
-const { create } = require('../../../../../vendor/dist/@apm-js-collab/code-transformer')
-const {
-  awaitContextCallback,
-  awaitContextCallbackAtTryStart,
-  configureGraphqlJitCompileObject,
-  configureGraphqlJitDeferredField,
-  configureGraphqlJitExecute,
-  configureGraphqlJitRuntime,
-  configureMercuriusRequest,
-  waitForAsyncEnd,
-} = require('./transforms')
 const instrumentations = require('./instrumentations')
 const { getRewriteTarget } = require('./targets')
 
-// `dc-polyfill` is referenced from injected `require()` (CJS) and `import`
-// (ESM) statements that the transformer splices into the rewritten module.
-// `require()` accepts an absolute filesystem path; the ESM resolver rejects it
-// with `ERR_INVALID_MODULE_SPECIFIER` and needs a `file://` URL instead. We
-// pre-compute both forms here so each matcher hands the transformer a
-// specifier that is valid for the module type it is rewriting.
-let dcPolyfillCjs
-let dcPolyfillEsm
-
-try {
-  const resolved = require.resolve('dc-polyfill')
-  dcPolyfillCjs = resolved.replaceAll('\\', '/')
-  dcPolyfillEsm = pathToFileURL(resolved).href
-} catch {
-  // The `dc-polyfill` module is unavailable for some reason (like bundling).
-  // Let's just keep the default of using `diagnostics-channel` as a fallback
-  // which works for most Node versions.
-}
+/**
+ * @typedef {object} InstrumentationMatcher
+ * @property {(name: string, transform: Function) => void} addTransform
+ * @property {(moduleName: string, version: string|undefined, filePath: string) => Transformer|undefined} getTransformer
+ *
+ * @typedef {object} Transformer
+ * @property {(source: string, moduleType: 'cjs'|'esm') => { code: string, map?: string }} transform
+ */
 
 /**
  * @type {Record<string, string>} map of module base name to version
  */
 const moduleVersions = {}
 const disabled = new Set()
-const matcherCjs = create(instrumentations, dcPolyfillCjs)
-const matcherEsm = create(instrumentations, dcPolyfillEsm)
 
-for (const matcher of [matcherCjs, matcherEsm]) {
-  matcher.addTransform('awaitContextCallback', awaitContextCallback)
-  matcher.addTransform('awaitContextCallbackAtTryStart', awaitContextCallbackAtTryStart)
-  matcher.addTransform('waitForAsyncEnd', waitForAsyncEnd)
-  matcher.addTransform('configureGraphqlJitCompileObject', configureGraphqlJitCompileObject)
-  matcher.addTransform('configureGraphqlJitDeferredField', configureGraphqlJitDeferredField)
-  matcher.addTransform('configureGraphqlJitExecute', configureGraphqlJitExecute)
-  matcher.addTransform('configureGraphqlJitRuntime', configureGraphqlJitRuntime)
-  matcher.addTransform('configureMercuriusRequest', configureMercuriusRequest)
-}
+// Matchers are built on the first module that actually needs rewriting. The
+// vendored transformer is a quarter megabyte of bundle that an application
+// without a rewrite target never needs to parse, and an application that loads
+// targets of only one module type never needs the other matcher.
+/** @type {InstrumentationMatcher|undefined} */
+let matcherCjs
+/** @type {InstrumentationMatcher|undefined} */
+let matcherEsm
 
 // Keep the marker split: source-map scanners can read a contiguous token in
 // string literals as this file's own inline map.
@@ -82,8 +57,7 @@ function rewrite (content, filename, format, target) {
 
   if (disabled.has(moduleName)) return content
 
-  const matcher = moduleType === 'esm' ? matcherEsm : matcherCjs
-  const transformer = matcher.getTransformer(moduleName, version, filePath)
+  const transformer = getMatcher(moduleType).getTransformer(moduleName, version, filePath)
 
   if (!transformer) return content
 
@@ -105,6 +79,76 @@ function rewrite (content, filename, format, target) {
   return content
 }
 
+/**
+ * @param {'cjs'|'esm'} moduleType
+ * @returns {InstrumentationMatcher}
+ */
+function getMatcher (moduleType) {
+  if (moduleType === 'esm') {
+    matcherEsm ??= createMatcher(moduleType)
+
+    return matcherEsm
+  }
+
+  matcherCjs ??= createMatcher(moduleType)
+
+  return matcherCjs
+}
+
+/**
+ * @param {'cjs'|'esm'} moduleType
+ * @returns {InstrumentationMatcher}
+ */
+function createMatcher (moduleType) {
+  const { create } = require('../../../../../vendor/dist/@apm-js-collab/code-transformer')
+  const {
+    awaitContextCallback,
+    awaitContextCallbackAtTryStart,
+    configureGraphqlJitCompileObject,
+    configureGraphqlJitDeferredField,
+    configureGraphqlJitExecute,
+    configureGraphqlJitRuntime,
+    configureMercuriusRequest,
+    waitForAsyncEnd,
+  } = require('./transforms')
+
+  const matcher = create(instrumentations, getDcPolyfillSpecifier(moduleType))
+
+  matcher.addTransform('awaitContextCallback', awaitContextCallback)
+  matcher.addTransform('awaitContextCallbackAtTryStart', awaitContextCallbackAtTryStart)
+  matcher.addTransform('waitForAsyncEnd', waitForAsyncEnd)
+  matcher.addTransform('configureGraphqlJitCompileObject', configureGraphqlJitCompileObject)
+  matcher.addTransform('configureGraphqlJitDeferredField', configureGraphqlJitDeferredField)
+  matcher.addTransform('configureGraphqlJitExecute', configureGraphqlJitExecute)
+  matcher.addTransform('configureGraphqlJitRuntime', configureGraphqlJitRuntime)
+  matcher.addTransform('configureMercuriusRequest', configureMercuriusRequest)
+
+  return matcher
+}
+
+/**
+ * `dc-polyfill` is referenced from injected `require()` (CJS) and `import`
+ * (ESM) statements that the transformer splices into the rewritten module.
+ * `require()` accepts an absolute filesystem path; the ESM resolver rejects it
+ * with `ERR_INVALID_MODULE_SPECIFIER` and needs a `file://` URL instead. Each
+ * matcher therefore hands the transformer the form that is valid for the
+ * module type it is rewriting.
+ *
+ * @param {'cjs'|'esm'} moduleType
+ * @returns {string|undefined} `undefined` when `dc-polyfill` cannot be resolved
+ */
+function getDcPolyfillSpecifier (moduleType) {
+  try {
+    const resolved = require.resolve('dc-polyfill')
+
+    return moduleType === 'esm' ? pathToFileURL(resolved).href : resolved.replaceAll('\\', '/')
+  } catch {
+    // The `dc-polyfill` module is unavailable for some reason (like bundling).
+    // Let's just keep the default of using `diagnostics-channel` as a fallback
+    // which works for most Node versions.
+  }
+}
+
 /** @typedef {{ buffer: ArrayBuffer | SharedArrayBuffer, byteLength: number, byteOffset: number }} BufferView */
 
 /**

packages/datadog-instrumentations/test/helpers/rewriter/index.spec.js

@@ -1,6 +1,7 @@
 'use strict'
 
-const { mkdtempSync, readFileSync, writeFileSync } = require('node:fs')
+const { spawnSync } = require('node:child_process')
+const { mkdirSync, mkdtempSync, readFileSync, writeFileSync } = require('node:fs')
 const { tmpdir } = require('node:os')
 const { resolve, join, dirname } = require('node:path')
 const Module = require('node:module')
@@ -1200,3 +1201,61 @@ describe('rewriter source-map trailer', () => {
     }
   })
 })
+
+describe('rewriter initialization', () => {
+  const repositoryRoot = resolve(__dirname, '../../../../..')
+  const transformerPath = join(repositoryRoot, 'vendor', 'dist', '@apm-js-collab', 'code-transformer')
+  const loaderPath = resolve(__dirname, '../../../src/helpers/rewriter/loader')
+
+  it('loads the code transformer on the first rewrite instead of at startup', () => {
+    const root = mkdtempSync(join(tmpdir(), 'dd-rewriter-defer-'))
+    const targetDirectory = join(root, 'node_modules', 'ai')
+    const untargetedDirectory = join(root, 'node_modules', 'untargeted')
+
+    mkdirSync(join(targetDirectory, 'dist'), { recursive: true })
+    writeFileSync(join(targetDirectory, 'package.json'), '{"version":"4.0.0","main":"dist/index.js"}')
+    writeFileSync(join(targetDirectory, 'dist', 'index.js'), `
+      function getTracer () { return 'tracer' }
+      module.exports = { getTracer }
+    `)
+
+    mkdirSync(untargetedDirectory, { recursive: true })
+    writeFileSync(join(untargetedDirectory, 'package.json'), '{"version":"1.0.0"}')
+    writeFileSync(join(untargetedDirectory, 'index.js'), 'module.exports = {}\n')
+
+    writeFileSync(join(root, 'main.js'), `
+      const transformerPath = require.resolve(${JSON.stringify(transformerPath)})
+
+      require(${JSON.stringify(loaderPath)})
+
+      const loadedAfterHook = require.cache[transformerPath] !== undefined
+
+      require('untargeted')
+
+      const loadedAfterUntargetedModule = require.cache[transformerPath] !== undefined
+
+      const { tracingChannel } = require(${JSON.stringify(require.resolve('dc-polyfill'))})
+      let starts = 0
+
+      tracingChannel('orchestrion:ai:getTracer').subscribe({ start () { starts++ } })
+      require('ai').getTracer()
+
+      console.log(JSON.stringify({
+        loadedAfterHook,
+        loadedAfterUntargetedModule,
+        loadedAfterTargetModule: require.cache[transformerPath] !== undefined,
+        starts,
+      }))
+    `)
+
+    const result = spawnSync(process.execPath, [join(root, 'main.js')], { cwd: root, encoding: 'utf8' })
+
+    assert.strictEqual(result.status, 0, result.stderr)
+    assert.deepStrictEqual(JSON.parse(result.stdout), {
+      loadedAfterHook: false,
+      loadedAfterUntargetedModule: false,
+      loadedAfterTargetModule: true,
+      starts: 1,
+    })
+  })
+})