[DO NOT MERGE] Try WASM rewriter again for startup time

LICENSE-3rdparty.csv

@@ -6,6 +6,7 @@ require,@datadog/native-iast-taint-tracking,Apache license 2.0,Copyright 2018 Da
 require,@datadog/openfeature-node-server,Apache license 2.0,Copyright 2024 Datadog Inc.
 require,@datadog/pprof,Apache license 2.0,Copyright 2019 Google Inc.
 require,@datadog/sketches-js,Apache license 2.0,Copyright 2020 Datadog Inc.
+require,@datadog/source-map-support,MIT,Copyright 2014 Evan Wallace
 require,@datadog/wasm-js-rewriter,Apache license 2.0,Copyright 2018 Datadog Inc.
 require,@opentelemetry/api,Apache license 2.0,Copyright OpenTelemetry Authors
 require,@opentelemetry/api-logs,Apache license 2.0,Copyright OpenTelemetry Authors
@@ -15,6 +16,8 @@ require,@isaacs/ttlcache,Blue Oak,Copyright Isaac Z. Schlueter and Contributors
 require,crypto-randomuuid,MIT,Copyright 2021 Node.js Foundation and contributors
 require,dc-polyfill,MIT,Copyright 2023 Datadog Inc.
 require,escape-string-regexp,MIT,Copyright Sindre Sorhus
+require,escodegen,BSD-2-Clause,Copyright 2012 Yusuke Suzuki and other contributors
+require,esquery,MIT,Copyright 2013 Joel Feenstra
 require,ignore,MIT,Copyright 2013 Kael Zhang and contributors
 require,import-in-the-middle,Apache license 2.0,Copyright 2021 Datadog Inc.
 require,istanbul-lib-coverage,BSD-3-Clause,Copyright 2012-2015 Yahoo! Inc.
@@ -23,6 +26,7 @@ require,jsonpath-plus,MIT,Copyright (c) 2011-2019 Stefan Goessner, Subbu Allamar
 require,limiter,MIT,Copyright 2011 John Hurliman
 require,lodash.sortby,MIT,Copyright JS Foundation and other contributors
 require,lru-cache,ISC,Copyright (c) 2010-2022 Isaac Z. Schlueter and Contributors
+require,meriyah,ISC,Copyright 2019 KFlash and others
 require,module-details-from-path,MIT,Copyright 2016 Thomas Watson Steen
 require,mutexify,MIT,Copyright (c) 2014 Mathias Buus
 require,opentracing,MIT,Copyright 2016 Resonance Labs Inc

loader-hook.mjs

@@ -2,6 +2,7 @@ import regexpEscape from 'escape-string-regexp'
 import * as iitm from 'import-in-the-middle/hook.mjs'
 import hooks from './packages/datadog-instrumentations/src/helpers/hooks.js'
 import configHelper from './packages/dd-trace/src/config-helper.js'
+import * as rewriterLoader from './packages/datadog-instrumentations/src/helpers/rewriter/loader.mjs'
 
 // For some reason `getEnvironmentVariable` is not otherwise available to ESM.
 const env = configHelper.getEnvironmentVariable
@@ -17,6 +18,10 @@ function initialize (data = {}) {
   return iitm.initialize(data)
 }
 
+function load (url, context, nextLoad) {
+  return rewriterLoader.load(url, context, (url, context) => iitm.load(url, context, nextLoad))
+}
+
 function addInstrumentations (data) {
   const instrumentations = Object.keys(hooks)
 
@@ -48,5 +53,5 @@ function addExclusions (data) {
   )
 }
 
-export { initialize }
-export { load, getFormat, resolve, getSource } from 'import-in-the-middle/hook.mjs'
+export { initialize, load }
+export { getFormat, resolve, getSource } from 'import-in-the-middle/hook.mjs'

package.json

@@ -128,6 +128,7 @@
     "@datadog/openfeature-node-server": "0.1.0-preview.15",
     "@datadog/pprof": "5.12.0",
     "@datadog/sketches-js": "2.1.1",
+    "@datadog/source-map-support": "npm:source-map-support@^0.5.21",
     "@datadog/wasm-js-rewriter": "4.0.1",
     "@isaacs/ttlcache": "^2.0.1",
     "@opentelemetry/api": ">=1.0.0 <1.10.0",
@@ -137,6 +138,8 @@
     "crypto-randomuuid": "^1.0.0",
     "dc-polyfill": "^0.1.10",
     "escape-string-regexp": "^5.0.0",
+    "escodegen": "^2.1.0",
+    "esquery": "^1.6.0",
     "ignore": "^7.0.5",
     "import-in-the-middle": "^1.14.2",
     "istanbul-lib-coverage": "^3.2.2",
@@ -145,6 +148,7 @@
     "limiter": "^1.1.5",
     "lodash.sortby": "^4.7.0",
     "lru-cache": "^10.4.3",
+    "meriyah": "^6.1.4",
     "module-details-from-path": "^1.0.4",
     "mutexify": "^1.4.0",
     "opentracing": ">=0.14.7",

packages/datadog-instrumentations/index.js

@@ -2,3 +2,4 @@
 
 require('./src/helpers/bundler-register')
 require('./src/helpers/register')
+require('./src/helpers/rewriter/loader')

packages/datadog-instrumentations/src/helpers/instrument.js

@@ -2,6 +2,7 @@
 
 const dc = require('dc-polyfill')
 const instrumentations = require('./instrumentations')
+const rewriterInstrumentations = require('./rewriter/instrumentations.json')
 const { AsyncResource } = require('async_hooks')
 
 const channelMap = {}
@@ -22,6 +23,14 @@ exports.tracingChannel = function (name) {
   return tc
 }
 
