switch to js implementation of orchestrion

LICENSE-3rdparty.csv

@@ -7,6 +7,7 @@
 "@datadog/openfeature-node-server","https://github.com/DataDog/openfeature-js-client","['Apache-2.0']","['DataDog']"
 "@datadog/pprof","https://github.com/DataDog/pprof-nodejs","['Apache-2.0']","['Google Inc.']"
 "@datadog/sketches-js","https://github.com/DataDog/sketches-js","['Apache-2.0']","['DataDog']"
+"@datadog/source-map","npm:@datadog/source-map","[]","[]"
 "@datadog/wasm-js-rewriter","https://github.com/DataDog/dd-wasm-js-rewriter","['Apache-2.0']","['Datadog Inc.']"
 "@isaacs/ttlcache","https://github.com/isaacs/ttlcache","['BlueOak-1.0.0']","['Isaac Z. Schlueter']"
 "@jsep-plugin/assignment","https://github.com/EricSmekens/jsep","['MIT']","['Shelly']"
@@ -31,13 +32,16 @@
 "acorn","https://github.com/acornjs/acorn","['MIT']","['acornjs']"
 "acorn-import-attributes","https://github.com/xtuc/acorn-import-attributes","['MIT']","['Sven Sauleau']"
 "argparse","https://github.com/nodeca/argparse","['Python-2.0']","['nodeca']"
+"astring","https://github.com/davidbonnet/astring","['MIT']","['David Bonnet']"
 "cjs-module-lexer","https://github.com/nodejs/cjs-module-lexer","['MIT']","['Guy Bedford']"
 "crypto-randomuuid","npm:crypto-randomuuid","['MIT']","['Stephen Belanger']"
 "dc-polyfill","https://github.com/DataDog/dc-polyfill","['MIT']","['Thomas Hunter II']"
 "dd-trace","https://github.com/DataDog/dd-trace-js","['(Apache-2.0 OR BSD-3-Clause)']","['Datadog Inc. <info@datadoghq.com>']"
 "delay","https://github.com/sindresorhus/delay","['MIT']","['Sindre Sorhus']"
 "detect-newline","https://github.com/sindresorhus/detect-newline","['MIT']","['Sindre Sorhus']"
 "escape-string-regexp","https://github.com/sindresorhus/escape-string-regexp","['MIT']","['Sindre Sorhus']"
+"esquery","https://github.com/estools/esquery","['BSD-3-Clause']","['Joel Feenstra']"
+"estraverse","https://github.com/estools/estraverse","['BSD-2-Clause']","['estools']"
 "fast-fifo","https://github.com/mafintosh/fast-fifo","['MIT']","['Mathias Buus']"
 "ignore","https://github.com/kaelzhang/node-ignore","['MIT']","['kael']"
 "import-in-the-middle","https://github.com/nodejs/import-in-the-middle","['Apache-2.0']","['Bryan English']"
@@ -50,6 +54,7 @@
 "lodash.sortby","https://github.com/lodash/lodash","['MIT']","['John-David Dalton']"
 "long","https://github.com/dcodeIO/long.js","['Apache-2.0']","['Daniel Wirtz']"
 "lru-cache","https://github.com/isaacs/node-lru-cache","['ISC']","['Isaac Z. Schlueter']"
+"meriyah","https://github.com/meriyah/meriyah","['ISC']","['Kenny F.']"
 "module-details-from-path","https://github.com/watson/module-details-from-path","['MIT']","['Thomas Watson']"
 "mutexify","https://github.com/mafintosh/mutexify","['MIT']","['Mathias Buus']"
 "node-addon-api","https://github.com/nodejs/node-addon-api","['MIT']","['nodejs']"

loader-hook.mjs

@@ -2,6 +2,7 @@ import regexpEscapeModule from './packages/node_modules/escape-string-regexp/ind
 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'
 
 const regexpEscape = regexpEscapeModule.default
 
@@ -19,6 +20,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)
 
@@ -50,5 +55,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'

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')
 const { AsyncResource } = require('async_hooks')
 
 const channelMap = {}
@@ -22,6 +23,15 @@ exports.tracingChannel = function (name) {
   return tc
 }
 
