dd-trace 5.87.0 → 5.89.0

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/bullmq.js

@@ -0,0 +1,108 @@
+'use strict'
+
+module.exports = [
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/cjs/classes/queue.js',
+    },
+    functionQuery: {
+      methodName: 'add',
+      className: 'Queue',
+      kind: 'Async',
+    },
+    channelName: 'Queue_add',
+  },
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/cjs/classes/queue.js',
+    },
+    functionQuery: {
+      methodName: 'addBulk',
+      className: 'Queue',
+      kind: 'Async',
+    },
+    channelName: 'Queue_addBulk',
+  },
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/cjs/classes/worker.js',
+    },
+    functionQuery: {
+      methodName: 'callProcessJob',
+      className: 'Worker',
+      kind: 'Async',
+    },
+    channelName: 'Worker_callProcessJob',
+  },
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/cjs/classes/flow-producer.js',
+    },
+    functionQuery: {
+      methodName: 'add',
+      className: 'FlowProducer',
+      kind: 'Async',
+    },
+    channelName: 'FlowProducer_add',
+  },
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/esm/classes/queue.js',
+    },
+    functionQuery: {
+      methodName: 'add',
+      className: 'Queue',
+      kind: 'Async',
+    },
+    channelName: 'Queue_add',
+  },
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/esm/classes/queue.js',
+    },
+    functionQuery: {
+      methodName: 'addBulk',
+      className: 'Queue',
+      kind: 'Async',
+    },
+    channelName: 'Queue_addBulk',
+  },
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/esm/classes/worker.js',
+    },
+    functionQuery: {
+      methodName: 'callProcessJob',
+      className: 'Worker',
+      kind: 'Async',
+    },
+    channelName: 'Worker_callProcessJob',
+  },
+  {
+    module: {
+      name: 'bullmq',
+      versionRange: '>=5.66.0',
+      filePath: 'dist/esm/classes/flow-producer.js',
+    },
+    functionQuery: {
+      methodName: 'add',
+      className: 'FlowProducer',
+      kind: 'Async',
+    },
+    channelName: 'FlowProducer_add',
+  },
+]

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

@@ -1,6 +1,7 @@
 'use strict'
 
 module.exports = [
-  ...require('./langchain'),
+  ...require('./ai'),
   ...require('./bullmq'),
+  ...require('./langchain'),
 ]

package/packages/datadog-instrumentations/src/helpers/rewriter/orchestrion/compiler.js

@@ -0,0 +1,74 @@
+'use strict'
+
+const log = require('../../../../../dd-trace/src/log')
+
+// eslint-disable-next-line camelcase, no-undef
+const runtimeRequire = typeof __webpack_require__ === 'function' ? __non_webpack_require__ : require
+
+const compiler = {
+  parse: (sourceText, options) => {
+    try {
+      // TODO: Figure out ESBuild `createRequire` issue and remove this hack.
+      const oxc = runtimeRequire(['oxc', 'parser'].join('-'))
+
+      compiler.parse = (sourceText, options) => {
+        const { program, errors } = oxc.parseSync('index.js', sourceText, {
+          ...options,
+          preserveParens: false,
+        })
+
+        if (errors?.length > 0) throw errors[0]
+
+        return program
+      }
+    } catch (e) {
+      log.error(e)
+
+      // Fallback for when OXC is not available.
+      const meriyah = require('../../../../../../vendor/dist/meriyah')
+
+      compiler.parse = (sourceText, { range, sourceType } = {}) => {
+        return meriyah.parse(sourceText.toString(), {
+          loc: range,
+          ranges: range,
+          module: sourceType === 'module',
+        })
+      }
+    }
+
+    return compiler.parse(sourceText, options)
+  },
+
+  generate: (...args) => {
+    const astring = require('../../../../../../vendor/dist/astring')
+
+    compiler.generate = astring.generate
+
+    return compiler.generate(...args)
+  },
+
+  traverse: (ast, query, visitor) => {
+    const esquery = require('../../../../../../vendor/dist/esquery').default
+
+    compiler.traverse = (ast, query, visitor) => {
+      return esquery.traverse(ast, esquery.parse(query), visitor)
+    }
+
+    return compiler.traverse(ast, query, visitor)
+  },
+
+  query: (ast, query) => {
+    const esquery = require('../../../../../../vendor/dist/esquery').default
+
+    compiler.query = esquery.query
+
+    return compiler.query(ast, query)
+  },
+}
+
+module.exports = {
+  parse: (...args) => compiler.parse(...args),
+  generate: (...args) => compiler.generate(...args),
+  traverse: (...args) => compiler.traverse(...args),
+  query: (...args) => compiler.query(...args),
+}

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

