dd-trace 5.79.0 → 5.81.0

package/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('../../../../../vendor/dist/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('../../../../../vendor/dist/@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 }

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

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

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

@@ -0,0 +1,237 @@
+'use strict'
+
+module.exports = [
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.cjs'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.cjs'
+    },
+    functionQuery: {
+      methodName: 'batch',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_batch'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/chat_models.cjs'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async',
+      className: 'BaseChatModel'
+    },
+    channelName: 'BaseChatModel_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/llms.cjs'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async'
+    },
+    channelName: 'BaseLLM_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.cjs'
+    },
+    functionQuery: {
+      methodName: 'embedQuery',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedQuery'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.cjs'
+    },
+    functionQuery: {
+      methodName: 'embedDocuments',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedDocuments'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/tools/index.cjs'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'StructuredTool'
+    },
+    channelName: 'Tool_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.cjs'
+    },
+    functionQuery: {
+      methodName: 'similaritySearch',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearch'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.cjs'
+    },
+    functionQuery: {
+      methodName: 'similaritySearchWithScore',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearchWithScore'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.js'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.js'
+    },
+    functionQuery: {
+      methodName: 'batch',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_batch'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/chat_models.js'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async',
+      className: 'BaseChatModel'
+    },
+    channelName: 'BaseChatModel_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/llms.js'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async'
+    },
+    channelName: 'BaseLLM_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.js'
+    },
+    functionQuery: {
+      methodName: 'embedQuery',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedQuery'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.js'
+    },
+    functionQuery: {
+      methodName: 'embedDocuments',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedDocuments'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/tools/index.js'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'StructuredTool'
+    },
+    channelName: 'Tool_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.js'
+    },
+    functionQuery: {
+      methodName: 'similaritySearch',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearch'
+  },
+  {
+
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.js'
+    },
+    functionQuery: {
+      methodName: 'similaritySearchWithScore',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearchWithScore'
+  }
+]

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

@@ -0,0 +1,9 @@
+'use strict'
+
+const Module = require('module')
+const shimmer = require('../../../../datadog-shimmer')
+const { rewrite } = require('./')
+
+shimmer.wrap(Module.prototype, '_compile', compile => function (content, filename, format) {
+  return compile.call(this, rewrite(content, filename, format), filename, format)
+})

package/packages/datadog-instrumentations/src/helpers/rewriter/loader.mjs

@@ -0,0 +1,11 @@
+import { rewrite } from './index.js'
+
+async function load (url, context, nextLoad) {
+  const result = await nextLoad(url, context)
+
+  result.source = rewrite(result.source, url, context.format)
+
+  return result
+}
+
+export { load }

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

