dd-trace 6.13.0 → 6.14.0

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

@@ -18,6 +18,8 @@ const {
   getEfdRetryCountForDuration,
   hasEfdRetries,
 } = require('../../dd-trace/src/ci-visibility/efd-retry-policy')
+const { FINAL_FLUSH_FALLBACK_DELAY, FINAL_FLUSH_TIMEOUT } =
+  require('../../dd-trace/src/ci-visibility/final-flush')
 const {
   getCoveredFilesFromCoverage,
   JEST_WORKER_TRACE_PAYLOAD_CODE,
@@ -95,10 +97,10 @@ const itrSkippedSuitesCh = channel('ci:jest:itr:skipped-suites')
 // https://github.com/jestjs/jest/blob/1d682f21c7a35da4d3ab3a1436a357b980ebd0fa/packages/jest-worker/src/types.ts#L37
 const CHILD_MESSAGE_CALL = 1
 
-// Maximum time we'll wait for the tracer to flush
-// The exporter has a 10-second bounded final-flush deadline. Leave enough time
-// for its completion callback before Jest's --forceExit fallback takes over.
-const FLUSH_TIMEOUT = 12_000
+// Maximum time we'll wait for the tracer to flush.
+// Let the exporter enforce its hard deadline and invoke its completion callback
+// before Jest's --forceExit fallback takes over.
+const FLUSH_TIMEOUT = FINAL_FLUSH_TIMEOUT + FINAL_FLUSH_FALLBACK_DELAY + 5000
 const JEST_SESSION_STATE = Symbol.for('dd-trace:jest:session')
 const JEST_BAIL_REPORTER_PATH = require.resolve('./jest/bail-reporter')
 const DD_JEST_HANDLE_TEST_EVENT_WRAPPED = Symbol('dd-trace:jest:handle-test-event-wrapped')
@@ -342,16 +344,17 @@ function getTestStats (testStatuses) {
 function formatIgnoredFailuresSummary (ignoredFailures) {
   if (!ignoredFailures?.efdFailureCount) return ''
 
-  const items = ignoredFailures.efdNames.map(text => ({ text, suffix: 'Early Flake Detection' }))
-
-  if (items.length === 0) return ''
-
-  const shown = items.slice(0, MAX_IGNORED_TEST_NAMES)
-  const more = items.length - shown.length
+  const shown = ignoredFailures.efdNames.slice(0, MAX_IGNORED_TEST_NAMES)
+  const more = ignoredFailures.efdNames.length - shown.length
   const moreSuffix = more > 0 ? `\n  ... and ${more} more` : ''
-  const formattedItems = shown
-    .map(({ text, suffix }) => `  • ${text}${suffix ? ` (${suffix})` : ''}`)
-    .join('\n') + moreSuffix
+  let formattedItems = ''
+  let isFirstItem = true
+  for (const text of shown) {
+    if (!isFirstItem) formattedItems += '\n'
+    formattedItems += `  • ${text} (Early Flake Detection)`
+    isFirstItem = false
+  }
+  formattedItems += moreSuffix
 
   return `${ignoredFailures.efdFailureCount} test failure(s) were ignored. Exit code set to 0.\n\n${formattedItems}`
 }
@@ -764,7 +767,8 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
             }
             return test.call(this, testName, testFn, ...callArgs)
           }
-          const eachBind = bind(eachTest).apply(this, eachArgs)
+          const boundTest = bind(eachTest)
+          const eachBind = boundTest.apply(this, eachArgs)
           return eachBind.apply(this, testArgs)
         }
       }
@@ -877,7 +881,8 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
             }
             return concurrentTest.call(this, testName, testFn, ...callArgs)
           }
-          const eachBind = bind(concurrentEachTest, false, needsEachError).apply(this, eachArgs)
+          const boundTest = bind(concurrentEachTest, false, needsEachError)
+          const eachBind = boundTest.apply(this, eachArgs)
           return eachBind.apply(this, testArgs)
         }
       }
