dd-trace 5.102.0 → 5.103.0

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

@@ -7,6 +7,10 @@ const {
   channel,
   addHook,
 } = require('./helpers/instrument')
+const {
+  clientToCluster,
+  cloneMessages,
+} = require('./helpers/kafka')
 
 const producerStartCh = channel('apm:kafkajs:produce:start')
 const producerCommitCh = channel('apm:kafkajs:produce:commit')
@@ -24,6 +28,26 @@ const batchConsumerErrorCh = channel('apm:kafkajs:consume-batch:error')
 
 const disabledHeaderWeakSet = new WeakSet()
 
+addHook({ name: 'kafkajs', file: 'src/producer/index.js', versions: ['>=1.4'] }, (createProducer) =>
+  shimmer.wrapFunction(createProducer, original => function wrappedCreateProducer (params) {
+    const producer = original(params)
+    if (params?.cluster) {
+      clientToCluster.set(producer, params.cluster)
+    }
+    return producer
+  })
+)
+
+addHook({ name: 'kafkajs', file: 'src/consumer/index.js', versions: ['>=1.4'] }, (createConsumer) =>
+  shimmer.wrapFunction(createConsumer, original => function wrappedCreateConsumer (params) {
+    const consumer = original(params)
+    if (params?.cluster) {
+      clientToCluster.set(consumer, params.cluster)
+    }
+    return consumer
+  })
+)
+
 addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKafka) => {
   class Kafka extends BaseKafka {
     constructor (options) {
@@ -34,79 +58,95 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
     }
   }
 
-  shimmer.wrap(Kafka.prototype, 'producer', createProducer => function (...args) {
-    const producer = createProducer.apply(this, args)
-    const send = producer.send
+  shimmer.wrap(Kafka.prototype, 'producer', createProducer => function () {
+    const producer = createProducer.apply(this, arguments)
+    const originalSend = producer.send
     const bootstrapServers = this._brokers
-
-    const kafkaClusterIdPromise = getKafkaClusterId(this)
+    const cluster = clientToCluster.get(producer)
 
     producer.send = function (...args) {
-      const wrappedSend = (clusterId) => {
-        const { topic, messages = [] } = args[0]
-
-        const ctx = {
-          bootstrapServers,
-          clusterId,
-          disableHeaderInjection: disabledHeaderWeakSet.has(producer),
-          messages,
-          topic,
-        }
+      if (!producerStartCh.hasSubscribers) {
+        return originalSend.apply(this, args)
+      }
 
-        for (const message of messages) {
-          if (message !== null && typeof message === 'object' && !ctx.disableHeaderInjection) {
-            message.headers = message.headers || {}
-          }
-        }
+      // Fast path: kafkajs has fetched metadata, so clusterId is already on
+      // the broker pool.
+      const metadata = cluster?.brokerPool?.metadata
+      if (metadata) {
+        return runSend.call(this, args, metadata.clusterId)
+      }
 
-        return producerStartCh.runStores(ctx, () => {
-          try {
-            const result = send.apply(this, args)
-            result.then(
-              (res) => {
-                ctx.result = res
-                producerFinishCh.publish(ctx)
-                producerCommitCh.publish(ctx)
-              },
-              (err) => {
-                ctx.error = err
-                if (err) {
-                  // Fixes bug where we would inject message headers for kafka brokers that don't support headers
-                  // (version <0.11). On the error, we disable header injection.
-                  // Unfortunately the error name / type is not more specific.
-                  // This approach is implemented by other tracers as well.
-                  if (err.name === 'KafkaJSProtocolError' && err.type === 'UNKNOWN') {
-                    disabledHeaderWeakSet.add(producer)
-                    log.error(
-                      // eslint-disable-next-line @stylistic/max-len
-                      'Kafka Broker responded with UNKNOWN_SERVER_ERROR (-1). Please look at broker logs for more information. Tracer message header injection for Kafka is disabled.'
-                    )
-                  }
-                  producerErrorCh.publish(err)
-                }
-                producerFinishCh.publish(ctx)
-              })
+      // Slow path, taken at most once per producer connect cycle. Prime the
+      // metadata fetch kafkajs's send would do internally a few stack frames
+      // later. `sharedPromiseTo` collapses our call and kafkajs's call into a
+      // single round trip, so total latency is unchanged.
+      if (typeof cluster?.refreshMetadataIfNecessary !== 'function') {
+        return runSend.call(this, args)
+      }
+      return cluster.refreshMetadataIfNecessary().then(
+        () => runSend.call(this, args, cluster.brokerPool?.metadata?.clusterId),
+        () => runSend.call(this, args)
+      )
+    }
 
-            return result
-          } catch (e) {
-            ctx.error = e
-            producerErrorCh.publish(ctx)
-            producerFinishCh.publish(ctx)
-            throw e
-          }
-        })
+    function runSend (args, clusterId) {
+      const arg0 = args[0]
+      const topic = arg0?.topic
+      const inputMessages = Array.isArray(arg0?.messages) ? arg0.messages : []
+      const disableHeaderInjection = disabledHeaderWeakSet.has(producer)
+
+      // Hand kafkajs and the plugin a shallow clone so injection writes to
+      // tracer-owned objects instead of the caller's. With injection
+      // disabled the clone must not seed `headers: {}` either: brokers that
+      // reject any header field cannot recover otherwise.
+      let messages = inputMessages
+      if (inputMessages.length > 0) {
+        messages = cloneMessages(inputMessages, !disableHeaderInjection)
+        args[0] = { ...arg0, messages }
       }
 
-      if (isPromise(kafkaClusterIdPromise)) {
-        // promise is not resolved
-        return kafkaClusterIdPromise.then((clusterId) => {
-          return wrappedSend(clusterId)
-        })
+      const ctx = {
+        bootstrapServers,
+        clusterId,
+        disableHeaderInjection,
+        messages,
+        topic,
       }
 
-      // promise is already resolved
-      return wrappedSend(kafkaClusterIdPromise)
+      return producerStartCh.runStores(ctx, () => {
+        try {
+          const result = originalSend.apply(this, args)
+          result.then(
+            (res) => {
+              ctx.result = res
+              producerFinishCh.publish(ctx)
+              producerCommitCh.publish(ctx)
+            },
+            (error) => {
+              ctx.error = error
+              if (error) {
+                if (error.name === 'KafkaJSProtocolError' && error.type === 'UNKNOWN') {
+                  disabledHeaderWeakSet.add(producer)
+                  log.error(
+                    // eslint-disable-next-line @stylistic/max-len
+                    'Kafka Broker responded with UNKNOWN_SERVER_ERROR (-1). Please look at broker logs for more information. Tracer message header injection for Kafka is disabled.'
+                  )
+                }
+                producerErrorCh.publish(error)
+              }
+              producerFinishCh.publish(ctx)
+            }
+          )
+          return result
+        } catch (error) {
+          ctx.error = error
+          producerErrorCh.publish(ctx)
+          producerFinishCh.publish(ctx)
+          throw error
+        }
+      })
     }
+
     return producer
   })
 
@@ -115,24 +155,26 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
       return createConsumer.apply(this, args)
     }
 
-    const kafkaClusterIdPromise = getKafkaClusterId(this)
-    let resolvedClusterId = null
+    const consumer = createConsumer.apply(this, arguments)
+    const cluster = clientToCluster.get(consumer)
+    const groupId = arguments[0].groupId
+
+    const readClusterId = () => cluster?.brokerPool?.metadata?.clusterId
 
-    const eachMessageExtractor = (args, clusterId) => {
+    const eachMessageExtractor = (args) => {
       const { topic, partition, message } = args[0]
-      return { topic, partition, message, groupId, clusterId }
+      return { topic, partition, message, groupId, clusterId: readClusterId() }
     }
 
-    const eachBatchExtractor = (args, clusterId) => {
+    const eachBatchExtractor = (args) => {
       const { batch } = args[0]
       const { topic, partition, messages } = batch
-      return { topic, partition, messages, groupId, clusterId }
+      return { topic, partition, messages, groupId, clusterId: readClusterId() }
     }
 
-    const consumer = createConsumer.apply(this, args)
-
     consumer.on(consumer.events.COMMIT_OFFSETS, (event) => {
       const { payload: { groupId: commitGroupId, topics } } = event
+      const clusterId = readClusterId()
       const commitList = []
       for (const { topic, partitions } of topics) {
         for (const { partition, offset } of partitions) {
@@ -141,7 +183,7 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
             partition,
             offset,
             topic,
-            clusterId: resolvedClusterId,
+            clusterId,
           })
         }
       }
@@ -149,118 +191,67 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
     })
 
     const run = consumer.run
-    const groupId = args[0].groupId
 
     consumer.run = function ({ eachMessage, eachBatch, ...runArgs }) {
-      const wrapConsume = (clusterId) => {
-        // In kafkajs COMMIT_OFFSETS always happens in the context of one synchronous run
-        // So this will always reference a correct cluster id
-        resolvedClusterId = clusterId
-        return run({
-          eachMessage: wrappedCallback(
-            eachMessage,
-            consumerStartCh,
-            consumerFinishCh,
-            consumerErrorCh,
-            eachMessageExtractor,
-            clusterId
-          ),
-          eachBatch: wrappedCallback(
-            eachBatch,
-            batchConsumerStartCh,
-            batchConsumerFinishCh,
-            batchConsumerErrorCh,
-            eachBatchExtractor,
-            clusterId
-          ),
-          ...runArgs,
-        })
-      }
-
-      if (isPromise(kafkaClusterIdPromise)) {
-        // promise is not resolved
-        return kafkaClusterIdPromise.then((clusterId) => {
-          return wrapConsume(clusterId)
-        })
-      }
-
-      // promise is already resolved
-      return wrapConsume(kafkaClusterIdPromise)
+      return run({
+        eachMessage: wrappedCallback(
+          eachMessage,
+          consumerStartCh,
+          consumerFinishCh,
+          consumerErrorCh,
+          eachMessageExtractor
+        ),
+        eachBatch: wrappedCallback(
+          eachBatch,
+          batchConsumerStartCh,
+          batchConsumerFinishCh,
+          batchConsumerErrorCh,
+          eachBatchExtractor
+        ),
+        ...runArgs,
+      })
     }
+
     return consumer
   })
+
   return Kafka
 })
 
