dd-trace 5.53.0 → 5.55.0

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

@@ -2,8 +2,7 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -60,33 +59,37 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
 
     producer.send = function () {
       const wrappedSend = (clusterId) => {
-        const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
+        const { topic, messages = [] } = arguments[0]
+
+        const ctx = {
+          bootstrapServers,
+          clusterId,
+          disableHeaderInjection: disabledHeaderWeakSet.has(producer),
+          messages,
+          topic
+        }
 
-        return innerAsyncResource.runInAsyncScope(() => {
-          if (!producerStartCh.hasSubscribers) {
-            return send.apply(this, arguments)
+        for (const message of messages) {
+          if (message !== null && typeof message === 'object' && !ctx.disableHeaderInjection) {
+            message.headers = message.headers || {}
           }
+        }
 
+        return producerStartCh.runStores(ctx, () => {
           try {
-            const { topic, messages = [] } = arguments[0]
-            producerStartCh.publish({
-              topic,
-              messages,
-              bootstrapServers,
-              clusterId,
-              disableHeaderInjection: disabledHeaderWeakSet.has(producer)
-            })
             const result = send.apply(this, arguments)
             result.then(
-              innerAsyncResource.bind(res => {
-                producerFinishCh.publish(undefined)
-                producerCommitCh.publish(res)
-              }),
-              innerAsyncResource.bind(err => {
+              (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.
-                  // Tnfortunately the error name / type is not more specific.
+                  // 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)
@@ -96,28 +99,28 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
                   }
                   producerErrorCh.publish(err)
                 }
-                producerFinishCh.publish(undefined)
+                producerFinishCh.publish(ctx)
               })
-            )
 
             return result
           } catch (e) {
-            producerErrorCh.publish(e)
-            producerFinishCh.publish(undefined)
+            ctx.error = e
+            producerErrorCh.publish(ctx)
+            producerFinishCh.publish(ctx)
             throw e
           }
         })
       }
 
-      if (!isPromise(kafkaClusterIdPromise)) {
-        // promise is already resolved
-        return wrappedSend(kafkaClusterIdPromise)
-      } else {
+      if (isPromise(kafkaClusterIdPromise)) {
         // promise is not resolved
         return kafkaClusterIdPromise.then((clusterId) => {
           return wrappedSend(clusterId)
         })
       }
+
+      // promise is already resolved
+      return wrappedSend(kafkaClusterIdPromise)
     }
     return producer
   })
@@ -170,15 +173,15 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
         })
       }
 
-      if (!isPromise(kafkaClusterIdPromise)) {
-        // promise is already resolved
-        return wrapConsume(kafkaClusterIdPromise)
-      } else {
+      if (isPromise(kafkaClusterIdPromise)) {
         // promise is not resolved
         return kafkaClusterIdPromise.then((clusterId) => {
           return wrapConsume(clusterId)
         })
       }
+
+      // promise is already resolved
+      return wrapConsume(kafkaClusterIdPromise)
     }
     return consumer
   })