@@ -926,10 +931,12 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
 
       let jestEach
       try {
-        jestEach = createRequire(path.join(this.rootDir, 'package.json'))('jest-each')
+        const requireFromPath = createRequire(path.join(this.rootDir, 'package.json'))
+        jestEach = requireFromPath('jest-each')
       } catch {
         try {
-          jestEach = createRequire(path.join(process.cwd(), 'package.json'))('jest-each')
+          const requireFromPath = createRequire(path.join(process.cwd(), 'package.json'))
+          jestEach = requireFromPath('jest-each')
         } catch {
           jestEach = undefined
         }
@@ -3067,9 +3074,7 @@ function getCliWrapper (isNewJestVersion) {
         result = await runCLI.apply(this, arguments)
       } catch (error) {
         try {
-          await waitForTestSessionFinish(getTestSessionFinishPayload('fail', error, {
-            isTestSessionFinalizationError: true,
-          }))
+          await waitForTestSessionFinish(getTestSessionFinishPayload('fail', error))
         } catch (finalizationError) {
           log.error('Jest test session finalization error: %s', finalizationError)
         }

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

@@ -3,6 +3,11 @@
 const shimmer = require('../../datadog-shimmer')
 const { addHook, channel } = require('./helpers/instrument')
 const { wrapThen } = require('./helpers/promise')
+const {
+  wrapPoolRelease,
+  wrapPromisePoolAcquire,
+  wrapPromisePoolQueryMethod,
+} = require('./helpers/pool-acquire')
 
 const startRawQueryCh = channel('datadog:knex:raw:start')
 const rawQuerySubscribes = channel('datadog:knex:raw:subscribes')
@@ -32,6 +37,14 @@ addHook({
   versions: ['>=2'],
   file: 'lib/knex-builder/Knex.js',
 }, Knex => {
+  shimmer.wrap(Knex.Client.prototype, 'acquireConnection', acquireConnection => wrapPromisePoolAcquire(
+    acquireConnection,
+    resolveKnexDriver,
+    resolveKnexConfig,
+    knexHasIdleConnection
+  ))
+  shimmer.wrap(Knex.Client.prototype, 'releaseConnection', wrapPoolRelease)
+
   shimmer.wrap(Knex.Client.prototype, 'raw', raw => function (...args) {
     if (!startRawQueryCh.hasSubscribers) {
       return raw.apply(this, args)
@@ -77,6 +90,34 @@ addHook({
   return Knex
 })
 
+addHook({
+  name: 'knex',
+  versions: ['>=2'],
+  file: 'lib/execution/runner.js',
+}, Runner => {
+  shimmer.wrap(Runner.prototype, 'ensureConnection', ensureConnection => wrapPromisePoolQueryMethod(
+    ensureConnection,
+    runner => runner.connection || runner.builder?._connection ? undefined : runner.client,
+    runner => resolveKnexDriver(runner.client)
+  ))
+
+  return Runner
+})
+
+addHook({
+  name: 'knex',
+  versions: ['>=2'],
+  file: 'lib/execution/transaction.js',
+}, Transaction => {
+  shimmer.wrap(Transaction.prototype, 'acquireConnection', acquireConnection => wrapPromisePoolQueryMethod(
+    acquireConnection,
+    (transaction, args) => args[0]?.connection ? undefined : transaction.client,
+    transaction => resolveKnexDriver(transaction.client)
+  ))
+
+  return Transaction
+})
+
 function wrapCallbackWithFinish (callback, finish, context) {
   if (typeof callback !== 'function') return callback
 
@@ -84,3 +125,38 @@ function wrapCallbackWithFinish (callback, finish, context) {
     finish(context, () => callback.apply(this, args))
   })
 }
+
+/**
+ * @param {{ pool?: { numFree?: () => number, numPendingAcquires?: () => number } }} client
+ * @returns {boolean}
+ */
+function knexHasIdleConnection (client) {
+  const pool = client.pool
+  if (typeof pool?.numFree !== 'function') return false
+  return pool.numFree() > pool.numPendingAcquires()
+}
+
+/**
+ * @param {{ driverName?: string }|undefined} client
+ * @returns {'mysql'|'mysql2'|'pg'|undefined}
+ */
+function resolveKnexDriver (client) {
+  switch (client?.driverName) {
+    case 'mysql':
+      return 'mysql'
+    case 'mysql2':
+      return 'mysql2'
+    case 'pg':
+    case 'cockroachdb':
+    case 'pg-redshift':
+      return 'pg'
+  }
+}
+
+/**
+ * @param {{ connectionSettings?: Record<string, unknown> }} client
+ * @returns {Record<string, unknown>}
+ */
+function resolveKnexConfig (client) {
+  return client.connectionSettings ?? {}
+}

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

@@ -11,6 +11,7 @@ const {
   runOutsidePoolQueryAcquire,
   runPoolAcquireError,
   runWithPoolWait,
+  setPromisePoolConnectionConfig,
   takePoolWaitTime,
   wrapPoolQueryMethod,
 } = require('./helpers/pool-acquire')
@@ -97,6 +98,8 @@ function createWrapQuery (options) {
       if (!startCh.hasSubscribers) return query.apply(this, arguments)
 
       const ctx = { sql, conf: options }
+      const poolWaitTime = takePoolWaitTime(this)
+      if (poolWaitTime !== undefined) ctx.poolWaitTime = poolWaitTime
 
       return startCh.runStores(ctx, query, this, ...arguments)
         .then(result => {
@@ -146,8 +149,11 @@ function createWrapQueryCallback (options) {
 function wrapConnection (promiseMethod, Connection) {
   return function (options) {
     Connection.apply(this, arguments)
+    setPromisePoolConnectionConfig(this, options)
 
+    const query = this[promiseMethod]
     shimmer.wrap(this, promiseMethod, createWrapQuery(options))
+    if (this.query === query) this.query = this[promiseMethod]
     shimmer.wrap(this, '_queryCallback', createWrapQueryCallback(options))
   }
 }

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

@@ -11,6 +11,60 @@ const MINIMUM_MOCHA_VERSION = DD_MAJOR >= 6 ? '>=8.0.0' : '>=5.2.0'
 const parameterizedTestCh = channel('ci:mocha:test:parameterize')
 const patched = new WeakSet()
 
+/**
+ * Registers every Mocha package view that can expose the constructor.
+ *
+ * @param {string[]} versions
+ * @param {(Mocha: Function, frameworkVersion: string) => Function} wrapMochaRun
+ * @returns {void}
+ */
+function addMochaRunHooks (versions, wrapMochaRun) {
+  const patchedMochaConstructors = new WeakSet()
+
+  /**
+   * @param {Function} Mocha
+   * @param {string} frameworkVersion
+   * @returns {Function}
+   */
+  function wrapMochaOnce (Mocha, frameworkVersion) {
+    if (patchedMochaConstructors.has(Mocha)) return Mocha
+    patchedMochaConstructors.add(Mocha)
+    return wrapMochaRun(Mocha, frameworkVersion)
+  }
+
+  /**
+   * @param {object|Function} MochaPackage
+   * @param {string} frameworkVersion
+   * @returns {object|Function}
+   */
+  function wrapMochaPackage (MochaPackage, frameworkVersion) {
+    wrapMochaOnce(MochaPackage.default ?? MochaPackage, frameworkVersion)
+    return MochaPackage
+  }
+
+  addHook({
+    name: 'mocha',
+    versions,
+    filePattern: String.raw`lib/mocha\.(?:c?js)$`,
+  }, wrapMochaOnce)
+
+  // Mocha 12's ESM package root can load lib/mocha.cjs before the CJS hook observes it. ESM imports report the
+  // package root while synchronous require reports index.js, so hook both views. Keeping the namespace intact makes
+  // the ESM hook skip its separate default-export callback, and wrapMochaOnce deduplicates loaders exposing both.
+  addHook({
+    name: 'mocha',
+    versions: ['>=12.0.0'],
+    patchDefault: false,
+  }, wrapMochaPackage)
+
+  addHook({
+    name: 'mocha',
+    versions: ['>=12.0.0'],
+    file: 'index.js',
+    patchDefault: false,
+  }, wrapMochaPackage)
+}
+
 // mocha-each support
 addHook({
   name: 'mocha-each',
@@ -38,7 +92,8 @@ addHook({
   name: 'mocha',
   versions: [MINIMUM_MOCHA_VERSION],
   file: 'lib/suite.js',
-}, (Suite) => {
+}, (SuitePackage) => {
+  const Suite = SuitePackage.Suite ?? SuitePackage
   shimmer.wrap(Suite.prototype, 'addTest', addTest => function (test) {
     const callSites = getCallSites()
     const testCallSite = callSites.find(site => site.getFileName() === test.file)
@@ -47,5 +102,7 @@ addHook({
     }
     return addTest.apply(this, arguments)
   })
-  return Suite
+  return SuitePackage
 })
+
+module.exports = { addMochaRunHooks }

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

@@ -35,6 +35,7 @@ const {
   isMarkedAsUnskippable,
 } = require('../../../dd-trace/src/plugins/util/test')
 
+const { addMochaRunHooks } = require('./common')
 const {
   isNewTest,
   getTestProperties,
@@ -60,8 +61,6 @@ const {
   adjustRunnerFailuresForTestOptimization,
 } = require('./utils')
 
-require('./common')
-
 const MINIMUM_MOCHA_VERSION = DD_MAJOR >= 6 ? '>=8.0.0' : '>=5.2.0'
 
 /**
@@ -536,7 +535,10 @@ function stopCurrentHook (runner, hook) {
     if (test) {
       if (hook.parent?._afterEach?.includes(hook) || hook.parent?._afterAll?.includes(hook)) {
         markTestTerminal(runner, test)
-        if (!test._ddTestFinishStarted) runnerTestEndHandlers.get(runner)?.(test)
+        if (!test._ddTestFinishStarted) {
+          const onTestEnd = runnerTestEndHandlers.get(runner)
+          onTestEnd?.(test)
+        }
       } else {
         markTestPending(runner, test)
       }
@@ -763,14 +765,18 @@ function wrapRunnerEmit (Runner) {
             const test = hook.ctx?.currentTest
             if (test && hook.parent?._afterEach?.includes(hook)) {
               markTestTerminal(this, test)
-              if (!test._ddTestFinishStarted) runnerTestEndHandlers.get(this)?.(test)
+              if (!test._ddTestFinishStarted) {
+                const onTestEnd = runnerTestEndHandlers.get(this)
+                onTestEnd?.(test)
+              }
             }
           } else if (event === 'pending' || event === 'pass' || event === 'fail' || event === 'retry' ||
             event === 'test end') {
             const test = arguments[1]
             stopAfterEachHooks(this, test)
             if (event === 'test end' && !test._ddTestFinishStarted) {
-              runnerTestEndHandlers.get(this)?.(test)
+              const onTestEnd = runnerTestEndHandlers.get(this)
+              onTestEnd?.(test)
             }
           }
         }
@@ -1063,13 +1069,15 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
 // 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.
-addHook({
-  name: 'mocha',
-  versions: [MINIMUM_MOCHA_VERSION],
-  file: 'lib/mocha.js',
-}, (Mocha, frameworkVersion) => {
+/**
+ * @param {Function} Mocha
+ * @param {string} frameworkVersion
+ * @returns {Function}
+ */
+function wrapMochaRun (Mocha, frameworkVersion) {
   warnDeprecatedMochaVersion(frameworkVersion)
 
+  // Shimmer is required because run must return its Runner while execution is paused and resumed after configuration.
   shimmer.wrap(Mocha.prototype, 'run', run => function (...args) {
     // Workers do not need to request any data, just run the tests
     if (!testFinishCh.hasSubscribers || getEnvironmentVariable('MOCHA_WORKER_ID') || this.options.parallel) {
@@ -1121,12 +1129,14 @@ addHook({
     return runner
   })
   return Mocha
-})
+}
+
+addMochaRunHooks([MINIMUM_MOCHA_VERSION], wrapMochaRun)
 
 addHook({
   name: 'mocha',
   versions: [MINIMUM_MOCHA_VERSION],
-  file: 'lib/cli/run-helpers.js',
+  filePattern: String.raw`lib/cli/run-helpers\.(?:c?js)$`,
 }, (run) => {
   // `runMocha` is an async function
   shimmer.wrap(run, 'runMocha', runMocha => function (...args) {
@@ -1157,7 +1167,7 @@ addHook({
 addHook({
   name: 'mocha',
   versions: [MINIMUM_MOCHA_VERSION],
-  file: 'lib/runner.js',
+  filePattern: String.raw`lib/runner\.(?:c?js)$`,
 }, function (Runner, frameworkVersion) {
   if (patched.has(Runner)) return Runner
 
@@ -1588,7 +1598,7 @@ addHook({
 addHook({
   name: 'mocha',
   versions: ['>=8.0.0'],
-  file: 'lib/nodejs/parallel-buffered-runner.js',
+  filePattern: String.raw`lib/nodejs/parallel-buffered-runner\.(?:c?js)$`,
 }, (ParallelBufferedRunner, frameworkVersion) => {
   shimmer.wrap(ParallelBufferedRunner.prototype, 'run', run => function (cb, { files, options = {} }) {
     if (!testFinishCh.hasSubscribers) {
@@ -1667,7 +1677,7 @@ addHook({
 addHook({
   name: 'mocha',
   versions: ['>=8.0.0'],
-  file: 'lib/nodejs/buffered-worker-pool.js',
+  filePattern: String.raw`lib/nodejs/buffered-worker-pool\.(?:c?js)$`,
 }, (BufferedWorkerPoolPackage, frameworkVersion) => {
   const { BufferedWorkerPool } = BufferedWorkerPoolPackage
 

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

@@ -507,7 +507,8 @@ function inheritDatadogPropertiesFromRetriedTest (test) {
 }
 
 function runnableWrapper (RunnablePackage, libraryConfig) {
-  shimmer.wrap(RunnablePackage.prototype, 'run', run => function (...args) {
+  const Runnable = RunnablePackage.Runnable ?? RunnablePackage
+  shimmer.wrap(Runnable.prototype, 'run', run => function (...args) {
     if (!testFinishCh.hasSubscribers) {
       return run.apply(this, args)
     }

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

@@ -6,6 +6,11 @@ const {
   WEBDRIVERIO_WORKER_EVENT,
   WEBDRIVERIO_WORKER_ORIGIN,
 } = require('../../../dd-trace/src/ci-visibility/exporters/test-worker/webdriverio')
+const { FINAL_FLUSH_TIMEOUT } = require('../../../dd-trace/src/ci-visibility/final-flush')
+
+const SCREENSHOT_UPLOAD = 'dd:test-optimization:webdriverio:screenshot:upload'
+const SCREENSHOT_UPLOAD_RESPONSE = 'dd:test-optimization:webdriverio:screenshot:upload:response'
+const SCREENSHOT_UPLOAD_TIMEOUT_MS = FINAL_FLUSH_TIMEOUT + 5000
 
 /**
  * Sends a message over WebdriverIO's worker IPC envelope.
@@ -30,10 +35,103 @@ function sendWebdriverioWorkerMessage (message, onError, onDone) {
   })
 }
 
+let screenshotUploadRequestId = 0
+const screenshotUploadRequests = new Map()
+
+/**
+ * Removes shared screenshot response listeners when there are no pending requests.
+ *
+ * @returns {void}
+ */
+function removeScreenshotUploadListeners () {
+  if (screenshotUploadRequests.size !== 0) return
+
+  process.off('message', onScreenshotUploadResponse)
+  process.off('disconnect', onScreenshotUploadDisconnect)
+}
+
+/**
+ * Completes one pending screenshot upload request.
+ *
+ * @param {string} requestId
+ * @param {Error} [error]
+ * @returns {void}
+ */
+function finishScreenshotUploadRequest (requestId, error) {
+  const request = screenshotUploadRequests.get(requestId)
+  if (!request) return
+
+  screenshotUploadRequests.delete(requestId)
+  clearTimeout(request.timeout)
+  removeScreenshotUploadListeners()
+  request.onDone(error)
+}
+
+/**
+ * Dispatches one coordinator screenshot response to its pending request.
+ *
+ * @param {object} message
+ * @returns {void}
+ */
+function onScreenshotUploadResponse (message) {
+  if (message?.name !== SCREENSHOT_UPLOAD_RESPONSE) return
+
+  const { error: errorMessage, requestId } = message.content || {}
+  if (!requestId) return
+
+  finishScreenshotUploadRequest(requestId, errorMessage ? new Error(errorMessage) : undefined)
+}
+
+/**
+ * Fails every pending screenshot upload after coordinator disconnect.
+ *
+ * @returns {void}
+ */
+function onScreenshotUploadDisconnect () {
+  for (const requestId of screenshotUploadRequests.keys()) {
+    finishScreenshotUploadRequest(
+      requestId,
+      new Error('WebdriverIO coordinator disconnected during screenshot upload')
+    )
+  }
+}
+
+/**
+ * Requests one screenshot upload from the WebdriverIO coordinator.
+ *
+ * @param {object} content - Screenshot upload metadata
+ * @param {(error?: Error) => void} onDone - Upload completion callback
+ * @returns {void}
+ */
+function requestWebdriverioScreenshotUpload (content, onDone) {
+  const requestId = `${process.pid}-${++screenshotUploadRequestId}`
+  const timeout = setTimeout(() => {
+    finishScreenshotUploadRequest(requestId, new Error('WebdriverIO screenshot upload timed out'))
+  }, SCREENSHOT_UPLOAD_TIMEOUT_MS)
+  timeout.unref?.()
+  if (screenshotUploadRequests.size === 0) {
+    process.on('message', onScreenshotUploadResponse)
+    process.once('disconnect', onScreenshotUploadDisconnect)
+  }
+  screenshotUploadRequests.set(requestId, { onDone, timeout })
+  sendWebdriverioWorkerMessage({
+    origin: 'datadog',
+    name: SCREENSHOT_UPLOAD,
+    content: { ...content, requestId },
+  }, error => finishScreenshotUploadRequest(
+    requestId,
+    error || new Error('WebdriverIO screenshot upload IPC failed')
+  ))
+}
+
 module.exports = {
   CONFIGURATION_REQUEST: 'dd:test-optimization:webdriverio:configuration:request',
   CONFIGURATION_RESPONSE: 'dd:test-optimization:webdriverio:configuration:response',
   createWebdriverioWorkerMessage,
+  requestWebdriverioScreenshotUpload,
+  SCREENSHOT_UPLOAD,
+  SCREENSHOT_UPLOAD_RESPONSE,
+  SCREENSHOT_UPLOAD_TIMEOUT_MS,
   sendWebdriverioWorkerMessage,
   SUITE_FINISH: 'dd:test-optimization:webdriverio:test-suite:finish',
   WORKER_READY: 'dd:test-optimization:webdriverio:worker:ready',

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

@@ -7,6 +7,7 @@ const { getEnvironmentVariable } = require('../../../dd-trace/src/config/helper'
 const log = require('../../../dd-trace/src/log')
 const { DD_MAJOR } = require('../../../../version')
 
+const { addMochaRunHooks } = require('./common')
 const {
   runnableWrapper,
   getOnTestHandler,
@@ -29,7 +30,6 @@ const {
   WEBDRIVERIO_WORKER_ENV,
   WORKER_READY,
 } = require('./webdriverio-protocol')
-require('./common')
 
 const MINIMUM_MOCHA_VERSION = DD_MAJOR >= 6 ? '>=8.0.0' : '>=5.2.0'
 
@@ -360,13 +360,15 @@ function isFailedTestReplayEnabled () {
   return config.isTestDynamicInstrumentationEnabled && config.isDiEnabled
 }
 
-addHook({
-  name: 'mocha',
-  versions: ['>=8.0.0'],
-  file: 'lib/mocha.js',
-}, (Mocha, frameworkVersion) => {
+/**
+ * @param {Function} Mocha
+ * @param {string} frameworkVersion
+ * @returns {Function}
+ */
+function wrapMochaRun (Mocha, frameworkVersion) {
   reportWebdriverioWorkerReady(frameworkVersion)
 
+  // Shimmer is required because run must return its Runner while execution is paused and resumed after configuration.
   shimmer.wrap(Mocha.prototype, 'run', run => function (...args) {
     applyMochaOptions(this.options)
     if (!isWebdriverioWorker || !workerFinishCh.hasSubscribers) {
@@ -438,13 +440,15 @@ addHook({
   })
 
   return Mocha
-})
+}
+
+addMochaRunHooks(['>=8.0.0'], wrapMochaRun)
 
 // Runner is also hooked in mocha/main.js, but in here we only generate test events.
 addHook({
   name: 'mocha',
   versions: [MINIMUM_MOCHA_VERSION],
-  file: 'lib/runner.js',
+  filePattern: String.raw`lib/runner\.(?:c?js)$`,
 }, function (Runner) {
   shimmer.wrap(Runner.prototype, 'runTests', runTests => getRunTestsWrapper(runTests, config))
 

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

@@ -27,8 +27,6 @@ addHook({ name: 'net' }, (net) => {
     }
 
     const options = getOptions(args)
-    const lastIndex = args.length - 1
-    const callback = args[lastIndex]
 
     if (!options) return connect.apply(this, args)
 
@@ -38,42 +36,39 @@ addHook({ name: 'net' }, (net) => {
     const errorCh = protocol === 'ipc' ? errorICPCh : errorTCPCh
     const ctx = { options }
 
-    if (typeof callback === 'function') {
-      args[lastIndex] = function (...args) {
-        return finishCh.runStores(ctx, callback, this, ...args)
-      }
-    }
-
     return startCh.runStores(ctx, () => {
-      setupListeners(this, protocol, ctx, finishCh, errorCh)
+      const cleanup = setupListeners(this, protocol, ctx, finishCh, errorCh)
 
-      const emit = this.emit
+      const originalEmit = this.emit
       let pendingReadyEvents = 2
       // Named `emit`/arity-1 mirrors the socket method so the per-socket wrap
       // skips its name/length rewrite.
-      this.emit = shimmer.wrapFunction(emit, originalEmit => function emit (eventName) {
+      this.emit = shimmer.wrapFunction(originalEmit, wrappedEmit => function emit (eventName) {
         switch (eventName) {
           case 'ready':
           case 'connect':
             if (--pendingReadyEvents === 0) this.emit = originalEmit
             return readyCh.runStores(ctx, () => {
-              return Reflect.apply(originalEmit, this, arguments)
+              return Reflect.apply(wrappedEmit, this, arguments)
             })
           case 'error':
           case 'close':
             this.emit = originalEmit
-            return Reflect.apply(originalEmit, this, arguments)
+            return Reflect.apply(wrappedEmit, this, arguments)
           default:
-            return Reflect.apply(originalEmit, this, arguments)
+            return Reflect.apply(wrappedEmit, this, arguments)
         }
       })
 
       try {
         return connect.apply(this, args)
-      } catch (err) {
-        errorCh.publish(err)
+      } catch (error) {
+        cleanup()
+        this.emit = originalEmit
+        errorCh.publish(error)
+        finishCh.runStores(ctx, () => {})
 
-        throw err
+        throw error
       }
     })
   })
@@ -142,4 +137,6 @@ function setupListeners (socket, protocol, ctx, finishCh, errorCh) {
   for (const event of events) {
     socket.once(event, wrapListener)
   }
+
+  return cleanupOtherListeners
 }