-const wrappedCallback = (fn, startCh, finishCh, errorCh, extractArgs, clusterId) => {
-  return typeof fn === 'function'
-    ? function (...args) {
-      const extractedArgs = extractArgs(args, clusterId)
-      const ctx = {
-        extractedArgs,
-      }
+const wrappedCallback = (fn, startCh, finishCh, errorCh, extractArgs) => {
+  if (typeof fn !== 'function') return fn
+  return function (...args) {
+    const ctx = {
+      extractedArgs: extractArgs(args),
+    }
 
-      return startCh.runStores(ctx, () => {
-        try {
-          const result = fn.apply(this, args)
-          if (result && typeof result.then === 'function') {
-            result.then(
-              (res) => {
-                ctx.result = res
-                finishCh.publish(ctx)
-              },
-              (err) => {
-                ctx.error = err
-                if (err) {
-                  errorCh.publish(ctx)
-                }
-                finishCh.publish(ctx)
-              })
-          } else {
-            finishCh.publish(ctx)
-          }
-          return result
-        } catch (e) {
-          ctx.error = e
-          errorCh.publish(ctx)
+    return startCh.runStores(ctx, () => {
+      try {
+        const result = fn.apply(this, args)
+        if (result && typeof result.then === 'function') {
+          result.then(
+            (res) => {
+              ctx.result = res
+              finishCh.publish(ctx)
+            },
+            (error) => {
+              ctx.error = error
+              if (error) {
+                errorCh.publish(ctx)
+              }
+              finishCh.publish(ctx)
+            }
+          )
+        } else {
           finishCh.publish(ctx)
-          throw e
         }
-      })
-    }
-    : fn
-}
-
-const getKafkaClusterId = (kafka) => {
-  if (kafka._ddKafkaClusterId) {
-    return kafka._ddKafkaClusterId
-  }
-
-  if (!kafka.admin) {
-    return null
-  }
-
-  const admin = kafka.admin()
-
-  if (!admin.describeCluster) {
-    return null
-  }
-
-  return admin.connect()
-    .then(() => {
-      return admin.describeCluster()
-    })
-    .then((clusterInfo) => {
-      const clusterId = clusterInfo?.clusterId
-      kafka._ddKafkaClusterId = clusterId
-      admin.disconnect()
-      return clusterId
-    })
-    .catch((error) => {
-      throw error
+        return result
+      } catch (error) {
+        ctx.error = error
+        errorCh.publish(ctx)
+        finishCh.publish(ctx)
+        throw error
+      }
     })
-}
-
-function isPromise (obj) {
-  return !!obj && (typeof obj === 'object' || typeof obj === 'function') && typeof obj.then === 'function'
+  }
 }

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

@@ -469,6 +469,10 @@ addHook({
     }
 
     const { suitesByTestFile, numSuitesByTestFile } = getSuitesByTestFile(this.suite)
+    // Root-level tests (direct children of root, no describe wrapper) keyed by file.
+    // Populated during the root 'suite' event so the normal finish path can include them
+    // in mixed-file status calculation.
+    const rootTestsByFile = new Map()
 
     this.once('start', getOnStartHandler(frameworkVersion))
 
@@ -489,6 +493,30 @@ addHook({
 
     this.on('suite', function (suite) {
       if (suite.root || !suite.tests.length) {
+        // This branch can be triggered when we have top level it(...) inside test files.
+        // In that case, they all (even if they are from different files) are going to be
+        // children of the root suite.
+        // Note: We could have suites that contain top level it(...) and also it(...) nested
+        // inside describe(...) ("mixed case"). Duplication is avoided by the context guard
+        // below. Since 'suite' fires for root first, in the mixed case the ctx is created
+        // here and the describe-based handler finds it already set.
+        if (suite.root && suite.tests.length > 0) {
+          const files = new Set(suite.tests.map(test => test.file).filter(Boolean))
+          for (const file of files) {
+            rootTestsByFile.set(file, suite.tests.filter(t => t.file === file))
+            if (testFileToSuiteCtx.get(file)) continue
+            const isUnskippable = unskippableSuites.includes(file)
+            isForcedToRun = isUnskippable && suitesToSkip.includes(getTestSuitePath(file, process.cwd()))
+            const ctx = {
+              testSuiteAbsolutePath: file,
+              isUnskippable,
+              isForcedToRun,
+              itrCorrelationId,
+            }
+            testFileToSuiteCtx.set(file, ctx)
+            testSuiteStartCh.runStores(ctx, () => {})
+          }
+        }
         return
       }
       let ctx = testFileToSuiteCtx.get(suite.file)
@@ -508,6 +536,44 @@ addHook({
 
     this.on('suite end', function (suite) {
       if (suite.root) {
+        // Symmetric to the suite start fix
+        const fileToTests = new Map()
+        for (const test of suite.tests) {
+          if (!test.file) continue
+          if (!fileToTests.has(test.file)) fileToTests.set(test.file, [])
+          fileToTests.get(test.file).push(test)
+        }
+        for (const [file, tests] of fileToTests) {
+          // Mixed case: if a file appears in suitesByTestFile (pre-populated before the run),
+          // its numSuitesByTestFile counter hits zero when its last describe-based suite ends
+          // and the normal path below fires testSuiteFinishCh. Since root is last when
+          // 'suite end' fires, any such file has already been handled — skipping it here
+          // avoids duplication.
+          if (suitesByTestFile[file]) continue
+          let status = 'pass'
+          if (tests.every(test => test.isPending())) {
+            status = 'skip'
+          } else {
+            for (const test of tests) {
+              if (test.state === 'failed' || test.timedOut) {
+                status = 'fail'
+                break
+              }
+            }
+          }
+          if (global.__coverage__) {
+            const coverageFiles = getCoveredFilenamesFromCoverage(global.__coverage__)
+            testSuiteCodeCoverageCh.publish({ coverageFiles, suiteFile: file })
+            mergeCoverage(global.__coverage__, originalCoverageMap)
+            resetCoverage(global.__coverage__)
+          }
+          const ctx = testFileToSuiteCtx.get(file)
+          if (ctx) {
+            testSuiteFinishCh.publish({ status, ...ctx.currentStore }, () => {})
+          } else {
+            log.warn('No ctx found for suite', file)
+          }
+        }
         return
       }
       const suitesInTestFile = suitesByTestFile[suite.file]
@@ -517,8 +583,9 @@ addHook({
         return
       }
 
+      const rootTests = rootTestsByFile.get(suite.file) || []
       let status = 'pass'
-      if (suitesInTestFile.every(suite => suite.pending)) {
+      if (suitesInTestFile.every(suite => suite.pending) && rootTests.every(test => test.isPending())) {
         status = 'skip'
       } else {
         // has to check every test in the test file
@@ -530,6 +597,11 @@ addHook({
             }
           })
         })
+        for (const test of rootTests) {
+          if (test.state === 'failed' || test.timedOut) {
+            status = 'fail'
+          }
+        }
       }
 
       if (global.__coverage__) {

package/packages/datadog-instrumentations/src/mongodb-core.js

@@ -11,6 +11,14 @@ const startCh = channel('apm:mongodb:query:start')
 const finishCh = channel('apm:mongodb:query:finish')
 const errorCh = channel('apm:mongodb:query:error')
 
+// Per-Connection cached topology shape (mongodb >= 4). The Connection's `address` is immutable
+// for the lifetime of the connection, so we synthesize the `{ s: { options } }` envelope the
+// plugin expects only once per connection. A WeakMap keeps the cache off the foreign Connection
+// instance — no extra own-key visible to `Reflect.ownKeys`, `Object.freeze`, or another tracer's
+// instrumentation walking the connection.
+/** @type {WeakMap<object, { s: { options: { host?: string, port?: string } } }>} */
+const topologyCache = new WeakMap()
+
 addHook({ name: 'mongodb-core', versions: ['2 - 3.1.9'] }, Server => {
   const serverProto = Server.Server.prototype
   shimmer.wrap(serverProto, 'command', command => wrapCommand(command, 'command'))
@@ -88,21 +96,36 @@ function wrapUnifiedCommand (command, operation, name) {
 }
 
 function wrapConnectionCommand (command, operation, name, instrumentFn = instrument) {
+  const opts = { name }
   return function (ns, ops) {
     if (!startCh.hasSubscribers) {
       return command.apply(this, arguments)
     }
-    const hostParts = typeof this.address === 'string' ? this.address.split(':') : ''
-    const options = hostParts.length === 2
-      ? { host: hostParts[0], port: hostParts[1] }
-      : {} // no port means the address is a random UUID so no host either
-    const topology = { s: { options } }
-
-    ns = `${ns.db}.${ns.collection}`
-    return instrumentFn(operation, command, this, arguments, topology, ns, ops, { name })
+    let topology = topologyCache.get(this)
+    if (topology === undefined) {
+      topology = synthesizeTopology(this.address)
+      topologyCache.set(this, topology)
+    }
+    return instrumentFn(operation, command, this, arguments, topology, `${ns.db}.${ns.collection}`, ops, opts)
   }
 }
 
+/**
+ * @param {string} address
+ * @returns {{ s: { options: { host?: string, port?: string } } }}
+ */
+function synthesizeTopology (address) {
+  if (typeof address === 'string') {
+    const colon = address.indexOf(':')
+    // Match the previous `.split(':')` length-2 check: exactly one colon with non-empty parts on both sides.
+    if (colon > 0 && colon < address.length - 1 && !address.includes(':', colon + 1)) {
+      return { s: { options: { host: address.slice(0, colon), port: address.slice(colon + 1) } } }
+    }
+  }
+  // No port means the address is a random UUID, an IPv6 form, or otherwise unparseable, so no host either.
+  return { s: { options: {} } }
+}
+
 function wrapQuery (query, operation, name) {
   return function (...args) {
     if (!startCh.hasSubscribers) {
@@ -171,6 +194,8 @@ function instrument (operation, command, instance, args, server, ns, ops, option
   })
 }
 
+module.exports = { synthesizeTopology }
+
 function instrumentPromise (operation, command, instance, args, server, ns, ops, options = {}) {
   const name = options.name || (ops && Object.keys(ops)[0])
 

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

@@ -15,6 +15,10 @@ const errorCh = channel('apm:pg:query:error')
 const startPoolQueryCh = channel('datadog:pg:pool:query:start')
 const finishPoolQueryCh = channel('datadog:pg:pool:query:finish')
 
+// Drivers like pg-promise reuse the same prepared-statement query object across executions; cache
+// the un-injected `text` so the wrap doesn't capture a previous DBM injection as the new original.
+const originalTextCache = new WeakMap()
+
 addHook({ name: 'pg', versions: ['>=8.0.3'], file: 'lib/native/client.js' }, Client => {
   shimmer.wrap(Client.prototype, 'query', query => wrapQuery(query))
   return Client
@@ -43,20 +47,14 @@ function wrapQuery (query) {
       : { text: args[0] }
 
     const textPropObj = pgQuery.cursor ?? pgQuery
-    const textProp = Object.getOwnPropertyDescriptor(textPropObj, 'text')
     const stream = typeof textPropObj.read === 'function'
 
-    // Only alter `text` property if safe to do so. Initially, it's a property, not a getter.
-    let originalText
-    if (!textProp || textProp.configurable) {
+    let originalText = originalTextCache.get(textPropObj)
+    if (originalText === undefined) {
       originalText = textPropObj.text
-
-      Object.defineProperty(textPropObj, 'text', {
-        get () {
-          return this?.__ddInjectableQuery || originalText
-        },
-      })
+      originalTextCache.set(textPropObj, originalText)
     }
+
     const abortController = new AbortController()
     const ctx = {
       params: this.connectionParameters,
@@ -109,6 +107,22 @@ function wrapQuery (query) {
         return Promise.reject(error)
       }
 
+      const injected = ctx.injected
+      if (injected !== undefined) {
+        // Skip the per-read getter trampoline when `text` is a configurable, writable data
+        // property (the pg / pg-cursor common shape). Accessor descriptors and read-only data
+        // still go through `defineProperty(get)` so `get text ()` query objects keep working.
+        const textProp = Object.getOwnPropertyDescriptor(textPropObj, 'text')
+        if (textProp?.configurable === true && textProp.writable === true) {
+          textPropObj.text = injected
+        } else if (textProp === undefined || textProp.configurable === true) {
+          Object.defineProperty(textPropObj, 'text', {
+            configurable: true,
+            get () { return injected },
+          })
+        }
+      }
+
       args[0] = pgQuery
 
       const retval = query.apply(this, args)