@@ -188,30 +191,35 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
 const wrappedCallback = (fn, startCh, finishCh, errorCh, extractArgs, clusterId) => {
   return typeof fn === 'function'
     ? function (...args) {
-      const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
-      return innerAsyncResource.runInAsyncScope(() => {
-        const extractedArgs = extractArgs(args, clusterId)
+      const extractedArgs = extractArgs(args, clusterId)
+      const ctx = {
+        extractedArgs
+      }
 
-        startCh.publish(extractedArgs)
+      return startCh.runStores(ctx, () => {
         try {
           const result = fn.apply(this, args)
           if (result && typeof result.then === 'function') {
             result.then(
-              innerAsyncResource.bind(() => finishCh.publish(undefined)),
-              innerAsyncResource.bind(err => {
+              (res) => {
+                ctx.result = res
+                finishCh.publish(ctx)
+              },
+              (err) => {
+                ctx.error = err
                 if (err) {
-                  errorCh.publish(err)
+                  errorCh.publish(ctx)
                 }
-                finishCh.publish(undefined)
+                finishCh.publish(ctx)
               })
-            )
           } else {
-            finishCh.publish(undefined)
+            finishCh.publish(ctx)
           }
           return result
         } catch (e) {
-          errorCh.publish(e)
-          finishCh.publish(undefined)
+          ctx.error = e
+          errorCh.publish(ctx)
+          finishCh.publish(ctx)
           throw e
         }
       })

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

@@ -22,6 +22,10 @@ function patch (file) {
   })
 }
 
+function finish () {
+  finishRawQueryCh.publish()
+}
+
 addHook({
   name: 'knex',
   versions: ['>=2'],
@@ -41,10 +45,6 @@ addHook({
 
     const asyncResource = new AsyncResource('bound-anonymous-fn')
 
-    function finish () {
-      finishRawQueryCh.publish()
-    }
-
     return asyncResource.runInAsyncScope(() => {
       startRawQueryCh.publish({ sql, dialect: this.dialect })
 

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

@@ -112,10 +112,9 @@ function wrapMiddleware (fn, layer) {
             throw err
           }
         )
-      } else {
-        fulfill(ctx)
-        return result
       }
+      fulfill(ctx)
+      return result
     } catch (e) {
       fulfill(ctx, e)
       throw e

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

@@ -8,6 +8,7 @@ const {
 const shimmer = require('../../datadog-shimmer')
 
 function isString (value) {
+  // eslint-disable-next-line unicorn/no-instanceof-builtins
   return typeof value === 'string' || value instanceof String
 }
 
@@ -61,10 +62,8 @@ addHook({ name: 'ldapjs', versions: ['>=2'] }, ldapjs => {
       let filter
       if (isString(options)) {
         filter = options
-      } else if (options !== null && typeof options === 'object' && options.filter) {
-        if (isString(options.filter)) {
-          filter = options.filter
-        }
+      } else if (options !== null && typeof options === 'object' && options.filter && isString(options.filter)) {
+        filter = options.filter
       }
       ldapSearchCh.publish({ base, filter })
     }

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

@@ -1,49 +1,39 @@
 'use strict'
 
-const { channel, addHook, AsyncResource } = require('./helpers/instrument')
+const { channel, addHook } = require('./helpers/instrument')
 
 const shimmer = require('../../datadog-shimmer')
 
+const commandAddCh = channel('apm:mariadb:command:add')
+const connectionStartCh = channel('apm:mariadb:connection:start')
+const connectionFinishCh = channel('apm:mariadb:connection:finish')
 const startCh = channel('apm:mariadb:query:start')
 const finishCh = channel('apm:mariadb:query:finish')
 const errorCh = channel('apm:mariadb:query:error')
 const skipCh = channel('apm:mariadb:pool:skip')
-const unskipCh = channel('apm:mariadb:pool:unskip')
 
-function wrapCommandStart (start, callbackResource) {
+function wrapCommandStart (start, ctx) {
   return shimmer.wrapFunction(start, start => function () {
     if (!startCh.hasSubscribers) return start.apply(this, arguments)
 
-    const resolve = callbackResource.bind(this.resolve)
-    const reject = callbackResource.bind(this.reject)
-
-    const asyncResource = callbackResource.runInAsyncScope(() => new AsyncResource('bound-anonymous-fn'))
-
+    const { reject, resolve } = this
     shimmer.wrap(this, 'resolve', function wrapResolve () {
       return function () {
-        asyncResource.runInAsyncScope(() => {
-          finishCh.publish()
-        })
-
-        return resolve.apply(this, arguments)
+        return finishCh.runStores(ctx, resolve, this, ...arguments)
       }
     })
 
     shimmer.wrap(this, 'reject', function wrapReject () {
       return function (error) {
-        asyncResource.runInAsyncScope(() => {
-          errorCh.publish(error)
-          finishCh.publish()
-        })
+        ctx.error = error
 
-        return reject.apply(this, arguments)
+        errorCh.publish(ctx)
+
+        return finishCh.runStores(ctx, reject, this, ...arguments)
       }
     })
 
-    return asyncResource.runInAsyncScope(() => {
-      startCh.publish({ sql: this.sql, conf: this.opts })
-      return start.apply(this, arguments)
-    })
+    return startCh.runStores(ctx, start, this, ...arguments)
   })
 }
 
@@ -52,11 +42,13 @@ function wrapCommand (Command) {
     constructor (...args) {
       super(...args)
 
-      const callbackResource = new AsyncResource('bound-anonymous-fn')
+      if (!this.start) return
 
-      if (this.start) {
-        this.start = wrapCommandStart(this.start, callbackResource)
-      }
+      const ctx = { sql: this.sql, conf: this.opts }
+
+      commandAddCh.publish(ctx)
+
+      this.start = wrapCommandStart(this.start, ctx)
     }
   }
 }
@@ -66,21 +58,19 @@ function createWrapQuery (options) {
     return function (sql) {
       if (!startCh.hasSubscribers) return query.apply(this, arguments)
 
-      const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-      return asyncResource.runInAsyncScope(() => {
-        startCh.publish({ sql, conf: options })
-
-        return query.apply(this, arguments)
-          .then(result => {
-            finishCh.publish()
-            return result
-          }, error => {
-            errorCh.publish(error)
-            finishCh.publish()
-            throw error
-          })
-      }, 'bound-anonymous-fn')
+      const ctx = { sql, conf: options }
+
+      return startCh.runStores(ctx, query, this, ...arguments)
+        .then(result => {
+          ctx.result = result
+          finishCh.publish(ctx)
+          return result
+        }, error => {
+          ctx.error
+          errorCh.publish(ctx)
+          finishCh.publish(ctx)
+          throw error
+        })
     }
   }
 }
@@ -91,30 +81,24 @@ function createWrapQueryCallback (options) {
       if (!startCh.hasSubscribers) return query.apply(this, arguments)
 
       const cb = arguments[arguments.length - 1]
-      const asyncResource = new AsyncResource('bound-anonymous-fn')
-      const callbackResource = new AsyncResource('bound-anonymous-fn')
+      const ctx = { sql, conf: options }
 
       if (typeof cb !== 'function') {
-        arguments.length = arguments.length + 1
+        arguments.length += 1
       }
 
-      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (err) {
+      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => function (err) {
         if (err) {
-          errorCh.publish(err)
-        }
-
-        finishCh.publish()
-
-        if (typeof cb === 'function') {
-          return callbackResource.runInAsyncScope(() => cb.apply(this, arguments))
+          ctx.error = err
+          errorCh.publish(ctx)
         }
-      }))
 
-      return asyncResource.runInAsyncScope(() => {
-        startCh.publish({ sql, conf: options })
+        return typeof cb === 'function'
+          ? finishCh.runStores(ctx, cb, this, ...arguments)
+          : finishCh.publish(ctx)
+      })
 
-        return query.apply(this, arguments)
-      }, 'bound-anonymous-fn')
+      return startCh.runStores(ctx, query, this, ...arguments)
     }
   }
 }
@@ -144,12 +128,7 @@ function wrapPoolBase (PoolBase) {
 // and/or orphan spans.
 function wrapPoolMethod (createConnection) {
   return function () {
-    skipCh.publish()
-    try {
-      return createConnection.apply(this, arguments)
-    } finally {
-      unskipCh.publish()
-    }
+    return skipCh.runStores({}, createConnection, this, ...arguments)
   }
 }
 
@@ -158,8 +137,13 @@ function wrapPoolGetConnectionMethod (getConnection) {
     const cb = arguments[arguments.length - 1]
     if (typeof cb !== 'function') return getConnection.apply(this, arguments)
 
-    const callbackResource = new AsyncResource('bound-anonymous-fn')
-    arguments[arguments.length - 1] = callbackResource.bind(cb)
+    const ctx = {}
+
+    arguments[arguments.length - 1] = function () {
+      return connectionFinishCh.runStores(ctx, cb, this, ...arguments)
+    }
+
+    connectionStartCh.publish(ctx)
 
     return getConnection.apply(this, arguments)
   }