dd-trace 5.21.0 → 5.22.0

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

@@ -1,14 +1,10 @@
 'use strict'
 
-const {
-  channel,
-  addHook
-} = require('./helpers/instrument')
+const { addHook } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
-const startCh = channel('apm:openai:request:start')
-const finishCh = channel('apm:openai:request:finish')
-const errorCh = channel('apm:openai:request:error')
+const tracingChannel = require('dc-polyfill').tracingChannel
+const ch = tracingChannel('apm:openai:request')
 
 const V4_PACKAGE_SHIMS = [
   {
@@ -110,33 +106,18 @@ addHook({ name: 'openai', file: 'dist/api.js', versions: ['>=3.0.0 <4'] }, expor
 
   for (const methodName of methodNames) {
     shimmer.wrap(exports.OpenAIApi.prototype, methodName, fn => function () {
-      if (!startCh.hasSubscribers) {
+      if (!ch.start.hasSubscribers) {
         return fn.apply(this, arguments)
       }
 
-      startCh.publish({
+      const ctx = {
         methodName,
         args: arguments,
         basePath: this.basePath,
         apiKey: this.configuration.apiKey
-      })
-
-      return fn.apply(this, arguments)
-        .then((response) => {
-          finish({
-            headers: response.headers,
-            body: response.data,
-            path: response.request.path,
-            method: response.request.method
-          })
-
-          return response
-        })
-        .catch(error => {
-          finish(undefined, error)
+      }
 
-          throw error
-        })
+      return ch.tracePromise(fn, ctx, this, ...arguments)
     })
   }
 
@@ -213,7 +194,7 @@ function convertBufferstoObjects (chunks = []) {
  * the chunks, and let the combined content be the final response.
  * This way, spans look the same as when not streamed.
  */
-function wrapStreamIterator (response, options, n) {
+function wrapStreamIterator (response, options, n, ctx) {
   let processChunksAsBuffers = false
   let chunks = []
   return function (itr) {
@@ -253,11 +234,13 @@ function wrapStreamIterator (response, options, n) {
                 }
               }
 
-              finish({
+              finish(ctx, {
                 headers: response.headers,
-                body,
-                path: response.url,
-                method: options.method
+                data: body,
+                request: {
+                  path: response.url,
+                  method: options.method
+                }
               })
             }
 
@@ -281,7 +264,7 @@ for (const shim of V4_PACKAGE_SHIMS) {
 
     for (const methodName of methods) {
       shimmer.wrap(targetPrototype, methodName, methodFn => function () {
-        if (!startCh.hasSubscribers) {
+        if (!ch.start.hasSubscribers) {
           return methodFn.apply(this, arguments)
         }
 
@@ -303,66 +286,72 @@ for (const shim of V4_PACKAGE_SHIMS) {
 
         const client = this._client || this.client
 
-        startCh.publish({
+        const ctx = {
           methodName: `${baseResource}.${methodName}`,
           args: arguments,
           basePath: client.baseURL,
           apiKey: client.apiKey
-        })
+        }
 
-        const apiProm = methodFn.apply(this, arguments)
+        return ch.start.runStores(ctx, () => {
+          const apiProm = methodFn.apply(this, arguments)
 
-        // wrapping `parse` avoids problematic wrapping of `then` when trying to call
-        // `withResponse` in userland code after. This way, we can return the whole `APIPromise`
-        shimmer.wrap(apiProm, 'parse', origApiPromParse => function () {
-          return origApiPromParse.apply(this, arguments)
+          // wrapping `parse` avoids problematic wrapping of `then` when trying to call
+          // `withResponse` in userland code after. This way, we can return the whole `APIPromise`
+          shimmer.wrap(apiProm, 'parse', origApiPromParse => function () {
+            return origApiPromParse.apply(this, arguments)
             // the original response is wrapped in a promise, so we need to unwrap it
-            .then(body => Promise.all([this.responsePromise, body]))
-            .then(([{ response, options }, body]) => {
-              if (stream) {
-                if (body.iterator) {
-                  shimmer.wrap(body, 'iterator', wrapStreamIterator(response, options, n))
+              .then(body => Promise.all([this.responsePromise, body]))
+              .then(([{ response, options }, body]) => {
+                if (stream) {
+                  if (body.iterator) {
+                    shimmer.wrap(body, 'iterator', wrapStreamIterator(response, options, n, ctx))
+                  } else {
+                    shimmer.wrap(
+                      body.response.body, Symbol.asyncIterator, wrapStreamIterator(response, options, n, ctx)
+                    )
+                  }
                 } else {
-                  shimmer.wrap(
-                    body.response.body, Symbol.asyncIterator, wrapStreamIterator(response, options, n)
-                  )
+                  finish(ctx, {
+                    headers: response.headers,
+                    data: body,
+                    request: {
+                      path: response.url,
+                      method: options.method
+                    }
+                  })
                 }
-              } else {
-                finish({
-                  headers: response.headers,
-                  body,
-                  path: response.url,
-                  method: options.method
-                })
-              }
 
-              return body
-            })
-            .catch(error => {
-              finish(undefined, error)
+                return body
+              })
+              .catch(error => {
+                finish(ctx, undefined, error)
 
-              throw error
-            })
-            .finally(() => {
+                throw error
+              })
+              .finally(() => {
               // maybe we don't want to unwrap here in case the promise is re-used?
               // other hand: we want to avoid resource leakage
-              shimmer.unwrap(apiProm, 'parse')
-            })
-        })
+                shimmer.unwrap(apiProm, 'parse')
+              })
+          })
 
-        return apiProm
+          return apiProm
+        })
       })
     }
     return exports
   })
 }
 
-function finish (response, error) {
+function finish (ctx, response, error) {
   if (error) {
-    errorCh.publish({ error })
+    ctx.error = error
+    ch.error.publish(ctx)
   }
 
-  finishCh.publish(response)
+  ctx.result = response
+  ch.asyncEnd.publish(ctx)
 }
 
 function getOption (args, option, defaultValue) {

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

@@ -31,10 +31,10 @@ addHook({ name: 'oracledb', versions: ['>=5'] }, oracledb => {
       if (arguments.length && typeof arguments[arguments.length - 1] === 'function') {
         const cb = arguments[arguments.length - 1]
         const outerAr = new AsyncResource('apm:oracledb:outer-scope')
-        arguments[arguments.length - 1] = function wrappedCb (err, result) {
+        arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => function wrappedCb (err, result) {
           finish(err)
           return outerAr.runInAsyncScope(() => cb.apply(this, arguments))
-        }
+        })
       }
 
       return new AsyncResource('apm:oracledb:inner-scope').runInAsyncScope(() => {
@@ -67,12 +67,12 @@ addHook({ name: 'oracledb', versions: ['>=5'] }, oracledb => {
   shimmer.wrap(oracledb, 'getConnection', getConnection => {
     return function wrappedGetConnection (connAttrs, callback) {
       if (callback) {
-        arguments[1] = (err, connection) => {
+        arguments[1] = shimmer.wrapFunction(callback, callback => (err, connection) => {
           if (connection) {
             connectionAttributes.set(connection, connAttrs)
           }
           callback(err, connection)
-        }
+        })
 
         getConnection.apply(this, arguments)
       } else {
@@ -86,12 +86,12 @@ addHook({ name: 'oracledb', versions: ['>=5'] }, oracledb => {
   shimmer.wrap(oracledb, 'createPool', createPool => {
     return function wrappedCreatePool (poolAttrs, callback) {
       if (callback) {
-        arguments[1] = (err, pool) => {
+        arguments[1] = shimmer.wrapFunction(callback, callback => (err, pool) => {
           if (pool) {
             poolAttributes.set(pool, poolAttrs)
           }
           callback(err, pool)
-        }
+        })
 
         createPool.apply(this, arguments)
       } else {
@@ -109,12 +109,12 @@ addHook({ name: 'oracledb', versions: ['>=5'] }, oracledb => {
         callback = arguments[arguments.length - 1]
       }
       if (callback) {
-        arguments[arguments.length - 1] = (err, connection) => {
+        arguments[arguments.length - 1] = shimmer.wrapFunction(callback, callback => (err, connection) => {
           if (connection) {
             connectionAttributes.set(connection, poolAttributes.get(this))
           }
           callback(err, connection)
-        }
+        })
         getConnection.apply(this, arguments)
       } else {
         return getConnection.apply(this, arguments).then((connection) => {

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

@@ -9,7 +9,7 @@ addHook({
   file: 'lib/passport-http/strategies/basic.js',
   versions: ['>=0.3.0']
 }, BasicStrategy => {
-  return shimmer.wrap(BasicStrategy, function () {
+  return shimmer.wrapFunction(BasicStrategy, BasicStrategy => function () {
     const type = 'http'
 
     if (typeof arguments[0] === 'function') {

package/packages/datadog-instrumentations/src/passport-local.js

@@ -9,7 +9,7 @@ addHook({
   file: 'lib/strategy.js',
   versions: ['>=1.0.0']
 }, Strategy => {
-  return shimmer.wrap(Strategy, function () {
+  return shimmer.wrapFunction(Strategy, Strategy => function () {
     const type = 'local'
 
     if (typeof arguments[0] === 'function') {

package/packages/datadog-instrumentations/src/passport-utils.js

@@ -11,7 +11,7 @@ function wrapVerifiedAndPublish (username, password, verified, type) {
   }
 
   // eslint-disable-next-line n/handle-callback-err
-  return shimmer.wrap(verified, function (err, user, info) {
+  return shimmer.wrapFunction(verified, verified => function (err, user, info) {
     const credentials = { type, username }
     passportVerifyChannel.publish({ credentials, user })
     return verified.apply(this, arguments)

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

@@ -122,7 +122,7 @@ function wrapPoolQuery (query) {
 
       const cb = arguments[arguments.length - 1]
       if (typeof cb === 'function') {
-        arguments[arguments.length - 1] = shimmer.wrap(cb, function () {
+        arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => function () {
           finish()
           return cb.apply(this, arguments)
         })

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

@@ -76,19 +76,19 @@ function wrapPrettyFactory (prettyFactory) {
 addHook({ name: 'pino', versions: ['2 - 3', '4', '>=5 <5.14.0'] }, pino => {
   const asJsonSym = (pino.symbols && pino.symbols.asJsonSym) || 'asJson'
 
-  return shimmer.wrap(pino, wrapPino(asJsonSym, wrapAsJson, pino))
+  return shimmer.wrapFunction(pino, pino => wrapPino(asJsonSym, wrapAsJson, pino))
 })
 
 addHook({ name: 'pino', versions: ['>=5.14.0 <6.8.0'] }, pino => {
   const mixinSym = pino.symbols.mixinSym
 
-  return shimmer.wrap(pino, wrapPino(mixinSym, wrapMixin, pino))
+  return shimmer.wrapFunction(pino, pino => wrapPino(mixinSym, wrapMixin, pino))
 })
 
 addHook({ name: 'pino', versions: ['>=6.8.0'] }, pino => {
   const mixinSym = pino.symbols.mixinSym
 
-  const wrapped = shimmer.wrap(pino, wrapPino(mixinSym, wrapMixin, pino))
+  const wrapped = shimmer.wrapFunction(pino, pino => wrapPino(mixinSym, wrapMixin, pino))
   wrapped.pino = wrapped
   wrapped.default = wrapped
 
@@ -101,5 +101,5 @@ addHook({ name: 'pino-pretty', file: 'lib/utils.js', versions: ['>=3'] }, utils
 })
 
 addHook({ name: 'pino-pretty', versions: ['1 - 2'] }, prettyFactory => {
-  return shimmer.wrap(prettyFactory, wrapPrettyFactory(prettyFactory))
+  return shimmer.wrapFunction(prettyFactory, wrapPrettyFactory)
 })

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

@@ -2,7 +2,7 @@ const semver = require('semver')
 
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
-const { parseAnnotations, getTestSuitePath, NUM_FAILED_TEST_RETRIES } = require('../../dd-trace/src/plugins/util/test')
+const { parseAnnotations, getTestSuitePath } = require('../../dd-trace/src/plugins/util/test')
 const log = require('../../dd-trace/src/log')
 
 const testStartCh = channel('ci:playwright:test:start')
@@ -36,8 +36,9 @@ const STATUS_TO_TEST_STATUS = {
 
 let remainingTestsByFile = {}
 let isEarlyFlakeDetectionEnabled = false
-let isFlakyTestRetriesEnabled = false
 let earlyFlakeDetectionNumRetries = 0
+let isFlakyTestRetriesEnabled = false
+let flakyTestRetriesCount = 0
 let knownTests = {}
 let rootDir = ''
 const MINIMUM_SUPPORTED_VERSION_EFD = '1.38.0'
@@ -407,6 +408,7 @@ function runnerHook (runnerExport, playwrightVersion) {
         isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
         earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
         isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
+        flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
       }
     } catch (e) {
       isEarlyFlakeDetectionEnabled = false
@@ -429,10 +431,10 @@ function runnerHook (runnerExport, playwrightVersion) {
 
     const projects = getProjectsFromRunner(this)
 
-    if (isFlakyTestRetriesEnabled) {
+    if (isFlakyTestRetriesEnabled && flakyTestRetriesCount > 0) {
       projects.forEach(project => {
         if (project.retries === 0) { // Only if it hasn't been set by the user
-          project.retries = NUM_FAILED_TEST_RETRIES
+          project.retries = flakyTestRetriesCount
         }
       })
     }

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

@@ -156,12 +156,12 @@ function start (client, command, args, url = {}) {
 }
 
 function wrapCallback (finishCh, errorCh, callback) {
-  return function (err) {
+  return shimmer.wrapFunction(callback, callback => function (err) {
     finish(finishCh, errorCh, err)
     if (callback) {
       return callback.apply(this, arguments)
     }
-  }
+  })
 }
 
 function finish (finishCh, errorCh, error) {

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

@@ -40,7 +40,7 @@ function wrapMiddleware (middleware) {
 function wrapFn (fn) {
   if (Array.isArray(fn)) return wrapMiddleware(fn)
 
-  return function (req, res, next) {
+  return shimmer.wrapFunction(fn, fn => function (req, res, next) {
     if (typeof next === 'function') {
       arguments[2] = wrapNext(req, next)
     }
@@ -72,16 +72,16 @@ function wrapFn (fn) {
     } finally {
       exitChannel.publish({ req })
     }
-  }
+  })
 }
 
 function wrapNext (req, next) {
-  return function () {
+  return shimmer.wrapFunction(next, next => function () {
     nextChannel.publish({ req })
     finishChannel.publish({ req })
 
     next.apply(this, arguments)
-  }
+  })
 }
 
 addHook({ name: 'restify', versions: ['>=3'], file: 'lib/server.js' }, Server => {

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

@@ -153,11 +153,11 @@ function wrapDeliveryUpdate (obj, update) {
   const asyncResource = context.asyncResource
   if (obj && asyncResource) {
     const cb = asyncResource.bind(update)
-    return AsyncResource.bind(function wrappedUpdate (settled, stateData) {
+    return shimmer.wrapFunction(cb, cb => AsyncResource.bind(function wrappedUpdate (settled, stateData) {
       const state = getStateFromData(stateData)
       dispatchReceiveCh.publish({ state })
       return cb.apply(this, arguments)
-    })
+    }))
   }
   return function wrappedUpdate (settled, stateData) {
     return update.apply(this, arguments)
@@ -178,7 +178,7 @@ function patchCircularBuffer (proto, Session) {
         }
         if (CircularBuffer && !patched.has(CircularBuffer.prototype)) {
           shimmer.wrap(CircularBuffer.prototype, 'pop_if', popIf => function (fn) {
-            arguments[0] = AsyncResource.bind(function (entry) {
+            arguments[0] = shimmer.wrapFunction(fn, fn => AsyncResource.bind(function (entry) {
               const context = contexts.get(entry)
               const asyncResource = context && context.asyncResource
 
@@ -198,7 +198,7 @@ function patchCircularBuffer (proto, Session) {
               }
 
               return shouldPop
-            })
+            }))
             return popIf.apply(this, arguments)
           })
           patched.add(CircularBuffer.prototype)

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

@@ -18,7 +18,7 @@ function createWrapRouterMethod (name) {
   function wrapLayerHandle (layer, original) {
     original._name = original._name || layer.name
 
-    const handle = shimmer.wrap(original, function () {
+    const handle = shimmer.wrapFunction(original, original => function () {
       if (!enterChannel.hasSubscribers) return original.apply(this, arguments)
 
       const matchers = layerMatchers.get(layer)
@@ -89,7 +89,7 @@ function createWrapRouterMethod (name) {
   }
 
   function wrapNext (req, next) {
-    return function (error) {
+    return shimmer.wrapFunction(next, next => function (error) {
       if (error && error !== 'route' && error !== 'router') {
         errorChannel.publish({ req, error })
       }
@@ -98,7 +98,7 @@ function createWrapRouterMethod (name) {
       finishChannel.publish({ req })
 
       next.apply(this, arguments)
-    }
+    })
   }
 
   function extractMatchers (fn) {
@@ -151,7 +151,7 @@ function createWrapRouterMethod (name) {
   }
 
   function wrapMethod (original) {
-    return function methodWithTrace (fn) {
+    return shimmer.wrapFunction(original, original => function methodWithTrace (fn) {
       const offset = this.stack ? [].concat(this.stack).length : 0
       const router = original.apply(this, arguments)
 
@@ -162,7 +162,7 @@ function createWrapRouterMethod (name) {
       wrapStack(this.stack, offset, extractMatchers(fn))
 
       return router
-    }
+    })
   }
 
   return wrapMethod

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

@@ -48,14 +48,14 @@ addHook({ name: 'sharedb', versions: ['>=1'], file: 'lib/agent.js' }, Agent => {
 
       callback = callbackResource.bind(callback)
 
-      arguments[1] = asyncResource.bind(function (error, res) {
+      arguments[1] = shimmer.wrapFunction(callback, callback => asyncResource.bind(function (error, res) {
         if (error) {
           errorCh.publish(error)
         }
         finishCh.publish({ request, res })
 
         return callback.apply(this, arguments)
-      })
+      }))
 
       try {
         return origHandleMessageFn.apply(this, arguments)

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

@@ -1,6 +1,5 @@
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
-const { NUM_FAILED_TEST_RETRIES } = require('../../dd-trace/src/plugins/util/test')
 
 // test hooks
 const testStartCh = channel('ci:vitest:test:start')
@@ -108,17 +107,19 @@ function getSortWrapper (sort) {
     // So we will use the sort from BaseSequencer. This means that a custom sequencer
     // will not work. This will be a known limitation.
     let isFlakyTestRetriesEnabled = false
+    let flakyTestRetriesCount = 0
 
     try {
       const { err, libraryConfig } = await getChannelPromise(libraryConfigurationCh)
       if (!err) {
         isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
+        flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
       }
     } catch (e) {
       isFlakyTestRetriesEnabled = false
     }
-    if (isFlakyTestRetriesEnabled && !this.ctx.config.retry) {
-      this.ctx.config.retry = NUM_FAILED_TEST_RETRIES
+    if (isFlakyTestRetriesEnabled && !this.ctx.config.retry && flakyTestRetriesCount > 0) {
+      this.ctx.config.retry = flakyTestRetriesCount
     }
 
     let testCodeCoverageLinesTotal

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

@@ -42,13 +42,12 @@ function wrapMethod (method, logCh) {
 
         if (patched.has(transport) || typeof transport.log !== 'function') continue
 
-        const log = transport.log
-        transport.log = function wrappedLog (level, msg, meta, callback) {
+        shimmer.wrap(transport, 'log', log => function wrappedLog (level, msg, meta, callback) {
           const payload = { message: meta || {} }
           logCh.publish(payload)
           arguments[2] = payload.message
           log.apply(this, arguments)
-        }
+        })
         patched.add(transport)
       }
     }

package/packages/datadog-plugin-cypress/src/cypress-plugin.js

@@ -28,8 +28,7 @@ const {
   TEST_SOURCE_FILE,
   TEST_IS_NEW,
   TEST_IS_RETRY,
-  TEST_EARLY_FLAKE_ENABLED,
-  NUM_FAILED_TEST_RETRIES
+  TEST_EARLY_FLAKE_ENABLED
 } = require('../../dd-trace/src/plugins/util/test')
 const { isMarkedAsUnskippable } = require('../../datadog-plugin-jest/src/util')
 const { ORIGIN_KEY, COMPONENT } = require('../../dd-trace/src/constants')
@@ -229,16 +228,16 @@ class CypressPlugin {
               isCodeCoverageEnabled,
               isEarlyFlakeDetectionEnabled,
               earlyFlakeDetectionNumRetries,
-              isFlakyTestRetriesEnabled
+              isFlakyTestRetriesEnabled,
+              flakyTestRetriesCount
             }
           } = libraryConfigurationResponse
           this.isSuitesSkippingEnabled = isSuitesSkippingEnabled
           this.isCodeCoverageEnabled = isCodeCoverageEnabled
           this.isEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
           this.earlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
-          this.isFlakyTestRetriesEnabled = isFlakyTestRetriesEnabled
-          if (this.isFlakyTestRetriesEnabled) {
-            this.cypressConfig.retries.runMode = NUM_FAILED_TEST_RETRIES
+          if (isFlakyTestRetriesEnabled) {
+            this.cypressConfig.retries.runMode = flakyTestRetriesCount
           }
         }
         return this.cypressConfig

package/packages/datadog-plugin-hapi/src/index.js

@@ -32,8 +32,8 @@ class HapiPlugin extends RouterPlugin {
       }
     })
 
-    this.addSub('apm:hapi:extension:enter', ({ req }) => {
-      this.enter(this._requestSpans.get(req))
+    this.addBind('apm:hapi:extension:start', ({ req }) => {
+      return this._requestSpans.get(req)
     })
   }
 }

package/packages/datadog-plugin-jest/src/index.js

@@ -154,6 +154,7 @@ class JestPlugin extends CiPlugin {
         config._ddEarlyFlakeDetectionNumRetries = this.libraryConfig?.earlyFlakeDetectionNumRetries ?? 0
         config._ddRepositoryRoot = this.repositoryRoot
         config._ddIsFlakyTestRetriesEnabled = this.libraryConfig?.isFlakyTestRetriesEnabled ?? false
+        config._ddFlakyTestRetriesCount = this.libraryConfig?.flakyTestRetriesCount
       })
     })