dd-trace 5.61.1 → 5.63.0

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

@@ -3,8 +3,7 @@
 const { errorMonitor } = require('events')
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -24,31 +23,50 @@ function wrapAllNames (names, action) {
   names.forEach(name => action(name))
 }
 
-// semver >=2 <3
-function wrapMaybeInvoke (_maybeInvoke) {
-  const wrapped = function (fn, args) {
-    if (!Array.isArray(args)) return _maybeInvoke.apply(this, arguments)
-
-    const callbackIndex = args.length - 1
-    const callback = args[callbackIndex]
+function wrapCallback (callback, ctx, channelPrefix) {
+  const callbackStartCh = channel(`${channelPrefix}:callback:start`)
+  const callbackFinishCh = channel(`${channelPrefix}:callback:finish`)
 
-    if (typeof callback === 'function') {
-      args[callbackIndex] = AsyncResource.bind(callback)
+  const wrapped = callbackStartCh.runStores(ctx, () => {
+    return function (...args) {
+      return callbackFinishCh.runStores(ctx, () => {
+        return callback.apply(this, args)
+      })
     }
-
-    return _maybeInvoke.apply(this, arguments)
-  }
+  })
+  Object.defineProperty(wrapped, '_dd_wrapped', { value: true })
   return wrapped
 }
 
 function wrapQuery (query) {
-  const wrapped = function (q, params, callback) {
-    if (typeof arguments[arguments.length - 1] === 'function') {
-      arguments[arguments.length - 1] = AsyncResource.bind(arguments[arguments.length - 1])
+  return function (q, params, callback) {
+    const cb = arguments[arguments.length - 1]
+    if (typeof cb === 'function') {
+      const ctx = {}
+      arguments[arguments.length - 1] = wrapCallback(cb, ctx, 'apm:couchbase:query')
     }
 
     return query.apply(this, arguments)
   }
+}
+
+function wrapCallbackFinish (callback, thisArg, _args, errorCh, finishCh, ctx, channelPrefix) {
+  const callbackStartCh = channel(`${channelPrefix}:callback:start`)
+  const callbackFinishCh = channel(`${channelPrefix}:callback:finish`)
+
+  const wrapped = callbackStartCh.runStores(ctx, () => {
+    return function finish (error, result) {
+      return callbackFinishCh.runStores(ctx, () => {
+        if (error) {
+          ctx.error = error
+          errorCh.publish(ctx)
+        }
+        finishCh.publish(ctx)
+        return callback.apply(thisArg, [error, result])
+      })
+    }
+  })
+  Object.defineProperty(wrapped, '_dd_wrapped', { value: true })
   return wrapped
 }
 
@@ -62,31 +80,24 @@ function wrap (prefix, fn) {
       return fn.apply(this, arguments)
     }
 
-    const callbackIndex = findCallbackIndex(arguments)
+    const callbackIndex = findCallbackIndex(arguments, 1)
 
     if (callbackIndex < 0) return fn.apply(this, arguments)
 
-    const callbackResource = new AsyncResource('bound-anonymous-fn')
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-    return asyncResource.runInAsyncScope(() => {
-      const cb = callbackResource.bind(arguments[callbackIndex])
+    const ctx = { bucket: { name: this.name || this._name }, seedNodes: this._dd_hosts }
+    return startCh.runStores(ctx, () => {
+      const cb = arguments[callbackIndex]
 
-      startCh.publish({ bucket: { name: this.name || this._name }, seedNodes: this._dd_hosts })
-
-      arguments[callbackIndex] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (error, result) {
-        if (error) {
-          errorCh.publish(error)
-        }
-        finishCh.publish(result)
-        return cb.apply(this, arguments)
-      }))
+      arguments[callbackIndex] = shimmer.wrapFunction(cb, (cb) => {
+        return wrapCallbackFinish(cb, this, arguments, errorCh, finishCh, ctx, prefix)
+      })
 
       try {
         return fn.apply(this, arguments)
       } catch (error) {
+        ctx.error = error
         error.stack // trigger getting the stack at the original throwing point
-        errorCh.publish(error)
+        errorCh.publish(ctx)
 
         throw error
       }
@@ -95,6 +106,26 @@ function wrap (prefix, fn) {
   return wrapped
 }
 
+// semver >=2 <3
+function wrapMaybeInvoke (_maybeInvoke, channelPrefix) {
+  return function (fn, args) {
+    if (!Array.isArray(args)) return _maybeInvoke.apply(this, arguments)
+
+    const callbackIndex = findCallbackIndex(args, 0)
+
+    if (callbackIndex === -1) return _maybeInvoke.apply(this, arguments)
+
+    const callback = args[callbackIndex]
+
+    if (typeof callback === 'function' && !callback._dd_wrapped) {
+      const ctx = {}
+      args[callbackIndex] = wrapCallback(callback, ctx, channelPrefix)
+    }
+
+    return _maybeInvoke.apply(this, arguments)
+  }
+}
+
 // semver >=3
 
 function wrapCBandPromise (fn, name, startData, thisArg, args) {
@@ -104,36 +135,36 @@ function wrapCBandPromise (fn, name, startData, thisArg, args) {
 
   if (!startCh.hasSubscribers) return fn.apply(thisArg, args)
 
-  const asyncResource = new AsyncResource('bound-anonymous-fn')
-  const callbackResource = new AsyncResource('bound-anonymous-fn')
-
-  return asyncResource.runInAsyncScope(() => {
-    startCh.publish(startData)
-
+  const ctx = startData
+  return startCh.runStores(ctx, () => {
     try {
       const cbIndex = findCallbackIndex(args, 1)
       if (cbIndex >= 0) {
         // v3 offers callback or promises event handling
         // NOTE: this does not work with v3.2.0-3.2.1 cluster.query, as there is a bug in the couchbase source code
-        const cb = callbackResource.bind(args[cbIndex])
-        args[cbIndex] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (error, result) {
-          if (error) {
-            errorCh.publish(error)
-          }
-          finishCh.publish({ result })
-          return cb.apply(thisArg, arguments)
-        }))
+        args[cbIndex] = shimmer.wrapFunction(args[cbIndex], (cb) => {
+          return wrapCallbackFinish(cb, thisArg, args, errorCh, finishCh, ctx, `apm:couchbase:${name}`)
+        })
       }
       const res = fn.apply(thisArg, args)
 
       // semver >=3 will always return promise by default
       res.then(
-        asyncResource.bind((result) => finishCh.publish({ result })),
-        asyncResource.bind((err) => errorCh.publish(err)))
+        (result) => {
+          ctx.result = result
+          finishCh.publish(ctx)
+        },
+        (err) => {
+          ctx.error = err
+          errorCh.publish(ctx)
+          finishCh.publish(ctx)
+        }
+      )
       return res
     } catch (e) {
       e.stack
-      errorCh.publish(e)
+      ctx.error = e
+      errorCh.publish(ctx)
       throw e
     }
   })
@@ -160,11 +191,14 @@ function wrapV3Query (query) {
 
 // semver >=2 <3
 addHook({ name: 'couchbase', file: 'lib/bucket.js', versions: ['^2.6.12'] }, Bucket => {
+  shimmer.wrap(Bucket.prototype, '_maybeInvoke', maybeInvoke => {
+    return wrapMaybeInvoke(maybeInvoke, 'apm:couchbase:bucket:maybeInvoke')
+  })
+
   const startCh = channel('apm:couchbase:query:start')
   const finishCh = channel('apm:couchbase:query:finish')
   const errorCh = channel('apm:couchbase:query:error')
 
-  shimmer.wrap(Bucket.prototype, '_maybeInvoke', maybeInvoke => wrapMaybeInvoke(maybeInvoke))
   shimmer.wrap(Bucket.prototype, 'query', query => wrapQuery(query))
 
   shimmer.wrap(Bucket.prototype, '_n1qlReq', _n1qlReq => function (host, q, adhoc, emitter) {
@@ -176,24 +210,25 @@ addHook({ name: 'couchbase', file: 'lib/bucket.js', versions: ['^2.6.12'] }, Buc
 
     const n1qlQuery = getQueryResource(q)
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-    return asyncResource.runInAsyncScope(() => {
-      startCh.publish({ resource: n1qlQuery, bucket: { name: this.name || this._name }, seedNodes: this._dd_hosts })
-
-      emitter.once('rows', asyncResource.bind(() => {
-        finishCh.publish()
-      }))
+    const ctx = { resource: n1qlQuery, bucket: { name: this.name || this._name }, seedNodes: this._dd_hosts }
+    return startCh.runStores(ctx, () => {
+      emitter.once('rows', () => {
+        finishCh.publish(ctx)
+      })
 
-      emitter.once(errorMonitor, asyncResource.bind((error) => {
-        errorCh.publish(error)
-        finishCh.publish()
-      }))
+      emitter.once(errorMonitor, (error) => {
+        if (!error) return
+        ctx.error = error
+        errorCh.publish(ctx)
+        finishCh.publish(ctx)
+      })
 
       try {
         return _n1qlReq.apply(this, arguments)
       } catch (err) {
         err.stack // trigger getting the stack at the original throwing point
-        errorCh.publish(err)
+        ctx.error = err
+        errorCh.publish(ctx)
 
         throw err
       }
@@ -208,9 +243,11 @@ addHook({ name: 'couchbase', file: 'lib/bucket.js', versions: ['^2.6.12'] }, Buc
 })
 
 addHook({ name: 'couchbase', file: 'lib/cluster.js', versions: ['^2.6.12'] }, Cluster => {
-  shimmer.wrap(Cluster.prototype, '_maybeInvoke', maybeInvoke => wrapMaybeInvoke(maybeInvoke))
-  shimmer.wrap(Cluster.prototype, 'query', query => wrapQuery(query))
+  shimmer.wrap(Cluster.prototype, '_maybeInvoke', maybeInvoke => {
+    return wrapMaybeInvoke(maybeInvoke, 'apm:couchbase:cluster:maybeInvoke')
+  })
 
+  shimmer.wrap(Cluster.prototype, 'query', query => wrapQuery(query))
   shimmer.wrap(Cluster.prototype, 'openBucket', openBucket => {
     return function () {
       const bucket = openBucket.apply(this, arguments)

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

@@ -1,7 +1,7 @@
 'use strict'
 
 const shimmer = require('../../datadog-shimmer')
-const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+const { addHook, channel } = require('./helpers/instrument')
 
 const errorChannel = channel('apm:fastify:middleware:error')
 const handleChannel = channel('apm:fastify:request:handle')
@@ -13,7 +13,12 @@ const responsePayloadReadCh = channel('datadog:fastify:response:finish')
 const pathParamsReadCh = channel('datadog:fastify:path-params:finish')
 const finishSetHeaderCh = channel('datadog:fastify:set-header:finish')
 
-const parsingResources = new WeakMap()
+// context management channels
+const preParsingCh = channel('datadog:fastify:pre-parsing:start')
+const preValidationCh = channel('datadog:fastify:pre-validation:start')
+const callbackFinishCh = channel('datadog:fastify:callback:execute')
+
+const parsingContexts = new WeakMap()
 const cookiesPublished = new WeakSet()
 const bodyPublished = new WeakSet()
 
@@ -51,6 +56,7 @@ function wrapAddHook (addHook) {
 
     arguments[arguments.length - 1] = shimmer.wrapFunction(fn, fn => function (request, reply, done) {
       const req = getReq(request)
+      const ctx = { req }
 
       try {
         // done callback is always the last argument
@@ -58,32 +64,26 @@ function wrapAddHook (addHook) {
 
         if (typeof doneCallback === 'function') {
           arguments[arguments.length - 1] = function (err) {
-            publishError(err, req)
+            ctx.error = err
+            publishError(ctx)
 
             const hasCookies = request.cookies && Object.keys(request.cookies).length > 0
 
             if (cookieParserReadCh.hasSubscribers && hasCookies && !cookiesPublished.has(req)) {
-              const res = getRes(reply)
-              const abortController = new AbortController()
-
-              cookieParserReadCh.publish({
-                req,
-                res,
-                abortController,
-                cookies: request.cookies
-              })
+              ctx.res = getRes(reply)
+              ctx.abortController = new AbortController()
+              ctx.cookies = request.cookies
 
+              cookieParserReadCh.publish(ctx)
               cookiesPublished.add(req)
 
-              if (abortController.signal.aborted) return
+              if (ctx.abortController.signal.aborted) return
             }
 
             if (name === 'onRequest' || name === 'preParsing') {
-              const parsingResource = new AsyncResource('bound-anonymous-fn')
-
-              parsingResources.set(req, parsingResource)
+              parsingContexts.set(req, ctx)
 
-              return parsingResource.runInAsyncScope(() => {
+              return callbackFinishCh.runStores(ctx, () => {
                 return doneCallback.apply(this, arguments)
               })
             }
@@ -96,12 +96,16 @@ function wrapAddHook (addHook) {
         const promise = fn.apply(this, arguments)
 
         if (promise && typeof promise.catch === 'function') {
-          return promise.catch(err => publishError(err, req))
+          return promise.catch(err => {
+            ctx.error = err
+            return publishError(ctx)
+          })
         }
 
         return promise
       } catch (e) {
-        throw publishError(e, req)
+        ctx.error = e
+        throw publishError(ctx)
       }
     })
 
@@ -116,7 +120,8 @@ function onRequest (request, reply, done) {
   const res = getRes(reply)
   const routeConfig = getRouteConfig(request)
 
-  handleChannel.publish({ req, res, routeConfig })
+  const ctx = { req, res, routeConfig }
+  handleChannel.publish(ctx)
 
   return done()
 }
@@ -127,18 +132,18 @@ function preHandler (request, reply, done) {
 
   const req = getReq(request)
   const res = getRes(reply)
+  const ctx = { req, res }
 
   const hasBody = request.body && Object.keys(request.body).length > 0
 
   // For multipart/form-data, the body is not available until after preValidation hook
   if (bodyParserReadCh.hasSubscribers && hasBody && !bodyPublished.has(req)) {
-    const abortController = new AbortController()
-
-    bodyParserReadCh.publish({ req, res, body: request.body, abortController })
-
+    ctx.abortController = new AbortController()
+    ctx.body = request.body
+    bodyParserReadCh.publish(ctx)
     bodyPublished.add(req)
 
-    if (abortController.signal.aborted) return
+    if (ctx.abortController.signal.aborted) return
   }
 
   reply.send = wrapSend(reply.send, req)
@@ -149,20 +154,17 @@ function preHandler (request, reply, done) {
 function preValidation (request, reply, done) {
   const req = getReq(request)
   const res = getRes(reply)
-  const parsingResource = parsingResources.get(req)
+  const ctx = parsingContexts.get(req)
+  ctx.res = res
 
   const processInContext = () => {
     let abortController
 
     if (queryParamsReadCh.hasSubscribers && request.query) {
       abortController ??= new AbortController()
-
-      queryParamsReadCh.publish({
-        req,
-        res,
-        abortController,
-        query: request.query
-      })
+      ctx.abortController = abortController
+      ctx.query = request.query
+      queryParamsReadCh.publish(ctx)
 
       if (abortController.signal.aborted) return
     }
@@ -170,8 +172,9 @@ function preValidation (request, reply, done) {
     // Analyze body before schema validation
     if (bodyParserReadCh.hasSubscribers && request.body && !bodyPublished.has(req)) {
       abortController ??= new AbortController()
-
-      bodyParserReadCh.publish({ req, res, body: request.body, abortController })
+      ctx.abortController = abortController
+      ctx.body = request.body
+      bodyParserReadCh.publish(ctx)
 
       bodyPublished.add(req)
 
@@ -180,13 +183,9 @@ function preValidation (request, reply, done) {
 
     if (pathParamsReadCh.hasSubscribers && request.params) {
       abortController ??= new AbortController()
-
-      pathParamsReadCh.publish({
-        req,
-        res,
-        abortController,
-        params: request.params
-      })
+      ctx.abortController = abortController
+      ctx.params = request.params
+      pathParamsReadCh.publish(ctx)
 
       if (abortController.signal.aborted) return
     }
@@ -194,9 +193,9 @@ function preValidation (request, reply, done) {
     done()
   }
 
-  if (!parsingResource) return processInContext()
+  if (!ctx) return processInContext()
 
-  parsingResource.runInAsyncScope(processInContext)
+  preValidationCh.runStores(ctx, processInContext)
 }
 
 function preParsing (request, reply, payload, done) {
@@ -205,19 +204,24 @@ function preParsing (request, reply, payload, done) {
   }
 
   const req = getReq(request)
-  const parsingResource = new AsyncResource('bound-anonymous-fn')
+  const ctx = { req }
+
+  parsingContexts.set(req, ctx)
 
-  parsingResources.set(req, parsingResource)
-  parsingResource.runInAsyncScope(() => done())
+  preParsingCh.runStores(ctx, () => done())
 }
 
 function wrapSend (send, req) {
   return function sendWithTrace (payload) {
+    const ctx = { req }
     if (payload instanceof Error) {
-      errorChannel.publish({ req, error: payload })
+      ctx.error = payload
+      errorChannel.publish(ctx)
     } else if (canPublishResponsePayload(payload)) {
       const res = getRes(this)
-      responsePayloadReadCh.publish({ req, res, body: payload })
+      ctx.res = res
+      ctx.body = payload
+      responsePayloadReadCh.publish(ctx)
     }
 
     return send.apply(this, arguments)
@@ -236,16 +240,17 @@ function getRouteConfig (request) {
   return request?.routeOptions?.config
 }
 
-function publishError (error, req) {
-  if (error) {
-    errorChannel.publish({ error, req })
+function publishError (ctx) {
+  if (ctx.error) {
+    errorChannel.publish(ctx)
   }
 
-  return error
+  return ctx.error
 }
 
 function onRoute (routeOptions) {
-  routeAddedChannel.publish({ routeOptions, onRoute })
+  const ctx = { routeOptions, onRoute }
+  routeAddedChannel.publish(ctx)
 }
 
 // send() payload types: https://fastify.dev/docs/latest/Reference/Reply/#senddata
@@ -282,7 +287,8 @@ function wrapReplyHeader (Reply) {
     const result = header.apply(this, arguments)
 
     if (finishSetHeaderCh.hasSubscribers && key && value) {
-      finishSetHeaderCh.publish({ name: key, value, res: getRes(this) })
+      const ctx = { name: key, value, res: getRes(this) }
+      finishSetHeaderCh.publish(ctx)
     }
 
     return result