@@ -0,0 +1,139 @@
+'use strict'
+
+const { parse } = require('./compiler')
+
+const tracingChannelPredicate = (node) => (
+  node.specifiers?.[0]?.local?.name === 'tr_ch_apm_tracingChannel' ||
+    node.declarations?.[0]?.id?.properties?.[0]?.value?.name === 'tr_ch_apm_tracingChannel'
+)
+
+const transforms = module.exports = {
+  tracingChannelImport ({ format }, 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")'
+
+    node.body.splice(index + 1, 0, parse(code, { module: format === 'module' }).body[0])
+  },
+
+  tracingChannelDeclaration (state, node) {
+    const { channelName, module: { name } } = state
+    const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
+
+    if (node.body.some(child => child.declarations?.[0]?.id?.name === channelVariable)) return
+
+    transforms.tracingChannelImport(state, node)
+
+    const index = node.body.findIndex(tracingChannelPredicate)
+    const code = `
+      const ${channelVariable} = tr_ch_apm_tracingChannel("orchestrion:${name}:${channelName}")
+    `
+
+    node.body.splice(index + 1, 0, parse(code).body[0])
+  },
+
+  traceCallback: traceAny,
+  tracePromise: traceAny,
+  traceSync: traceAny,
+}
+
+function traceAny (state, node, _parent, ancestry) {
+  const program = ancestry[ancestry.length - 1]
+
+  if (node.type === 'ClassDeclaration' || node.type === 'ClassExpression') {
+    traceInstanceMethod(state, node, program)
+  } else {
+    traceFunction(state, node, program)
+  }
+}
+
+function traceFunction (state, node, program) {
+  const { operator } = state
+
+  transforms.tracingChannelDeclaration(state, program)
+
+  node.body = wrap(state, {
+    type: 'ArrowFunctionExpression',
+    params: node.params,
+    body: node.body,
+    async: operator === 'tracePromise',
+    expression: false,
+    generator: false
+  })
+}
+
+function traceInstanceMethod (state, node, program) {
+  const { functionQuery, operator } = state
+  const { methodName } = functionQuery
+
+  const classBody = node.body
+
+  // If the method exists on the class, we return as it will be patched later
+  // while traversing child nodes later on.
+  if (classBody.body.some(({ key }) => key.name === methodName)) return
+
+  // Method doesn't exist on the class so we assume an instance method and
+  // wrap it in the constructor instead.
+  let ctor = classBody.body.find(({ kind }) => kind === 'constructor')
+
+  transforms.tracingChannelDeclaration(state, program)
+
+  if (!ctor) {
+    ctor = parse(
+      node.superClass
+        ? 'class A { constructor (...args) { super(...args) } }'
+        : 'class A { constructor () {} }'
+    ).body[0].body.body[0] // Extract constructor from dummy class body.
+
+    classBody.body.unshift(ctor)
+  }
+
+  const ctorBody = parse(`
+    const __apm$${methodName} = this["${methodName}"]
+    this["${methodName}"] = function () {}
+  `).body
+
+  // Extract only right-hand side function of line 2.
+  const fn = ctorBody[1].expression.right
+
+  fn.async = operator === 'tracePromise'
+  fn.body = wrap(state, { type: 'Identifier', name: `__apm$${methodName}` })
+
+  ctor.value.body.body.push(...ctorBody)
+}
+
+function wrap (state, node) {
+  const { channelName, operator, functionQuery: { index = -1 } } = state
+  const async = operator === 'tracePromise' ? 'async' : ''
+  const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
+  const tracedArgs = operator === 'traceCallback'
+    ? `__apm$original_args.splice(${index}, 1, arguments[${index >= 0 ? index : `arguments.length + ${index}`}])`
+    : '__apm$original_args'
+  const traceParams = operator === 'traceCallback'
+    ? `__apm$traced, ${index}`
+    : '__apm$traced'
+  const wrapper = parse(`
+    function wrapper () {
+      const __apm$original_args = arguments;
+      const __apm$traced = ${async} function () {
+        const __apm$wrapped = () => {};
+        const __apm$traced_args = ${tracedArgs};
+        return __apm$wrapped.apply(this, __apm$traced_args);
+      };
+      if (!${channelVariable}.hasSubscribers) return __apm$traced.apply(this, arguments);
+      return ${channelVariable}.${operator}(${traceParams}, {
+        arguments,
+        self: this,
+        moduleVersion: "1.0.0"
+      }, this, ...arguments);
+    }
+  `).body[0].body // Extract only block statement of function body.
+
+  // Replace the right-hand side assignment of `const __apm$wrapped = () => {}`.
+  wrapper.body[1].declarations[0].init.body.body[0].declarations[0].init = node
+
+  return wrapper
+}

package/packages/datadog-instrumentations/src/jest.js

