dd-trace 4.44.0 → 4.46.0

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

@@ -88,12 +88,10 @@ function wrapMethod (method, path, type, hasPeer) {
     return method
   }
 
-  const wrapped = function () {
+  const wrapped = shimmer.wrapFunction(method, method => function () {
     const args = ensureMetadata(this, arguments, 1)
     return callMethod(this, method, args, path, args[1], type, hasPeer)
-  }
-
-  Object.assign(wrapped, method)
+  })
 
   patched.add(wrapped)
 
@@ -101,7 +99,7 @@ function wrapMethod (method, path, type, hasPeer) {
 }
 
 function wrapCallback (ctx, callback = () => { }) {
-  return function (err) {
+  return shimmer.wrapFunction(callback, callback => function (err) {
     if (err) {
       ctx.error = err
       errorChannel.publish(ctx)
@@ -111,7 +109,7 @@ function wrapCallback (ctx, callback = () => { }) {
       return callback.apply(this, arguments)
       // No async end channel needed
     })
-  }
+  })
 }
 
 function createWrapEmit (ctx, hasPeer = false) {

package/packages/datadog-instrumentations/src/grpc/server.js

@@ -119,7 +119,7 @@ function wrapStream (call, ctx, onCancel) {
 }
 
 function wrapCallback (callback = () => {}, call, ctx, onCancel) {
-  return function (err, value, trailer, flags) {
+  return shimmer.wrapFunction(callback, callback => function (err, value, trailer, flags) {
     if (err) {
       ctx.error = err
       errorChannel.publish(ctx)
@@ -136,7 +136,7 @@ function wrapCallback (callback = () => {}, call, ctx, onCancel) {
       return callback.apply(this, arguments)
       // No async end channel needed
     })
-  }
+  })
 }
 
 function wrapSendStatus (sendStatus, ctx) {

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

@@ -1,12 +1,13 @@
 'use strict'
 
+const tracingChannel = require('dc-polyfill').tracingChannel
 const shimmer = require('../../datadog-shimmer')
-const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+const { addHook, channel } = require('./helpers/instrument')
 
 const handleChannel = channel('apm:hapi:request:handle')
 const routeChannel = channel('apm:hapi:request:route')
 const errorChannel = channel('apm:hapi:request:error')
-const enterChannel = channel('apm:hapi:extension:enter')
+const hapiTracingChannel = tracingChannel('apm:hapi:extension')
 
 function wrapServer (server) {
   return function (options) {
@@ -27,17 +28,17 @@ function wrapServer (server) {
 }
 
 function wrapStart (start) {
-  return function () {
+  return shimmer.wrapFunction(start, start => function () {
     if (this && typeof this.ext === 'function') {
       this.ext('onPreResponse', onPreResponse)
     }
 
     return start.apply(this, arguments)
-  }
+  })
 }
 
 function wrapExt (ext) {
-  return function (events, method, options) {
+  return shimmer.wrapFunction(ext, ext => function (events, method, options) {
     if (events !== null && typeof events === 'object') {
       arguments[0] = wrapEvents(events)
     } else {
@@ -45,7 +46,7 @@ function wrapExt (ext) {
     }
 
     return ext.apply(this, arguments)
-  }
+  })
 }
 
 function wrapDispatch (dispatch) {
@@ -91,19 +92,15 @@ function wrapEvents (events) {
 function wrapHandler (handler) {
   if (typeof handler !== 'function') return handler
 
-  return function (request, h) {
+  return shimmer.wrapFunction(handler, handler => function (request, h) {
     const req = request && request.raw && request.raw.req
 
     if (!req) return handler.apply(this, arguments)
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-    return asyncResource.runInAsyncScope(() => {
-      enterChannel.publish({ req })
-
+    return hapiTracingChannel.traceSync(() => {
       return handler.apply(this, arguments)
     })
-  }
+  })
 }
 
 function onPreResponse (request, h) {

package/packages/datadog-instrumentations/src/helpers/hooks.js

@@ -72,7 +72,6 @@ module.exports = {
   'microgateway-core': () => require('../microgateway-core'),
   mocha: () => require('../mocha'),
   'mocha-each': () => require('../mocha'),
-  workerpool: () => require('../mocha'),
   moleculer: () => require('../moleculer'),
   mongodb: () => require('../mongodb'),
   'mongodb-core': () => require('../mongodb-core'),
@@ -89,6 +88,7 @@ module.exports = {
   'node:http2': () => require('../http2'),
   'node:https': () => require('../http'),
   'node:net': () => require('../net'),
+  nyc: () => require('../nyc'),
   oracledb: () => require('../oracledb'),
   openai: () => require('../openai'),
   paperplane: () => require('../paperplane'),
@@ -113,5 +113,6 @@ module.exports = {
   undici: () => require('../undici'),
   vitest: { esmFirst: true, fn: () => require('../vitest') },
   when: () => require('../when'),
-  winston: () => require('../winston')
+  winston: () => require('../winston'),
+  workerpool: () => require('../mocha')
 }

package/packages/datadog-instrumentations/src/helpers/register.js

@@ -90,8 +90,15 @@ for (const packageName of names) {
       }
 
       if (matchesFile) {
-        const version = moduleVersion || getVersion(moduleBaseDir)
-        if (!Object.hasOwnProperty(namesAndSuccesses, name)) {
+        let version = moduleVersion
+        try {
+          version = version || getVersion(moduleBaseDir)
+        } catch (e) {
+          log.error(`Error getting version for "${name}": ${e.message}`)
+          log.error(e)
+          continue
+        }
+        if (typeof namesAndSuccesses[`${name}@${version}`] === 'undefined') {
           namesAndSuccesses[`${name}@${version}`] = false
         }
 

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

@@ -52,11 +52,11 @@ function patch (http, methodName) {
         let callback = args.callback
 
         if (callback) {
-          callback = function () {
+          callback = shimmer.wrapFunction(args.callback, cb => function () {
             return asyncStartChannel.runStores(ctx, () => {
-              return args.callback.apply(this, arguments)
+              return cb.apply(this, arguments)
             })
-          }
+          })
         }
 
         const options = args.options

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

@@ -12,8 +12,7 @@ const {
   getTestParametersString,
   addEfdStringToTestName,
   removeEfdStringFromTestName,
-  getIsFaultyEarlyFlakeDetection,
-  NUM_FAILED_TEST_RETRIES
+  getIsFaultyEarlyFlakeDetection
 } = require('../../dd-trace/src/plugins/util/test')
 const {
   getFormattedJestTestParameters,
@@ -132,6 +131,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
 
       this.isEarlyFlakeDetectionEnabled = this.testEnvironmentOptions._ddIsEarlyFlakeDetectionEnabled
       this.isFlakyTestRetriesEnabled = this.testEnvironmentOptions._ddIsFlakyTestRetriesEnabled
+      this.flakyTestRetriesCount = this.testEnvironmentOptions._ddFlakyTestRetriesCount
 
       if (this.isEarlyFlakeDetectionEnabled) {
         const hasKnownTests = !!knownTests.jest
@@ -149,7 +149,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
       if (this.isFlakyTestRetriesEnabled) {
         const currentNumRetries = this.global[RETRY_TIMES]
         if (!currentNumRetries) {
-          this.global[RETRY_TIMES] = NUM_FAILED_TEST_RETRIES
+          this.global[RETRY_TIMES] = this.flakyTestRetriesCount
         }
       }
     }
@@ -638,7 +638,7 @@ addHook({
 
 function jestAdapterWrapper (jestAdapter, jestVersion) {
   const adapter = jestAdapter.default ? jestAdapter.default : jestAdapter
-  const newAdapter = shimmer.wrap(adapter, function () {
+  const newAdapter = shimmer.wrapFunction(adapter, adapter => function () {
     const environment = arguments[2]
     if (!environment) {
       return adapter.apply(this, arguments)
@@ -773,6 +773,7 @@ addHook({
       _ddEarlyFlakeDetectionNumRetries,
       _ddRepositoryRoot,
       _ddIsFlakyTestRetriesEnabled,
+      _ddFlakyTestRetriesCount,
       ...restOfTestEnvironmentOptions
     } = testEnvironmentOptions
 

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

@@ -81,8 +81,8 @@ addHook({
 function wrapCallbackWithFinish (callback, finish) {
   if (typeof callback !== 'function') return callback
 
-  return function () {
+  return shimmer.wrapFunction(callback, callback => function () {
     finish()
     callback.apply(this, arguments)
-  }
+  })
 }

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

@@ -71,7 +71,7 @@ function wrapStack (layer) {
 
     middleware = original || middleware
 
-    const handler = shimmer.wrap(middleware, wrapMiddleware(middleware, layer))
+    const handler = shimmer.wrapFunction(middleware, middleware => wrapMiddleware(middleware, layer))
 
     originals.set(handler, middleware)
 
@@ -84,7 +84,7 @@ function wrapMiddleware (fn, layer) {
 
   const name = fn.name
 
-  return function (ctx, next) {
+  return shimmer.wrapFunction(fn, fn => function (ctx, next) {
     if (!ctx || !enterChannel.hasSubscribers) return fn.apply(this, arguments)
 
     const req = ctx.req
@@ -122,7 +122,7 @@ function wrapMiddleware (fn, layer) {
     } finally {
       exitChannel.publish({ req })
     }
-  }
+  })
 }
 
 function fulfill (ctx, error) {
@@ -142,11 +142,11 @@ function fulfill (ctx, error) {
 }
 
 function wrapNext (req, next) {
-  return function () {
+  return shimmer.wrapFunction(next, next => function () {
     nextChannel.publish({ req })
 
     return next.apply(this, arguments)
-  }
+  })
 }
 
 addHook({ name: 'koa', versions: ['>=2'] }, Koa => {

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

@@ -77,7 +77,7 @@ addHook({ name: 'ldapjs', versions: ['>=2'] }, ldapjs => {
     if (callbackIndex > -1) {
       const callback = arguments[callbackIndex]
       // eslint-disable-next-line n/handle-callback-err
-      arguments[callbackIndex] = shimmer.wrap(callback, function (err, corkedEmitter) {
+      arguments[callbackIndex] = shimmer.wrapFunction(callback, callback => function (err, corkedEmitter) {
         if (corkedEmitter !== null && typeof corkedEmitter === 'object' && typeof corkedEmitter.on === 'function') {
           wrapEmitter(corkedEmitter)
         }

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

@@ -11,7 +11,7 @@ const skipCh = channel('apm:mariadb:pool:skip')
 const unskipCh = channel('apm:mariadb:pool:unskip')
 
 function wrapCommandStart (start, callbackResource) {
-  return function () {
+  return shimmer.wrapFunction(start, start => function () {
     if (!startCh.hasSubscribers) return start.apply(this, arguments)
 
     const resolve = callbackResource.bind(this.resolve)
@@ -44,7 +44,7 @@ function wrapCommandStart (start, callbackResource) {
       startCh.publish({ sql: this.sql, conf: this.opts })
       return start.apply(this, arguments)
     })
-  }
+  })
 }
 
 function wrapCommand (Command) {
@@ -98,7 +98,7 @@ function createWrapQueryCallback (options) {
         arguments.length = arguments.length + 1
       }
 
-      arguments[arguments.length - 1] = asyncResource.bind(function (err) {
+      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (err) {
         if (err) {
           errorCh.publish(err)
         }
@@ -108,7 +108,7 @@ function createWrapQueryCallback (options) {
         if (typeof cb === 'function') {
           return callbackResource.runInAsyncScope(() => cb.apply(this, arguments))
         }
-      })
+      }))
 
       return asyncResource.runInAsyncScope(() => {
         startCh.publish({ sql, conf: options })
@@ -119,7 +119,7 @@ function createWrapQueryCallback (options) {
   }
 }
 
-function wrapConnection (Connection, promiseMethod) {
+function wrapConnection (promiseMethod, Connection) {
   return function (options) {
     Connection.apply(this, arguments)
 
@@ -170,13 +170,13 @@ addHook({ name, file: 'lib/pool.js', versions: ['>=3'] }, (Pool) => {
 })
 
 addHook({ name, file: 'lib/connection.js', versions: ['>=2.5.2 <3'] }, (Connection) => {
-  return shimmer.wrap(Connection, wrapConnection(Connection, '_queryPromise'))
+  return shimmer.wrapFunction(Connection, wrapConnection.bind(null, '_queryPromise'))
 })
 
 addHook({ name, file: 'lib/connection.js', versions: ['>=2.0.4 <=2.5.1'] }, (Connection) => {
-  return shimmer.wrap(Connection, wrapConnection(Connection, 'query'))
+  return shimmer.wrapFunction(Connection, wrapConnection.bind(null, 'query'))
 })
 
 addHook({ name, file: 'lib/pool-base.js', versions: ['>=2.0.4 <3'] }, (PoolBase) => {
-  return shimmer.wrap(PoolBase, wrapPoolBase(PoolBase))
+  return shimmer.wrapFunction(PoolBase, wrapPoolBase)
 })

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

@@ -26,14 +26,14 @@ addHook({ name: 'memcached', versions: ['>=2.2'] }, Memcached => {
       const query = queryCompiler.apply(this, arguments)
       const callback = callbackResource.bind(query.callback)
 
-      query.callback = asyncResource.bind(function (err) {
+      query.callback = shimmer.wrapFunction(callback, callback => asyncResource.bind(function (err) {
         if (err) {
           errorCh.publish(err)
         }
         finishCh.publish()
 
         return callback.apply(this, arguments)
-      })
+      }))
       startCh.publish({ client, server, query })
 
       return query

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

@@ -40,7 +40,7 @@ function wrapPluginsFactory (pluginsFactory) {
 }
 
 function wrapNext (req, res, next) {
-  return function nextWithTrace (err) {
+  return shimmer.wrapFunction(next, next => function nextWithTrace (err) {
     const requestResource = requestResources.get(req)
 
     requestResource.runInAsyncScope(() => {
@@ -54,13 +54,13 @@ function wrapNext (req, res, next) {
     })
 
     return next.apply(this, arguments)
-  }
+  })
 }
 
 addHook({ name, versions, file: 'lib/config-proxy-middleware.js' }, configProxyFactory => {
-  return shimmer.wrap(configProxyFactory, wrapConfigProxyFactory(configProxyFactory))
+  return shimmer.wrapFunction(configProxyFactory, wrapConfigProxyFactory)
 })
 
 addHook({ name, versions, file: 'lib/plugins-middleware.js' }, pluginsFactory => {
-  return shimmer.wrap(pluginsFactory, wrapPluginsFactory(pluginsFactory))
+  return shimmer.wrapFunction(pluginsFactory, wrapPluginsFactory)
 })

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

@@ -15,7 +15,7 @@ addHook({
 
   patched.add(mochaEach)
 
-  return shimmer.wrap(mochaEach, function () {
+  return shimmer.wrapFunction(mochaEach, mochaEach => function () {
     const [params] = arguments
     const { it, ...rest } = mochaEach.apply(this, arguments)
     return {

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

@@ -47,6 +47,7 @@ const config = {}
 
 // We'll preserve the original coverage here
 const originalCoverageMap = createCoverageMap()
+let untestedCoverage
 
 // test channels
 const testStartCh = channel('ci:mocha:test:start')
@@ -66,6 +67,8 @@ const testSessionStartCh = channel('ci:mocha:session:start')
 const testSessionFinishCh = channel('ci:mocha:session:finish')
 const itrSkippedSuitesCh = channel('ci:mocha:itr:skipped-suites')
 
+const getCodeCoverageCh = channel('ci:nyc:get-coverage')
+
 function getFilteredSuites (originalSuites) {
   return originalSuites.reduce((acc, suite) => {
     const testPath = getTestSuitePath(suite.file, process.cwd())
@@ -131,6 +134,9 @@ function getOnEndHandler (isParallel) {
     let testCodeCoverageLinesTotal
     if (global.__coverage__) {
       try {
+        if (untestedCoverage) {
+          originalCoverageMap.merge(fromCoverageMapToCoverage(untestedCoverage))
+        }
         testCodeCoverageLinesTotal = originalCoverageMap.getCoverageSummary().lines.pct
       } catch (e) {
         // ignore errors
@@ -153,6 +159,84 @@ function getOnEndHandler (isParallel) {
   })
 }
 
+function getExecutionConfiguration (runner, onFinishRequest) {
+  const mochaRunAsyncResource = new AsyncResource('bound-anonymous-fn')
+
+  const onReceivedSkippableSuites = ({ err, skippableSuites, itrCorrelationId: responseItrCorrelationId }) => {
+    if (err) {
+      suitesToSkip = []
+    } else {
+      suitesToSkip = skippableSuites
+      itrCorrelationId = responseItrCorrelationId
+    }
+    // We remove the suites that we skip through ITR
+    const filteredSuites = getFilteredSuites(runner.suite.suites)
+    const { suitesToRun } = filteredSuites
+
+    isSuitesSkipped = suitesToRun.length !== runner.suite.suites.length
+
+    log.debug(
+      () => `${suitesToRun.length} out of ${runner.suite.suites.length} suites are going to run.`
+    )
+
+    runner.suite.suites = suitesToRun
+
+    skippedSuites = Array.from(filteredSuites.skippedSuites)
+
+    onFinishRequest()
+  }
+
+  const onReceivedKnownTests = ({ err, knownTests: receivedKnownTests }) => {
+    if (err) {
+      knownTests = []
+      isEarlyFlakeDetectionEnabled = false
+    } else {
+      knownTests = receivedKnownTests
+    }
+
+    if (isSuitesSkippingEnabled) {
+      skippableSuitesCh.publish({
+        onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
+      })
+    } else {
+      onFinishRequest()
+    }
+  }
+
+  const onReceivedConfiguration = ({ err, libraryConfig }) => {
+    if (err || !skippableSuitesCh.hasSubscribers || !knownTestsCh.hasSubscribers) {
+      return onFinishRequest()
+    }
+
+    isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
+    isSuitesSkippingEnabled = libraryConfig.isSuitesSkippingEnabled
+    earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
+    isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
+
+    config.isEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
+    config.isSuitesSkippingEnabled = isSuitesSkippingEnabled
+    config.earlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
+    config.isFlakyTestRetriesEnabled = isFlakyTestRetriesEnabled
+    config.flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
+
+    if (isEarlyFlakeDetectionEnabled) {
+      knownTestsCh.publish({
+        onDone: mochaRunAsyncResource.bind(onReceivedKnownTests)
+      })
+    } else if (isSuitesSkippingEnabled) {
+      skippableSuitesCh.publish({
+        onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
+      })
+    } else {
+      onFinishRequest()
+    }
+  }
+
+  libraryConfigurationCh.publish({
+    onDone: mochaRunAsyncResource.bind(onReceivedConfiguration)
+  })
+}
+
 // In this hook we delay the execution with options.delay to grab library configuration,
 // skippable and known tests.
 // It is called but skipped in parallel mode.
@@ -161,7 +245,6 @@ addHook({
   versions: ['>=5.2.0'],
   file: 'lib/mocha.js'
 }, (Mocha) => {
-  const mochaRunAsyncResource = new AsyncResource('bound-anonymous-fn')
   shimmer.wrap(Mocha.prototype, 'run', run => function () {
     // Workers do not need to request any data, just run the tests
     if (!testStartCh.hasSubscribers || process.env.MOCHA_WORKER_ID || this.options.parallel) {
@@ -181,79 +264,17 @@ addHook({
       }
     })
 
-    const onReceivedSkippableSuites = ({ err, skippableSuites, itrCorrelationId: responseItrCorrelationId }) => {
-      if (err) {
-        suitesToSkip = []
-      } else {
-        suitesToSkip = skippableSuites
-        itrCorrelationId = responseItrCorrelationId
-      }
-      // We remove the suites that we skip through ITR
-      const filteredSuites = getFilteredSuites(runner.suite.suites)
-      const { suitesToRun } = filteredSuites
-
-      isSuitesSkipped = suitesToRun.length !== runner.suite.suites.length
-
-      log.debug(
-        () => `${suitesToRun.length} out of ${runner.suite.suites.length} suites are going to run.`
-      )
-
-      runner.suite.suites = suitesToRun
-
-      skippedSuites = Array.from(filteredSuites.skippedSuites)
-
-      global.run()
-    }
-
-    const onReceivedKnownTests = ({ err, knownTests: receivedKnownTests }) => {
-      if (err) {
-        knownTests = []
-        isEarlyFlakeDetectionEnabled = false
-      } else {
-        knownTests = receivedKnownTests
-      }
-
-      if (isSuitesSkippingEnabled) {
-        skippableSuitesCh.publish({
-          onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
+    getExecutionConfiguration(runner, () => {
+      if (getCodeCoverageCh.hasSubscribers) {
+        getCodeCoverageCh.publish({
+          onDone: (receivedCodeCoverage) => {
+            untestedCoverage = receivedCodeCoverage
+            global.run()
+          }
         })
       } else {
         global.run()
       }
-    }
-
-    const onReceivedConfiguration = ({ err, libraryConfig }) => {
-      if (err || !skippableSuitesCh.hasSubscribers || !knownTestsCh.hasSubscribers) {
-        return global.run()
-      }
-
-      isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
-      isSuitesSkippingEnabled = libraryConfig.isSuitesSkippingEnabled
-      earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
-      isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
-
-      config.isEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
-      config.isSuitesSkippingEnabled = isSuitesSkippingEnabled
-      config.earlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
-      config.isFlakyTestRetriesEnabled = isFlakyTestRetriesEnabled
-
-      if (isEarlyFlakeDetectionEnabled) {
-        knownTestsCh.publish({
-          onDone: mochaRunAsyncResource.bind(onReceivedKnownTests)
-        })
-      } else if (isSuitesSkippingEnabled) {
-        skippableSuitesCh.publish({
-          onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
-        })
-      } else {
-        global.run()
-      }
-    }
-
-    mochaRunAsyncResource.runInAsyncScope(() => {
-      libraryConfigurationCh.publish({
-        onDone: mochaRunAsyncResource.bind(onReceivedConfiguration)
-      })
     })
 
     return runner

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

@@ -3,8 +3,7 @@
 const {
   getTestSuitePath,
   removeEfdStringFromTestName,
-  addEfdStringToTestName,
-  NUM_FAILED_TEST_RETRIES
+  addEfdStringToTestName
 } = require('../../../dd-trace/src/plugins/util/test')
 const { channel, AsyncResource } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
@@ -114,7 +113,7 @@ function runnableWrapper (RunnablePackage, libraryConfig) {
     }
     // Flaky test retries does not work in parallel mode
     if (libraryConfig?.isFlakyTestRetriesEnabled) {
-      this.retries(NUM_FAILED_TEST_RETRIES)
+      this.retries(libraryConfig?.flakyTestRetriesCount)
     }
     // The reason why the wrapping logic is here is because we need to cover
     // `afterEach` and `beforeEach` hooks as well.

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

@@ -3,3 +3,7 @@ if (process.env.MOCHA_WORKER_ID) {
 } else {
   require('./mocha/main')
 }
+
+// TODO add appropriate calls to wrapFunction whenever we're adding a callback
+// wrapper. Right now this is less of an issue since that only has effect in
+// SSI, where CI Vis isn't supported.

package/packages/datadog-instrumentations/src/moleculer/server.js

@@ -24,7 +24,7 @@ function createMiddleware () {
     localAction (next, action) {
       const broker = this
 
-      return function datadogMiddleware (ctx) {
+      return shimmer.wrapFunction(next, next => function datadogMiddleware (ctx) {
         const actionResource = new AsyncResource('bound-anonymous-fn')
 
         return actionResource.runInAsyncScope(() => {
@@ -47,7 +47,7 @@ function createMiddleware () {
             finishChannel.publish()
           }
         })
-      }
+      })
     }
   }
 }