@@ -0,0 +1,43 @@
+'use strict'
+
+/*
+This folder 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.
+- Supports tracing iterator (sync/async) returning functions (sync/async).
+*/
+
+/* eslint-disable camelcase */
+
+const { InstrumentationMatcher } = require('./matcher')
+
+function create (configs, dc_module) {
+  return new InstrumentationMatcher(configs, dc_module)
+}
+
+module.exports = { create }

package/packages/datadog-instrumentations/src/helpers/rewriter/orchestrion/matcher.js

@@ -0,0 +1,49 @@
+'use strict'
+
+/* eslint-disable camelcase */
+
+const semifies = require('../../../../../../vendor/dist/semifies')
+const { Transformer } = require('./transformer')
+
+// TODO: addTransform
+
+class InstrumentationMatcher {
+  #configs = []
+  #dc_module = null
+  #transformers = {}
+
+  constructor (configs, dc_module) {
+    this.#configs = configs
+    this.#dc_module = dc_module || 'diagnostics_channel'
+  }
+
+  free () {
+    this.#transformers = {}
+  }
+
+  getTransformer (module_name, version, file_path) {
+    const id = `${module_name}/${file_path}@${version}`
+
+    if (this.#transformers[id]) return this.#transformers[id]
+
+    const configs = this.#configs.filter(({ module: { name, filePath, versionRange } }) =>
+      name === module_name &&
+      filePath === file_path &&
+      semifies(version, versionRange)
+    )
+
+    if (configs.length === 0) return
+
+    this.#transformers[id] = new Transformer(
+      module_name,
+      version,
+      file_path,
+      configs,
+      this.#dc_module
+    )
+
+    return this.#transformers[id]
+  }
+}
+
+module.exports = { InstrumentationMatcher }

package/packages/datadog-instrumentations/src/helpers/rewriter/orchestrion/transformer.js