@@ -1110,7 +1110,7 @@ function jestAdapterWrapper (jestAdapter, jestVersion) {
        * Child processes do not each request ITR configuration, so the jest's parent process
        * needs to pass them the configuration. This is done via _ddTestCodeCoverageEnabled, which
        * controls whether coverage is reported.
-      */
+       */
       if (environment.testEnvironmentOptions?._ddTestCodeCoverageEnabled) {
         const root = environment.repositoryRoot || environment.rootDir
 

package/packages/datadog-instrumentations/src/langchain.js

@@ -1,113 +1,7 @@
 'use strict'
-const { addHook } = require('./helpers/instrument')
 
-const shimmer = require('../../datadog-shimmer')
+const { addHook, getHooks } = require('./helpers/instrument')
 
-const tracingChannel = require('dc-polyfill').tracingChannel
-
-function wrap (obj, name, channelName, namespace) {
-  const channel = tracingChannel(channelName)
-  shimmer.wrap(obj, name, function (original) {
-    return function () {
-      if (!channel.start.hasSubscribers) {
-        return original.apply(this, arguments)
-      }
-      const ctx = { self: this, arguments }
-      if (namespace) {
-        ctx.namespace = namespace
-      }
-      return channel.tracePromise(original, ctx, this, ...arguments)
-    }
-  })
-}
-
-// langchain compiles into ESM and CommonJS, with ESM being the default and landing in the `.js` files
-// however, CommonJS ends up in `cjs` files, and are required under the hood with `.cjs` files
-// we patch each separately and explicitly to match against exports only once, and not rely on file regex matching
-const extensions = ['js', 'cjs']
-
-for (const extension of extensions) {
-  addHook({ name: '@langchain/core', file: `dist/runnables/base.${extension}`, versions: ['>=0.1'] }, exports => {
-    if (extension === 'cjs') {
-      wrap(exports.RunnableSequence.prototype, 'invoke', 'orchestrion:@langchain/core:RunnableSequence_invoke')
-      wrap(exports.RunnableSequence.prototype, 'batch', 'orchestrion:@langchain/core:RunnableSequence_batch')
-    }
-    return exports
-  })
-
-  addHook({
-    name: '@langchain/core',
-    file: `dist/language_models/chat_models.${extension}`,
-    versions: ['>=0.1']
-  }, exports => {
-    if (extension === 'cjs') {
-      wrap(exports.BaseChatModel.prototype, 'generate', 'orchestrion:@langchain/core:BaseChatModel_generate')
-    }
-    return exports
-  })
-
-  addHook({ name: '@langchain/core', file: `dist/language_models/llms.${extension}`, versions: ['>=0.1'] }, exports => {
-    if (extension === 'cjs') {
-      wrap(exports.BaseLLM.prototype, 'generate', 'orchestrion:@langchain/core:BaseLLM_generate')
-    }
-    return exports
-  })
-
-  addHook({ name: '@langchain/core', file: `dist/tools/index.${extension}`, versions: ['>=0.1'] }, exports => {
-    if (extension === 'cjs') {
-      wrap(exports.StructuredTool.prototype, 'invoke', 'orchestrion:@langchain/core:Tool_invoke')
-    }
-    return exports
-  })
-
-  addHook({ name: '@langchain/core', file: `dist/vectorstores.${extension}`, versions: ['>=0.1'] }, exports => {
-    if (extension === 'cjs') {
-      wrap(
-        exports.VectorStore.prototype, 'similaritySearch', 'orchestrion:@langchain/core:VectorStore_similaritySearch'
-      )
-      wrap(
-        exports.VectorStore.prototype, 'similaritySearchWithScore',
-        'orchestrion:@langchain/core:VectorStore_similaritySearchWithScore'
-      )
-    }
-
-    return exports
-  })
-
-  addHook({ name: '@langchain/core', file: `dist/embeddings.${extension}`, versions: ['>=0.1'] }, exports => {
-    if (extension === 'cjs') {
-      shimmer.wrap(exports, 'Embeddings', Embeddings => {
-        return class extends Embeddings {
-          constructor (...args) {
-            super(...args)
-
-            const namespace = ['langchain', 'embeddings']
-
-            if (this.constructor.name === 'OpenAIEmbeddings') {
-              namespace.push('openai')
-            }
-
-            wrap(this, 'embedQuery', 'apm:@langchain/core:Embeddings_embedQuery', namespace)
-            wrap(this, 'embedDocuments', 'apm:@langchain/core:Embeddings_embedDocuments', namespace)
-          }
-        }
-      })
-    } else {
-      const channel = tracingChannel('orchestrion:@langchain/core:Embeddings_constructor')
-      channel.subscribe({
-        end (ctx) {
-          const { self } = ctx
-          const namespace = ['langchain', 'embeddings']
-
-          if (self.constructor.name === 'OpenAIEmbeddings') {
-            namespace.push('openai')
-          }
-
-          wrap(self, 'embedQuery', 'apm:@langchain/core:Embeddings_embedQuery', namespace)
-          wrap(self, 'embedDocuments', 'apm:@langchain/core:Embeddings_embedDocuments', namespace)
-        }
-      })
-    }
-    return exports
-  })
+for (const hook of getHooks('@langchain/core')) {
+  addHook(hook, exports => exports)
 }

package/packages/datadog-instrumentations/src/mocha/main.js

@@ -1,6 +1,6 @@
 'use strict'
 
-const { createCoverageMap } = require('istanbul-lib-coverage')
+const { createCoverageMap } = require('../../../../vendor/dist/istanbul-lib-coverage')
 const { addHook, channel } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
 const { isMarkedAsUnskippable } = require('../../../datadog-plugin-jest/src/util')