dd-trace 5.103.0 → 5.104.0

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

@@ -18,6 +18,7 @@ const {
   getIsFaultyEarlyFlakeDetection,
   collectTestOptimizationSummariesFromTraces,
   logTestOptimizationSummary,
+  getTestOptimizationRequestResults,
 } = require('../../../dd-trace/src/plugins/util/test')
 
 const {
@@ -109,6 +110,29 @@ function isTestFailed (test) {
   return false
 }
 
+function getRootSuiteStatus (rootTests) {
+  let status = 'pass'
+  if (rootTests.every(t => t.isPending())) {
+    status = 'skip'
+  } else {
+    for (const test of rootTests) {
+      if (test.state === 'failed' || test.timedOut || test._ddHookFailed) {
+        status = 'fail'
+      }
+    }
+  }
+  return status
+}
+
+function haveRootTestsFinished (rootTests) {
+  for (const test of rootTests) {
+    if (!test.isPending() && !test.state && !test.timedOut && !test._ddHookFailed) {
+      return false
+    }
+  }
+  return true
+}
+
 function getFilteredSuites (originalSuites) {
   return originalSuites.reduce((acc, suite) => {
     const testPath = getTestSuitePath(suite.file, process.cwd())
@@ -226,11 +250,38 @@ function getOnEndHandler (isParallel) {
   }
 }
 
+function getRunStoresPromise (channelToPublishTo, ctx) {
+  return new Promise(resolve => {
+    channelToPublishTo.runStores({ ...ctx, onDone: resolve }, () => {})
+  })
+}
+
+function applyKnownTestsResponse ({ err, knownTests }) {
+  if (err) {
+    config.knownTests = []
+    config.isEarlyFlakeDetectionEnabled = false
+    config.isKnownTestsEnabled = false
+  } else {
+    config.knownTests = knownTests
+  }
+}
+
+function applyTestManagementTestsResponse ({ err, testManagementTests: receivedTestManagementTests }) {
+  if (err) {
+    config.testManagementTests = {}
+    config.isTestManagementTestsEnabled = false
+    config.testManagementAttemptToFixRetries = 0
+  } else {
+    config.testManagementTests = receivedTestManagementTests
+  }
+}
+
 function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFinishRequest) {
   const ctx = {
     isParallel,
     frameworkVersion,
   }
+  let skippableSuitesResponse
 
   const onReceivedSkippableSuites = ({ err, skippableSuites, itrCorrelationId: responseItrCorrelationId }) => {
     if (err) {
@@ -256,6 +307,16 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     })
   }
 
+  const requestSkippableSuites = () => {
+    if (skippableSuitesResponse) {
+      onReceivedSkippableSuites(skippableSuitesResponse)
+      return
+    }
+
+    ctx.onDone = onReceivedSkippableSuites
+    skippableSuitesCh.runStores(ctx, () => {})
+  }
+
   const onReceivedImpactedTests = ({ err, modifiedFiles: receivedModifiedFiles }) => {
     if (err) {
       config.modifiedFiles = []
@@ -264,8 +325,7 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
       config.modifiedFiles = receivedModifiedFiles
     }
     if (config.isSuitesSkippingEnabled) {
-      ctx.onDone = onReceivedSkippableSuites
-      skippableSuitesCh.runStores(ctx, () => {})
+      requestSkippableSuites()
     } else {
       mochaGlobalRunCh.runStores(ctx, () => {
         onFinishRequest()
@@ -273,44 +333,12 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     }
   }
 
-  const onReceivedTestManagementTests = ({ err, testManagementTests: receivedTestManagementTests }) => {
-    if (err) {
-      config.testManagementTests = {}
-      config.isTestManagementTestsEnabled = false
-      config.testManagementAttemptToFixRetries = 0
-    } else {
-      config.testManagementTests = receivedTestManagementTests
-    }
+  const continueAfterTestRequests = () => {
     if (config.isImpactedTestsEnabled) {
       ctx.onDone = onReceivedImpactedTests
       modifiedFilesCh.runStores(ctx, () => {})
     } else if (config.isSuitesSkippingEnabled) {
-      ctx.onDone = onReceivedSkippableSuites
-      skippableSuitesCh.runStores(ctx, () => {})
-    } else {
-      mochaGlobalRunCh.runStores(ctx, () => {
-        onFinishRequest()
-      })
-    }
-  }
-
-  const onReceivedKnownTests = ({ err, knownTests }) => {
-    if (err) {
-      config.knownTests = []
-      config.isEarlyFlakeDetectionEnabled = false
-      config.isKnownTestsEnabled = false
-    } else {
-      config.knownTests = knownTests
-    }
-    if (config.isTestManagementTestsEnabled) {
-      ctx.onDone = onReceivedTestManagementTests
-      testManagementTestsCh.runStores(ctx, () => {})
-    } else if (config.isImpactedTestsEnabled) {
-      ctx.onDone = onReceivedImpactedTests
-      modifiedFilesCh.runStores(ctx, () => {})
-    } else if (config.isSuitesSkippingEnabled) {
-      ctx.onDone = onReceivedSkippableSuites
-      skippableSuitesCh.runStores(ctx, () => {})
+      requestSkippableSuites()
     } else {
       mochaGlobalRunCh.runStores(ctx, () => {
         onFinishRequest()
@@ -336,23 +364,30 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     config.isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
     config.flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
 
-    if (config.isKnownTestsEnabled) {
-      ctx.onDone = onReceivedKnownTests
-      knownTestsCh.runStores(ctx, () => {})
-    } else if (config.isTestManagementTestsEnabled) {
-      ctx.onDone = onReceivedTestManagementTests
-      testManagementTestsCh.runStores(ctx, () => {})
-    } else if (config.isImpactedTestsEnabled) {
-      ctx.onDone = onReceivedImpactedTests
-      modifiedFilesCh.runStores(ctx, () => {})
-    } else if (config.isSuitesSkippingEnabled) {
-      ctx.onDone = onReceivedSkippableSuites
-      skippableSuitesCh.runStores(ctx, () => {})
-    } else {
-      mochaGlobalRunCh.runStores(ctx, () => {
-        onFinishRequest()
-      })
-    }
+    getTestOptimizationRequestResults({
+      isKnownTestsEnabled: config.isKnownTestsEnabled,
+      isTestManagementTestsEnabled: config.isTestManagementTestsEnabled,
+      isSuitesSkippingEnabled: config.isSuitesSkippingEnabled,
+      getKnownTests: () => getRunStoresPromise(knownTestsCh, ctx),
+      getTestManagementTests: () => getRunStoresPromise(testManagementTestsCh, ctx),
+      getSkippableSuites: () => getRunStoresPromise(skippableSuitesCh, ctx),
+    }).then(requestResults => {
+      const {
+        knownTestsResponse,
+        testManagementTestsResponse,
+        skippableSuitesResponse: requestSkippableSuitesResponse,
+      } = requestResults
+
+      if (knownTestsResponse) {
+        applyKnownTestsResponse(knownTestsResponse)
+      }
+      if (testManagementTestsResponse) {
+        applyTestManagementTestsResponse(testManagementTestsResponse)
+      }
+      skippableSuitesResponse = requestSkippableSuitesResponse
+
+      continueAfterTestRequests()
+    })
   }
 
   ctx.onDone = onReceivedConfiguration
@@ -473,22 +508,166 @@ addHook({
     // Populated during the root 'suite' event so the normal finish path can include them
     // in mixed-file status calculation.
     const rootTestsByFile = new Map()
+    // Counts how many original tests per pure-root file still need their final attempt.
+    // Hits zero when the last test's lifecycle completes, triggering the suite finish.
+    const rootPendingCountByFile = new Map()
+    const rootFinalizationPendingCountByFile = new Map()
+    const rootFallbackPendingFiles = new Set()
+    const rootFinalizationPendingTests = new WeakSet()
+    let pendingRootFinalizations = 0
+    let hasEnded = false
+    let hasFinishedRun = false
+    let endRunner
+
+    function updateRootTestForFinalAttempt (test) {
+      if (!test._retriedTest) return
+
+      const rootTests = rootTestsByFile.get(test.file)
+      if (!rootTests) return
+
+      const retriedTestIndex = rootTests.indexOf(test._retriedTest)
+      if (retriedTestIndex !== -1) {
+        rootTests[retriedTestIndex] = test
+      }
+    }
+
+    function finishRunIfReady () {
+      if (hasFinishedRun) return
+      if (hasEnded && pendingRootFinalizations === 0) {
+        hasFinishedRun = true
+        onEnd.call(endRunner)
+      }
+    }
+
+    function incrementPendingRootFinalization (test) {
+      if (!rootPendingCountByFile.has(test.file) || rootFinalizationPendingTests.has(test)) return
+
+      rootFinalizationPendingTests.add(test)
+      pendingRootFinalizations++
+      rootFinalizationPendingCountByFile.set(
+        test.file,
+        (rootFinalizationPendingCountByFile.get(test.file) || 0) + 1
+      )
+    }
+
+    function decrementPendingRootFinalization (test) {
+      if (!rootFinalizationPendingTests.has(test)) return
+
+      rootFinalizationPendingTests.delete(test)
+      pendingRootFinalizations--
+
+      const remaining = rootFinalizationPendingCountByFile.get(test.file) - 1
+      if (remaining > 0) {
+        rootFinalizationPendingCountByFile.set(test.file, remaining)
+      } else {
+        rootFinalizationPendingCountByFile.delete(test.file)
+      }
+
+      if (!rootFinalizationPendingCountByFile.has(test.file) && rootFallbackPendingFiles.delete(test.file)) {
+        finishRootSuiteFallbackForFile(test.file)
+      }
+
+      finishRunIfReady()
+    }
+
+    function finishRootSuiteForFile (file) {
+      const remaining = rootPendingCountByFile.get(file) - 1
+      if (remaining > 0) {
+        rootPendingCountByFile.set(file, remaining)
+        return
+      }
+      rootPendingCountByFile.delete(file)
+
+      const ctx = testFileToSuiteCtx.get(file)
+      if (!ctx) {
+        log.warn('No ctx found for suite', file)
+        return
+      }
+
+      const rootTests = rootTestsByFile.get(file) || []
+      const status = getRootSuiteStatus(rootTests)
+
+      if (global.__coverage__) {
+        const coverageFiles = getCoveredFilenamesFromCoverage(global.__coverage__)
+        testSuiteCodeCoverageCh.publish({ coverageFiles, suiteFile: file })
+        mergeCoverage(global.__coverage__, originalCoverageMap)
+        resetCoverage(global.__coverage__)
+      }
+
+      testSuiteFinishCh.publish({ status, ...ctx.currentStore }, () => {})
+    }
+
+    function finishRootSuiteFallbackForFile (file) {
+      const ctx = testFileToSuiteCtx.get(file)
+      if (!ctx || !rootPendingCountByFile.has(file)) return
+
+      const rootTests = rootTestsByFile.get(file) || []
+      const status = haveRootTestsFinished(rootTests) ? getRootSuiteStatus(rootTests) : 'fail'
+      rootPendingCountByFile.delete(file)
+      testSuiteFinishCh.publish({ status, ...ctx.currentStore }, () => {})
+    }
+
+    function finishRootSuiteAfterFinalAttempt (test) {
+      if (!test._ddIsFinalAttempt || !rootPendingCountByFile.has(test.file)) return
+
+      updateRootTestForFinalAttempt(test)
+      finishRootSuiteForFile(test.file)
+    }
+
+    const onEnd = getOnEndHandler(false)
 
     this.once('start', getOnStartHandler(frameworkVersion))
 
-    this.once('end', getOnEndHandler(false))
+    this.once('end', function () {
+      hasEnded = true
+      endRunner = this
+      finishRunIfReady()
+    })
+
+    // The job of this listener is to
+    // initialize the suite span tag in correct order
+    // (that is suiteA -> testA ... -> suiteB -> testB
+    // instead of suiteA -> suiteB -> testA -> ... -> testB)
+    // when the suite has tests that are in the top level
+    // (no describe(...))
+    this.on('test', function (test) {
+      const ctx = testFileToSuiteCtx.get(test.file)
+      if (ctx?._pendingRootStart) {
+        ctx._pendingRootStart = false
+        testSuiteStartCh.runStores(ctx, () => {})
+      }
+    })
 
     this.on('test', getOnTestHandler(true))
 
-    this.on('test end', getOnTestEndHandler(config))
+    this.on('test end', getOnTestEndHandler(config, {
+      onStart: incrementPendingRootFinalization,
+      onFinish: function (test) {
+        finishRootSuiteAfterFinalAttempt(test)
+        decrementPendingRootFinalization(test)
+      },
+    }))
 
     this.on('retry', getOnTestRetryHandler(config))
 
     // If the hook passes, 'hook end' will be emitted. Otherwise, 'fail' will be emitted
     this.on('hook end', getOnHookEndHandler(config))
 
+    this.on('hook end', function (hook) {
+      const test = hook.ctx?.currentTest
+      if (!test) return
+      finishRootSuiteAfterFinalAttempt(test)
+    })
+
     this.on('fail', getOnFailHandler(true, config))
 
+    this.on('fail', function (testOrHook) {
+      if (testOrHook.type !== 'hook') return
+      const test = testOrHook.ctx?.currentTest
+      if (!test) return
+      finishRootSuiteAfterFinalAttempt(test)
+    })
+
     this.on('pending', getOnPendingHandler())
 
     this.on('suite', function (suite) {
@@ -503,7 +682,13 @@ addHook({
         if (suite.root && suite.tests.length > 0) {
           const files = new Set(suite.tests.map(test => test.file).filter(Boolean))
           for (const file of files) {
-            rootTestsByFile.set(file, suite.tests.filter(t => t.file === file))
+            const testsForFile = suite.tests.filter(t => t.file === file)
+            rootTestsByFile.set(file, testsForFile)
+            // Only track the countdown for pure root-level files.
+            // Mixed files are finished by the normal 'suite end' path.
+            if (!suitesByTestFile[file]) {
+              rootPendingCountByFile.set(file, testsForFile.length)
+            }
             if (testFileToSuiteCtx.get(file)) continue
             const isUnskippable = unskippableSuites.includes(file)
             isForcedToRun = isUnskippable && suitesToSkip.includes(getTestSuitePath(file, process.cwd()))
@@ -512,9 +697,9 @@ addHook({
               isUnskippable,
               isForcedToRun,
               itrCorrelationId,
+              _pendingRootStart: true, // Now the suite start fires lazily on the first test event for this file
             }
             testFileToSuiteCtx.set(file, ctx)
-            testSuiteStartCh.runStores(ctx, () => {})
           }
         }
         return
@@ -536,42 +721,38 @@ addHook({
 
     this.on('suite end', function (suite) {
       if (suite.root) {
-        // Symmetric to the suite start fix
-        const fileToTests = new Map()
+        // Normal case: pure root-level files are finished by the 'test end' / 'hook end'
+        // listeners via finishRootSuiteForFile. Two edge cases remain here:
+        //
+        // 1. All-pending: no 'test' event fired, _pendingRootStart is still true.
+        //    Start and immediately finish with 'skip'.
+        //
+        // 2. Aborted mid-run (e.g. a beforeEach hook failure): Mocha skips remaining
+        //    tests and jumps straight to 'suite end'. rootPendingCountByFile still has
+        //    a nonzero count for the file because the last tests never ran. Finish it
+        //    as failed now.
+        //
+        // 3. Async finalization lagged behind Mocha's synchronous events (e.g. DI retry
+        //    wait): all tests have Mocha terminal state, but the final-attempt callback
+        //    did not run before root 'suite end'. Finish from the observed test states.
+        const processedFiles = new Set()
         for (const test of suite.tests) {
-          if (!test.file) continue
-          if (!fileToTests.has(test.file)) fileToTests.set(test.file, [])
-          fileToTests.get(test.file).push(test)
-        }
-        for (const [file, tests] of fileToTests) {
-          // Mixed case: if a file appears in suitesByTestFile (pre-populated before the run),
-          // its numSuitesByTestFile counter hits zero when its last describe-based suite ends
-          // and the normal path below fires testSuiteFinishCh. Since root is last when
-          // 'suite end' fires, any such file has already been handled — skipping it here
-          // avoids duplication.
-          if (suitesByTestFile[file]) continue
-          let status = 'pass'
-          if (tests.every(test => test.isPending())) {
-            status = 'skip'
-          } else {
-            for (const test of tests) {
-              if (test.state === 'failed' || test.timedOut) {
-                status = 'fail'
-                break
-              }
+          if (!test.file || processedFiles.has(test.file)) continue
+          processedFiles.add(test.file)
+          if (suitesByTestFile[test.file]) continue // mixed: handled by normal path
+          const ctx = testFileToSuiteCtx.get(test.file)
+          if (!ctx) continue
+          if (ctx._pendingRootStart) {
+            ctx._pendingRootStart = false
+            testSuiteStartCh.runStores(ctx, () => {})
+            testSuiteFinishCh.publish({ status: 'skip', ...ctx.currentStore }, () => {})
+          } else if (rootPendingCountByFile.has(test.file)) {
+            if (rootFinalizationPendingCountByFile.has(test.file)) {
+              rootFallbackPendingFiles.add(test.file)
+              continue
             }
-          }
-          if (global.__coverage__) {
-            const coverageFiles = getCoveredFilenamesFromCoverage(global.__coverage__)
-            testSuiteCodeCoverageCh.publish({ coverageFiles, suiteFile: file })
-            mergeCoverage(global.__coverage__, originalCoverageMap)
-            resetCoverage(global.__coverage__)
-          }
-          const ctx = testFileToSuiteCtx.get(file)
-          if (ctx) {
-            testSuiteFinishCh.publish({ status, ...ctx.currentStore }, () => {})
-          } else {
-            log.warn('No ctx found for suite', file)
+
+            finishRootSuiteFallbackForFile(test.file)
           }
         }
         return

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

@@ -503,13 +503,38 @@ function getTestFinishInfo (test, status, config, error) {
   }
 }
 
-function getOnTestEndHandler (config) {
+function getOnTestEndHandler (config, finalAttemptHandlers) {
   return async function (test) {
     if (test._ddShouldSkipEfdRetry) {
       return
     }
     const ctx = getTestContext(test)
     const status = getTestStatus(test)
+    const shouldFinishTest = ctx && (!getAfterEachHooks(test).length || (test._ddIsDisabled && !test._ddIsAttemptToFix))
+    let testFinishInfo
+    let isFinalAttempt = false
+
+    // If there are afterEach to be run, we don't finish the test yet.
+    // Disabled tests (marked pending by us) are finished immediately without waiting for afterEach hooks.
+    // In older mocha versions, pending tests don't run afterEach hooks, so we can't rely on
+    // getOnHookEndHandler to finish the test. This mirrors Jest's approach where the skip handler
+    // directly sets finalStatus without waiting for hooks
+    if (!ctx && test.isPending()) {
+      test._ddIsFinalAttempt = true
+      isFinalAttempt = true
+    }
+
+    if (shouldFinishTest) {
+      testFinishInfo = getTestFinishInfo(test, status, config, ctx.err || test.err)
+      if (testFinishInfo.finalStatus !== undefined) {
+        test._ddIsFinalAttempt = true
+        isFinalAttempt = true
+      }
+    }
+
+    if (isFinalAttempt) {
+      finalAttemptHandlers?.onStart?.(test)
+    }
 
     // After finishing it might take a bit for the snapshot to be handled.
     // This means that tests retried with DI are BREAKPOINT_HIT_GRACE_PERIOD_MS slower at least.
@@ -521,13 +546,7 @@ function getOnTestEndHandler (config) {
       })
     }
 
-    // If there are afterEach to be run, we don't finish the test yet.
-    // Disabled tests (marked pending by us) are finished immediately without waiting for afterEach hooks.
-    // In older mocha versions, pending tests don't run afterEach hooks, so we can't rely on
-    // getOnHookEndHandler to finish the test. This mirrors Jest's approach where the skip handler
-    // directly sets finalStatus without waiting for hooks
-    if (ctx && (!getAfterEachHooks(test).length || (test._ddIsDisabled && !test._ddIsAttemptToFix))) {
-      const testFinishInfo = getTestFinishInfo(test, status, config, ctx.err || test.err)
+    if (shouldFinishTest) {
       testFinishCh.publish({
         status,
         hasBeenRetried: isMochaRetry(test),
@@ -536,6 +555,10 @@ function getOnTestEndHandler (config) {
         ...ctx.currentStore,
       })
     }
+
+    if (isFinalAttempt) {
+      finalAttemptHandlers?.onFinish?.(test)
+    }
   }
 }
 
@@ -552,6 +575,9 @@ function getOnHookEndHandler (config) {
         // skip to avoid double-publishing
         if (ctx && (!test._ddIsDisabled || test._ddIsAttemptToFix)) {
           const testFinishInfo = getTestFinishInfo(test, status, config, ctx.err || test.err)
+          if (testFinishInfo.finalStatus !== undefined) {
+            test._ddIsFinalAttempt = true
+          }
           testFinishCh.publish({
             status,
             hasBeenRetried: isMochaRetry(test),
@@ -583,6 +609,20 @@ function getOnFailHandler (isMain, config) {
         testContext.err = err
         errorCh.runStores(testContext, () => {})
         const testFinishInfo = getTestFinishInfo(test, 'fail', config, err)
+        // ATR never retries hook failures: this.retries(N) is set in runnableWrapper
+        // which only runs when the test function executes — hooks bypass that path,
+        // so _retries stays at -1 and getIsLastRetry returns false, leaving finalStatus
+        // undefined. We must also mark the attempt final when no clone-based retry
+        // mechanism (EFD original, EFD clone, ATF) has queued further attempts.
+        const noCloneRetries = !test._ddIsEfdRetry &&
+          !((test._ddIsNew || test._ddIsModified) && config.isEarlyFlakeDetectionEnabled) &&
+          !test._ddIsAttemptToFix
+        if (testFinishInfo.finalStatus !== undefined || noCloneRetries) {
+          test._ddIsFinalAttempt = true
+        }
+        // test.state is never set to 'failed' for hook failures (Mocha marks the hook,
+        // not the test). Flag it so finishRootSuiteForFile can compute the correct status.
+        test._ddHookFailed = true
         testFinishCh.publish({
           status: 'fail',
           hasBeenRetried: isMochaRetry(test),

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

@@ -174,7 +174,7 @@ function instrument (operation, command, instance, args, server, ns, ops, option
     name,
   }
   return startCh.runStores(ctx, () => {
-    args[index] = shimmer.wrapFunction(callback, callback => function (err, res) {
+    args[index] = shimmer.wrapCallback(callback, callback => function (err, res) {
       if (err) {
         ctx.error = err
         errorCh.publish(ctx)

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

@@ -125,21 +125,19 @@ addHook({
                     const resolve = args[0]
                     const reject = args[1]
 
-                    args[0] = shimmer.wrapFunction(resolve, resolve => function wrappedResolve (...args) {
-                      finishCh.publish(ctx)
-
-                      if (resolve) {
+                    if (typeof resolve === 'function') {
+                      args[0] = shimmer.wrapCallback(resolve, resolve => function wrappedResolve (...args) {
+                        finishCh.publish(ctx)
                         return resolve.apply(this, args)
-                      }
-                    })
-
-                    args[1] = shimmer.wrapFunction(reject, reject => function wrappedReject (...args) {
-                      finishCh.publish(ctx)
+                      })
+                    }
 
-                      if (reject) {
+                    if (typeof reject === 'function') {
+                      args[1] = shimmer.wrapCallback(reject, reject => function wrappedReject (...args) {
+                        finishCh.publish(ctx)
                         return reject.apply(this, args)
-                      }
-                    })
+                      })
+                    }
 
                     return originalThen.apply(this, args)
                   }

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

@@ -29,7 +29,7 @@ addHook({ name: 'mysql', file: 'lib/Connection.js', versions: ['>=2'] }, Connect
 
         if (res._callback) {
           const cb = res._callback
-          res._callback = shimmer.wrapFunction(cb, cb => function (error, result) {
+          res._callback = shimmer.wrapCallback(cb, cb => function (error, result) {
             if (error) {
               ctx.error = error
               errorCh.publish(ctx)
@@ -86,7 +86,7 @@ addHook({ name: 'mysql', file: 'lib/Pool.js', versions: ['>=2'] }, Pool => {
     return startPoolQueryCh.runStores(ctx, () => {
       const cb = args[args.length - 1]
       if (typeof cb === 'function') {
-        args[args.length - 1] = shimmer.wrapFunction(cb, cb => function (...args) {
+        args[args.length - 1] = shimmer.wrapCallback(cb, cb => function (...args) {
           return finishPoolQueryCh.runStores(ctx, cb, this, ...args)
         })
       }

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

@@ -176,7 +176,7 @@ function wrapConnection (Connection, version) {
         if (typeof this.onResult === 'function') {
           const onResult = this.onResult
 
-          this.onResult = shimmer.wrapFunction(onResult, onResult => function (error) {
+          this.onResult = shimmer.wrapCallback(onResult, onResult => function (error) {
             if (error) {
               ctx.error = error
               errorCh.publish(ctx)

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

@@ -193,7 +193,7 @@ function wrapPoolQuery (query) {
       }
 
       if (typeof cb === 'function') {
-        args[args.length - 1] = shimmer.wrapFunction(cb, cb => function (...args) {
+        args[args.length - 1] = shimmer.wrapCallback(cb, cb => function (...args) {
           finish(ctx)
           return cb.apply(this, args)
         })