+exports.getHooks = function getHooks (names) {
+  names = [names].flat()
+
+  return rewriterInstrumentations
+    .filter(inst => names.includes(inst.moduleName))
+    .map(inst => ({ name: inst.moduleName, versions: [inst.versionRange], file: inst.filePath }))
+}
+
 /**
  * @param {object} args
  * @param {string|string[]} args.name module name

packages/datadog-instrumentations/src/helpers/register.js

@@ -10,6 +10,7 @@ const checkRequireCache = require('./check-require-cache')
 const telemetry = require('../../../dd-trace/src/guardrails/telemetry')
 const { isInServerlessEnvironment } = require('../../../dd-trace/src/serverless')
 const { getEnvironmentVariables } = require('../../../dd-trace/src/config-helper')
+const rewriter = require('./rewriter')
 
 const envs = getEnvironmentVariables()
 
@@ -48,6 +49,10 @@ if (DD_TRACE_DEBUG && DD_TRACE_DEBUG.toLowerCase() !== 'false') {
 const seenCombo = new Set()
 const allInstrumentations = {}
 
+for (const inst of disabledInstrumentations) {
+  rewriter.disable(inst)
+}
+
 // TODO: make this more efficient
 for (const packageName of names) {
   if (disabledInstrumentations.has(packageName)) continue

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

@@ -0,0 +1,149 @@
+'use strict'
+
+// The rewriter works effectively the same as Orchestrion with some additions:
+// - Supports an `astQuery` field to filter AST nodes with an esquery query.
+
+const { readFileSync } = require('fs')
+const { join } = require('path')
+const semifies = require('semifies')
+const transforms = require('./transforms')
+const log = require('../../../../dd-trace/src/log')
+const instrumentations = require('./instrumentations.json')
+
+const supported = {}
+const disabled = new Set()
+const sourceMaps = {}
+
+let parse
+let generate
+let esquery
+let sourceMapSupport
+
+const { Rewriter } = require('@datadog/wasm-js-rewriter')
+const rewriter = new Rewriter({
+  orchestrion: `
+version: 1
+dc_module: dc-polyfill
+instrumentations: []
+`
+})
+
+function rewrite (content, filename, format) {
+  if (!content) return content
+
+  try {
+    let ast
+
+    filename = filename.replace('file://', '')
+
+    for (const inst of instrumentations) {
+      const { astQuery, moduleName, versionRange, filePath, channelName } = inst
+      const transform = transforms[inst.operator]
+
+      if (disabled.has(moduleName)) continue
+      if (!filename.endsWith(`${moduleName}/${filePath}`)) continue
+      if (!transform) continue
+      if (!satisfies(filename, filePath, versionRange)) continue
+
+      parse ??= require('meriyah').parse
+      generate ??= require('escodegen').generate
+      esquery ??= require('esquery')
+
+      if (!sourceMapSupport) {
+        // Use an alias to ensure we have our own instance, otherwise there
+        // could be an existing user instance and the library doesn't support
+        // multiple calls to `install`.
+        sourceMapSupport = require('@datadog/source-map-support')
+        sourceMapSupport.install({
+          retrieveSourceMap: function (url) {
+            const map = sourceMaps[url]
+            return map ? { url, map } : null
+          }
+        })
+      }
+
+      try {
+        ast ??= parse(content.toString(), { loc: true, ranges: true, module: format === 'module' })
+      } catch (e) {
+        log.error(e)
+      }
+
+      const query = astQuery || functionQuery(inst)
+      const selector = esquery.parse(query)
+      const state = { channelName, format, moduleName }
+
+      esquery.traverse(ast, selector, (...args) => transform(state, ...args))
+    }
+
+    if (ast) {
+      const { code, map } = generate(ast, { sourceMap: filename, sourceMapWithCode: true })
+
+      sourceMaps[filename] = map.toString()
+
+      return code
+    }
+  } catch (e) {
+    log.error(e)
+  }
+
+  return content
+}
+
+function disable (instrumentation) {
+  disabled.add(instrumentation)
+}
+
+function satisfies (filename, filePath, versions) {
+  const [basename] = filename.split(filePath)
+
+  if (supported[basename] === undefined) {
+    try {
+      const pkg = JSON.parse(readFileSync(
+        join(basename, 'package.json'), 'utf8'
+      ))
+
+      supported[basename] = semifies(pkg.version, versions)
+    } catch {
+      supported[basename] = false
+    }
+  }
+
+  return supported[basename]
+}
+
+// TODO: Support index
+function functionQuery (inst) {
+  const { functionQuery } = inst
+  const { methodName, functionName, expressionName, className } = functionQuery
+  const kind = functionQuery.kind?.toLowerCase()
+
+  let queries = []
+
+  if (className) {
+    queries.push(
+      `[id.name="${className}"] > ClassBody > [key.name="${methodName}"] > [async]`
+    )
+  } else if (methodName) {
+    queries.push(
+      `ClassBody > [key.name="${methodName}"] > [async]`,
+      `Property[key.name="${methodName}"] > [async]`
+    )
+  }
+
+  if (functionName) {
+    queries.push(`FunctionDeclaration[id.name="${functionName}"]`)
+  } else if (expressionName) {
+    queries.push(
+      `FunctionExpression[id.name="${expressionName}"]`,
+      `ArrowFunctionExpression[id.name="${expressionName}"]`
+    )
+  }
+
+  if (kind) {
+    queries = queries.map(q => `${q}[async="${kind === 'async' ? 'true' : 'false'}"]`)
+  }
+
+  return queries.join(', ')
+}
+
+module.exports = { rewrite, disable }