+exports.getHooks = function getHooks (names) {
+  names = [names].flat()
+
+  return rewriterInstrumentations
+    .map(inst => inst.module)
+    .filter(({ name }) => names.includes(name))
+    .map(({ name, versionRange, filePath }) => ({ name, versions: [versionRange], file: 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/compiler.js

@@ -0,0 +1,27 @@
+'use strict'
+
+let meriyah
+let astring
+let esquery
+
+module.exports = {
+  parse: (...args) => {
+    meriyah ??= require('meriyah')
+
+    return meriyah.parse(...args)
+  },
+
+  generate: (...args) => {
+    astring ??= require('astring')
+
+    return astring.generate(...args)
+  },
+
+  traverse: (ast, query, visitor) => {
+    esquery ??= require('esquery').default
+
+    const selector = esquery.parse(query)
+
+    return esquery.traverse(ast, selector, visitor)
+  },
+}

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

@@ -0,0 +1,152 @@
+'use strict'
+
+/*
+This rewriter is basically a JavaScript version of Orchestrion-JS. The goal is
+not to replace Orchestrion-JS, but rather to make it easier and faster to write
+new integrations in the short-term, especially as many changes to the rewriter
+will be needed as all the patterns we need have not been identified yet. This
+will avoid the back and forth of having to make Rust changes to an external
+library for every integration change or addition that requires something new.
+
+In the meantime, we'll work concurrently on a change to Orchestrion-JS that
+adds an "arbitrary transform" or "plugin" system that can be used from
+JavaScript, in order to enable quick iteration while still using Orchestrion-JS.
+Once that's done we'll use that, so that we can remove this JS approach and
+return to using Orchestrion-JS.
+
+The long term goal is to backport any additional features we add to the JS
+rewriter (or using the plugin system in Orchestrion-JS once we're using that)
+to Orchestrion-JS  once we're confident that the implementation is fairly
+complete and has all features we need.
+
+Here is a list of the additions and changes in this rewriter compared to
+Orchestrion-JS that will need to be backported:
+
+(NOTE: Please keep this list up-to-date whenever new features are added)
+
+- Supports an `astQuery` field to filter AST nodes with an esquery query. This
+  is mostly meant to be used when experimenting or if what needs to be queried
+  is not a function. We'll see over time if something like this is needed to be
+  backported or if it can be replaced by simpler queries.
+- Supports replacing methods of child class instances in the base constructor.
+*/
+
+const { readFileSync } = require('fs')
+const { join } = require('path')
+const semifies = require('semifies')
+const transforms = require('./transforms')
+const { generate, parse, traverse } = require('./compiler')
+const log = require('../../../../dd-trace/src/log')
+const instrumentations = require('./instrumentations')
+const { getEnvironmentVariable } = require('../../../../dd-trace/src/config-helper')
+
+const NODE_OPTIONS = getEnvironmentVariable('NODE_OPTIONS')
+
+const supported = {}
+const disabled = new Set()
+
+// TODO: Source maps without `--enable-source-maps`.
+const enableSourceMaps = NODE_OPTIONS?.includes('--enable-source-maps') ||
+  process.execArgv?.some(arg => arg.includes('--enable-source-maps'))
+
+let SourceMapGenerator
+
+function rewrite (content, filename, format) {
+  if (!content) return content
+
+  try {
+    let ast
+
+    filename = filename.replace('file://', '')
+
+    for (const inst of instrumentations) {
+      const { astQuery, functionQuery = {}, module: { name, versionRange, filePath } } = inst
+      const { kind } = functionQuery
+      const operator = kind === 'Async' ? 'tracePromise' : kind === 'Callback' ? 'traceCallback' : 'traceSync'
+      const transform = transforms[operator]
+
+      if (disabled.has(name)) continue
+      if (!filename.endsWith(`${name}/${filePath}`)) continue
+      if (!transform) continue
+      if (!satisfies(filename, filePath, versionRange)) continue
+
+      ast ??= parse(content.toString(), { loc: true, ranges: true, module: format === 'module' })
+
+      const query = astQuery || fromFunctionQuery(functionQuery)
+      const state = { ...inst, format, functionQuery, operator }
+
+      traverse(ast, query, (...args) => transform(state, ...args))
+    }
+
+    if (ast) {
+      if (!enableSourceMaps) return generate(ast)
+
+      // TODO: Can we use the same version of `source-map` that DI uses?
+      SourceMapGenerator ??= require('@datadog/source-map').SourceMapGenerator
+
+      const sourceMap = new SourceMapGenerator({ file: filename })
+      const code = generate(ast, { sourceMap })
+      const map = Buffer.from(sourceMap.toString()).toString('base64')
+
+      return code + '\n' + `//# sourceMappingURL=data:application/json;base64,${map}`
+    }
+  } 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 fromFunctionQuery (functionQuery) {
+  const { methodName, functionName, expressionName, className } = functionQuery
+  const queries = []
+
+  if (className) {
+    queries.push(
+      `[id.name="${className}"]`,
+      `[id.name="${className}"] > ClassBody > [key.name="${methodName}"] > [async]`,
+      `[id.name="${className}"] > ClassExpression > 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}"][async]`)
+  } else if (expressionName) {
+    queries.push(
+      `FunctionExpression[id.name="${expressionName}"][async]`,
+      `ArrowFunctionExpression[id.name="${expressionName}"][async]`
+    )
+  }
+
+  return queries.join(', ')
+}
+
+module.exports = { rewrite, disable }

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

@@ -0,0 +1,5 @@
+'use strict'
+
+module.exports = [
+  ...require('./langchain')
+]