dd-trace 5.86.0 → 5.88.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/transformer.js

@@ -0,0 +1,21 @@
+'use strict'
+
+const transforms = require('./transforms')
+
+function transform (state, ...args) {
+  const operator = state.operator = getOperator(state)
+
+  transforms[operator](state, ...args)
+}
+
+function 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'
+  }
+}
+
+module.exports = { transform }

package/packages/datadog-instrumentations/src/helpers/rewriter/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 ({ sourceType }, node) {
     if (node.body.some(tracingChannelPredicate)) return
 
     const index = node.body.findIndex(child => child.directive === 'use strict')
-    const code = format === 'module'
+    const code = sourceType === '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])
+    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,6 +146,55 @@ 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 channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
@@ -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
+}

package/packages/datadog-instrumentations/src/http/client.js

@@ -14,6 +14,7 @@ const finishChannel = channel('apm:http:client:request:finish')
 const endChannel = channel('apm:http:client:request:end')
 const asyncStartChannel = channel('apm:http:client:request:asyncStart')
 const errorChannel = channel('apm:http:client:request:error')
+const responseFinishChannel = channel('apm:http:client:response:finish')
 
 const names = ['http', 'https', 'node:http', 'node:https']
 
@@ -39,6 +40,112 @@ function normalizeCallback (inputOptions, callback, inputURL) {
   return typeof inputOptions === 'function' ? [inputOptions, inputURL || {}] : [callback, inputOptions]
 }
 
+/**
+ * Wires the downstream response so we can observe when the customer consumes
+ * the body and when the stream finishes
+ *
+ * @param {object} ctx - Instrumentation context
+ * @param {import('http').IncomingMessage} res - The downstream response object.
+ * @returns {{ finalizeIfNeeded: () => void }|null} Cleanup helper used for drain.
+ */
+function setupResponseInstrumentation (ctx, res) {
+  const shouldInstrumentFinish = responseFinishChannel.hasSubscribers
+
+  if (!shouldInstrumentFinish) {
+    return null
+  }
+
+  let bodyConsumed = false
+  let finishCalled = false
+  let originalRead = null
+  let dataListenerAdded = false
+  let dataReadStarted = false
+
+  const { shouldCollectBody } = ctx
+  const bodyChunks = shouldCollectBody ? [] : null
+
+  const collectChunk = chunk => {
+    if (!shouldCollectBody || !chunk) return
+
+    if (typeof chunk === 'string') {
+      bodyChunks.push(chunk)
+    } else if (Buffer.isBuffer(chunk)) {
+      bodyChunks.push(chunk)
+    } else {
+      // Handle Uint8Array or other array-like types
+      bodyChunks.push(Buffer.from(chunk))
+    }
+  }
+
+  // Listen for body consumption
+  const onNewListener = (eventName) => {
+    if (eventName === 'data' || eventName === 'readable') {
+      bodyConsumed = true
+
+      // For 'data' events, add our own listener to collect chunks
+      if (eventName === 'data' && !dataListenerAdded && !dataReadStarted) {
+        dataListenerAdded = true
+        res.on('data', collectChunk)
+      }
+
+      // For 'readable' events, wrap the read() method
+      if (eventName === 'readable' && !originalRead && !dataListenerAdded && typeof res.read === 'function') {
+        originalRead = res.read
+        res.read = function () {
+          const chunk = originalRead.apply(this, arguments)
+          if (!dataListenerAdded) {
+            dataReadStarted = true
+            collectChunk(chunk)
+          }
+          return chunk
+        }
+      }
+    }
+  }
+
+  res.on('newListener', onNewListener)
+
+  // Cleanup function to restore original behavior
+  const cleanup = () => {
+    res.off('newListener', onNewListener)
+    res.off('data', collectChunk)
+
+    if (originalRead) {
+      res.read = originalRead
+      originalRead = null
+    }
+  }
+
+  const notifyFinish = () => {
+    if (finishCalled) return
+    finishCalled = true
+
+    // Combine collected chunks into a single body
+    let body = null
+    if (bodyChunks?.length) {
+      const firstChunk = bodyChunks[0]
+      body = typeof firstChunk === 'string'
+        ? bodyChunks.join('')
+        : Buffer.concat(bodyChunks)
+    }
+
+    responseFinishChannel.publish({ ctx, res, body })
+    cleanup()
+  }
+
+  res.once('end', notifyFinish)
+  res.once('close', notifyFinish)
+
+  return {
+    finalizeIfNeeded () {
+      if (!bodyConsumed) {
+        // Body not consumed, resume to complete the response
+        notifyFinish()
+      }
+    },
+  }
+}
+
 function patch (http, methodName) {
   shimmer.wrap(http, methodName, instrumentRequest)
 
@@ -103,7 +210,18 @@ function patch (http, methodName) {
                 ctx.res = res
                 res.once('end', finish)
                 res.once(errorMonitor, finish)
-                break
+
+                const instrumentation = setupResponseInstrumentation(ctx, res)
+
+                if (!instrumentation) {
+                  break
+                }
+
+                const result = emit.apply(this, arguments)
+
+                instrumentation.finalizeIfNeeded()
+
+                return result
               }
               case 'connect':
               case 'upgrade':