@@ -0,0 +1,121 @@
+'use strict'
+
+/* eslint-disable camelcase */
+
+const { generate, parse, traverse } = require('./compiler')
+const transforms = require('./transforms')
+
+let SourceMapConsumer
+let SourceMapGenerator
+
+class Transformer {
+  #module_name = null
+  #file_path = null
+  #configs = []
+  #dc_module = null
+
+  // TODO: module_name false for user module
+  constructor (module_name, _version, file_path, configs, dc_module) {
+    this.#module_name = module_name
+    this.#file_path = file_path
+    this.#configs = configs
+    this.#dc_module = dc_module
+  }
+
+  free () {
+    // Freeing is not needed for a JavaScript implementation.
+  }
+
+  transform (code, module_type, sourcemap) {
+    if (!code) return { code }
+
+    const sourceType = module_type === 'esm' ? 'module' : 'script'
+
+    let ast
+
+    for (const config of this.#configs) {
+      const { astQuery, functionQuery = {} } = config
+
+      ast ??= parse(code.toString(), { range: true, sourceType })
+
+      const query = astQuery || this.#fromFunctionQuery(functionQuery)
+      const state = { ...config, dcModule: this.#dc_module, sourceType, functionQuery }
+
+      state.operator = this.#getOperator(state)
+
+      traverse(ast, query, (...args) => this.#visit(state, ...args))
+    }
+
+    if (ast) {
+      SourceMapConsumer ??= require('../../../../../../vendor/dist/@datadog/source-map').SourceMapConsumer
+      SourceMapGenerator ??= require('../../../../../../vendor/dist/@datadog/source-map').SourceMapGenerator
+
+      const file = `${this.#module_name}/${this.#file_path}`
+      const sourceMapInput = sourcemap ? new SourceMapConsumer(sourcemap) : { file }
+      const sourceMap = new SourceMapGenerator(sourceMapInput)
+      const code = generate(ast, { sourceMap })
+      const map = sourceMap.toString()
+
+      return { code, map }
+    }
+
+    return { code }
+  }
+
+  #visit (state, ...args) {
+    const transform = transforms[state.operator]
+    const { index } = state.functionQuery
+
+    if (index !== undefined) {
+      state.functionIndex = ++state.functionIndex || 0
+
+      if (index !== state.functionIndex) return
+    }
+
+    transform(state, ...args)
+  }
+
+  #getOperator ({ functionQuery: { kind } }) {
+    switch (kind) {
+      case 'Async': return 'tracePromise'
+      case 'AsyncIterator': return 'traceAsyncIterator'
+      case 'Callback': return 'traceCallback'
+      case 'Iterator': return 'traceIterator'
+      case 'Sync': return 'traceSync'
+    }
+  }
+
+  #fromFunctionQuery (functionQuery) {
+    const { functionName, expressionName, className } = functionQuery
+    const method = functionQuery.methodName || functionQuery.privateMethodName
+    const type = functionQuery.privateMethodName ? 'PrivateIdentifier' : 'Identifier'
+    const queries = []
+
+    if (className) {
+      queries.push(
+        `[id.name="${className}"]`,
+        `[id.name="${className}"] > ClassExpression`,
+        `[id.name="${className}"] > ClassBody > [key.name="${method}"][key.type=${type}] > [async]`,
+        `[id.name="${className}"] > ClassExpression > ClassBody > [key.name="${method}"][key.type=${type}] > [async]`
+      )
+    } else if (method) {
+      queries.push(
+        `ClassBody > [key.name="${method}"][key.type=${type}] > [async]`,
+        `Property[key.name="${method}"][key.type=${type}] > [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 = { Transformer }

package/packages/datadog-instrumentations/src/helpers/rewriter/{transforms.js → orchestrion/transforms.js}

@@ -1,6 +1,6 @@
 'use strict'
 
-const { parse, query } = require('./compiler')
+const { parse, query, traverse } = require('./compiler')
 
 const tracingChannelPredicate = (node) => (
   node.specifiers?.[0]?.local?.name === 'tr_ch_apm_tracingChannel' ||
@@ -8,15 +8,15 @@ const tracingChannelPredicate = (node) => (
 )
 
 const transforms = module.exports = {
-  tracingChannelImport ({ format }, node) {
+  tracingChannelImport ({ dcModule, sourceType }, node) {
     if (node.body.some(tracingChannelPredicate)) return
 
     const index = node.body.findIndex(child => child.directive === 'use strict')
-    const code = format === 'module'
-      ? 'import { tracingChannel as tr_ch_apm_tracingChannel } from "diagnostics_channel"'
-      : 'const {tracingChannel: tr_ch_apm_tracingChannel} = require("diagnostics_channel")'
+    const code = sourceType === 'module'
+      ? `import { tracingChannel as tr_ch_apm_tracingChannel } from "${dcModule}"`
+      : `const {tracingChannel: tr_ch_apm_tracingChannel} = require("${dcModule}")`
 
-    node.body.splice(index + 1, 0, parse(code, { module: format === 'module' }).body[0])
+    node.body.splice(index + 1, 0, parse(code, { sourceType }).body[0])
   },
 
   tracingChannelDeclaration (state, node) {
@@ -35,7 +35,9 @@ const transforms = module.exports = {
     node.body.splice(index + 1, 0, parse(code).body[0])
   },
 
+  traceAsyncIterator: traceAny,
   traceCallback: traceAny,
+  traceIterator: traceAny,
   tracePromise: traceAny,
   traceSync: traceAny,
 }
@@ -51,18 +53,25 @@ function traceAny (state, node, _parent, ancestry) {
 }
 
 function traceFunction (state, node, program) {
-  const { operator } = state
-
   transforms.tracingChannelDeclaration(state, program)
 
   node.body = wrap(state, {
-    type: 'ArrowFunctionExpression',
+    type: 'FunctionExpression',
     params: node.params,
     body: node.body,
-    async: operator === 'tracePromise',
+    async: node.async,
     expression: false,
-    generator: false,
-  })
+    generator: node.generator,
+  }, program)
+
+  // The original function no longer contains any calls to `await` or `yield` as
+  // the function body is copied to the internal wrapped function, so we set
+  // these to false to avoid altering the return value of the wrapper. The old
+  // values are instead copied to the new AST node above.
+  node.generator = false
+  node.async = false
+
+  wrapSuper(state, node)
 }
 
 function traceInstanceMethod (state, node, program) {
@@ -100,15 +109,19 @@ function traceInstanceMethod (state, node, program) {
   const fn = ctorBody[1].expression.right
 
   fn.async = operator === 'tracePromise'
-  fn.body = wrap(state, { type: 'Identifier', name: `__apm$${methodName}` })
+  fn.body = wrap(state, { type: 'Identifier', name: `__apm$${methodName}` }, program)
+
+  wrapSuper(state, fn)
 
   ctor.value.body.body.push(...ctorBody)
 }
 
-function wrap (state, node) {
+function wrap (state, node, program) {
   const { channelName, operator } = state
 
+  if (operator === 'traceAsyncIterator') return wrapIterator(state, node, program)
   if (operator === 'traceCallback') return wrapCallback(state, node)
+  if (operator === 'traceIterator') return wrapIterator(state, node, program)
 
   const async = operator === 'tracePromise' ? 'async' : ''
   const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
@@ -133,12 +146,61 @@ function wrap (state, node) {
   return wrapper
 }
 
+function wrapSuper (_state, node) {
+  const members = new Set()
+
+  traverse(
+    node.body,
+    '[object.type=Super]',
+    (node, parent) => {
+      const { name } = node.property
+
+      let child
+
+      if (parent.callee) {
+        // This is needed because for generator functions we have to move the
+        // original function to a nested wrapped function, but we can't use an
+        // arrow function because arrow function cannot be generator functions,
+        // and `super` cannot be called from a nested function, so we have to
+        // rewrite any `super` call to not use the keyword.
+        const { expression } = parse(`__apm$super['${name}'].call(this)`).body[0]
+
+        parent.callee = child = expression.callee
+        parent.arguments.unshift(...expression.arguments)
+      } else {
+        parent.expression = child = parse(`__apm$super['${name}']`).body[0]
+      }
+
+      child.computed = parent.callee.computed
+      child.optional = parent.callee.optional
+
+      members.add(name)
+    }
+  )
+
+  for (const name of members) {
+    const member = parse(`
+      class Wrapper {
+        wrapper () {
+          __apm$super['${name}'] = super['${name}']
+        }
+      }
+    `).body[0].body.body[0].value.body.body[0]
+
+    node.body.body.unshift(member)
+  }
+
+  if (members.size > 0) {
+    node.body.body.unshift(parse('const __apm$super = {}').body[0])
+  }
+}
+
 function wrapCallback (state, node) {
-  const { channelName, functionQuery: { index = -1 } } = state
+  const { channelName, functionQuery: { callbackIndex = -1 } } = state
   const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
   const wrapper = parse(`
     function wrapper () {
-      const __apm$cb = Array.prototype.at.call(arguments, ${index});
+      const __apm$cb = Array.prototype.at.call(arguments, ${callbackIndex});
       const __apm$ctx = {
         arguments,
         self: this,
@@ -173,7 +235,7 @@ function wrapCallback (state, node) {
       if (typeof __apm$cb !== 'function') {
         return __apm$traced();
       }
-      Array.prototype.splice.call(arguments, ${index}, 1, __apm$wrappedCb);
+      Array.prototype.splice.call(arguments, ${callbackIndex}, 1, __apm$wrappedCb);
 
       return ${channelVariable}.start.runStores(__apm$ctx, () => {
         try {
@@ -194,3 +256,67 @@ function wrapCallback (state, node) {
 
   return wrapper
 }
+
+function wrapIterator (state, node, program) {
+  const { channelName, operator } = state
+  const baseChannel = channelName.replaceAll(':', '_')
+  const channelVariable = 'tr_ch_apm$' + baseChannel
+  const nextChannel = baseChannel + '_next'
+  const traceMethod = operator === 'traceAsyncIterator' ? 'tracePromise' : 'traceSync'
+  const traceNext = `tr_ch_apm$${nextChannel}.${traceMethod}`
+
+  transforms.tracingChannelDeclaration({ ...state, channelName: nextChannel }, program)
+
+  const wrapper = parse(`
+    function wrapper () {
+      const __apm$traced = () => {
+        const __apm$wrapped = () => {};
+        return __apm$wrapped.apply(this, arguments);
+      };
+
+      if (!${channelVariable}.start.hasSubscribers) return __apm$traced();
+
+      {
+        const wrap = iter => {
+          const { next: iterNext, return: iterReturn, throw: iterThrow } = iter;
+
+          iter.next = (...args) => ${traceNext}(iterNext, ctx, iter, ...args);
+          iter.return = (...args) => ${traceNext}(iterReturn, ctx, iter, ...args);
+          iter.throw = (...args) => ${traceNext}(iterThrow, ctx, iter, ...args);
+
+          return iter;
+        };
+        const ctx = {
+          arguments,
+          self: this,
+          moduleVersion: "1.0.0"
+        };
+        const iter = ${channelVariable}.traceSync(__apm$traced, ctx);
+
+        if (typeof iter.then !== 'function') return wrap(iter);
+
+        return iter.then(result => {
+          ctx.result = result;
+
+          ${channelVariable}.asyncStart.publish(ctx);
+          ${channelVariable}.asyncEnd.publish(ctx);
+
+          return wrap(result);
+        }, err => {
+          ctx.error = err;
+
+          ${channelVariable}.error.publish(ctx);
+          ${channelVariable}.asyncStart.publish(ctx);
+          ${channelVariable}.asyncEnd.publish(ctx);
+
+          return Promise.reject(err);
+        });
+      };
+    }
+  `).body[0].body // Extract only block statement of function body.
+
+  // Replace the right-hand side assignment of `const __apm$wrapped = () => {}`.
+  query(wrapper, '[id.name=__apm$wrapped]')[0].init = node
+
+  return wrapper
+}