dd-trace 5.21.0 → 5.23.0

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

@@ -11,12 +11,12 @@ const {
   fromCoverageMapToCoverage,
   getCoveredFilenamesFromCoverage,
   mergeCoverage,
-  resetCoverage
+  resetCoverage,
+  getIsFaultyEarlyFlakeDetection
 } = require('../../../dd-trace/src/plugins/util/test')
 
 const {
   isNewTest,
-  retryTest,
   getSuitesByTestFile,
   runnableWrapper,
   getOnTestHandler,
@@ -25,22 +25,21 @@ const {
   getOnHookEndHandler,
   getOnFailHandler,
   getOnPendingHandler,
-  testFileToSuiteAr
+  testFileToSuiteAr,
+  newTests,
+  getTestFullName,
+  getRunTestsWrapper
 } = require('./utils')
+
 require('./common')
 
 const testSessionAsyncResource = new AsyncResource('bound-anonymous-fn')
 const patched = new WeakSet()
-const newTests = {}
-let suitesToSkip = []
+
 const unskippableSuites = []
+let suitesToSkip = []
 let isSuitesSkipped = false
 let skippedSuites = []
-let isEarlyFlakeDetectionEnabled = false
-let isSuitesSkippingEnabled = false
-let isFlakyTestRetriesEnabled = false
-let earlyFlakeDetectionNumRetries = 0
-let knownTests = []
 let itrCorrelationId = ''
 let isForcedToRun = false
 const config = {}
@@ -69,6 +68,17 @@ const itrSkippedSuitesCh = channel('ci:mocha:itr:skipped-suites')
 
 const getCodeCoverageCh = channel('ci:nyc:get-coverage')
 
+// Tests from workers do not come with `isFailed` method
+function isTestFailed (test) {
+  if (test.isFailed) {
+    return test.isFailed()
+  }
+  if (test.isPending) {
+    return !test.isPending() && test.state !== 'failed'
+  }
+  return false
+}
+
 function getFilteredSuites (originalSuites) {
   return originalSuites.reduce((acc, suite) => {
     const testPath = getTestSuitePath(suite.file, process.cwd())
@@ -107,7 +117,7 @@ function getOnEndHandler (isParallel) {
       status = 'fail'
     }
 
-    if (!isParallel && isEarlyFlakeDetectionEnabled) {
+    if (config.isEarlyFlakeDetectionEnabled) {
       /**
        * If Early Flake Detection (EFD) is enabled the logic is as follows:
        * - If all attempts for a test are failing, the test has failed and we will let the test process fail.
@@ -116,7 +126,7 @@ function getOnEndHandler (isParallel) {
        * on flakiness (the test will be considered flaky), but you may choose to unblock the pipeline too.
        */
       for (const tests of Object.values(newTests)) {
-        const failingNewTests = tests.filter(test => test.isFailed())
+        const failingNewTests = tests.filter(test => isTestFailed(test))
         const areAllNewTestsFailing = failingNewTests.length === tests.length
         if (failingNewTests.length && !areAllNewTestsFailing) {
           this.stats.failures -= failingNewTests.length
@@ -153,13 +163,14 @@ function getOnEndHandler (isParallel) {
       hasForcedToRunSuites: isForcedToRun,
       hasUnskippableSuites: !!unskippableSuites.length,
       error,
-      isEarlyFlakeDetectionEnabled,
+      isEarlyFlakeDetectionEnabled: config.isEarlyFlakeDetectionEnabled,
+      isEarlyFlakeDetectionFaulty: config.isEarlyFlakeDetectionFaulty,
       isParallel
     })
   })
 }
 
-function getExecutionConfiguration (runner, onFinishRequest) {
+function getExecutionConfiguration (runner, isParallel, onFinishRequest) {
   const mochaRunAsyncResource = new AsyncResource('bound-anonymous-fn')
 
   const onReceivedSkippableSuites = ({ err, skippableSuites, itrCorrelationId: responseItrCorrelationId }) => {
@@ -186,15 +197,15 @@ function getExecutionConfiguration (runner, onFinishRequest) {
     onFinishRequest()
   }
 
-  const onReceivedKnownTests = ({ err, knownTests: receivedKnownTests }) => {
+  const onReceivedKnownTests = ({ err, knownTests }) => {
     if (err) {
-      knownTests = []
-      isEarlyFlakeDetectionEnabled = false
+      config.knownTests = []
+      config.isEarlyFlakeDetectionEnabled = false
     } else {
-      knownTests = receivedKnownTests
+      config.knownTests = knownTests
     }
 
-    if (isSuitesSkippingEnabled) {
+    if (config.isSuitesSkippingEnabled) {
       skippableSuitesCh.publish({
         onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
       })
@@ -208,21 +219,19 @@ function getExecutionConfiguration (runner, onFinishRequest) {
       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.isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
+    config.earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
+    config.earlyFlakeDetectionFaultyThreshold = libraryConfig.earlyFlakeDetectionFaultyThreshold
+    // ITR and auto test retries are not supported in parallel mode yet
+    config.isSuitesSkippingEnabled = !isParallel && libraryConfig.isSuitesSkippingEnabled
+    config.isFlakyTestRetriesEnabled = !isParallel && libraryConfig.isFlakyTestRetriesEnabled
+    config.flakyTestRetriesCount = !isParallel && libraryConfig.flakyTestRetriesCount
 
-    if (isEarlyFlakeDetectionEnabled) {
+    if (config.isEarlyFlakeDetectionEnabled) {
       knownTestsCh.publish({
         onDone: mochaRunAsyncResource.bind(onReceivedKnownTests)
       })
-    } else if (isSuitesSkippingEnabled) {
+    } else if (config.isSuitesSkippingEnabled) {
       skippableSuitesCh.publish({
         onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
       })
@@ -250,8 +259,8 @@ addHook({
       return run.apply(this, arguments)
     }
 
-    // `options.delay` does not work in parallel mode, so ITR and EFD can't work.
-    // TODO: use `lib/cli/run-helpers.js#runMocha` to get the data in parallel mode.
+    // `options.delay` does not work in parallel mode, so we can't delay the execution this way
+    // This needs to be both here and in `runMocha` hook. Read the comment in `runMocha` hook for more info.
     this.options.delay = true
 
     const runner = run.apply(this, arguments)
@@ -263,7 +272,19 @@ addHook({
       }
     })
 
-    getExecutionConfiguration(runner, () => {
+    getExecutionConfiguration(runner, false, () => {
+      if (config.isEarlyFlakeDetectionEnabled) {
+        const testSuites = this.files.map(file => getTestSuitePath(file, process.cwd()))
+        const isFaulty = getIsFaultyEarlyFlakeDetection(
+          testSuites,
+          config.knownTests?.mocha || {},
+          config.earlyFlakeDetectionFaultyThreshold
+        )
+        if (isFaulty) {
+          config.isEarlyFlakeDetectionEnabled = false
+          config.isEarlyFlakeDetectionFaulty = true
+        }
+      }
       if (getCodeCoverageCh.hasSubscribers) {
         getCodeCoverageCh.publish({
           onDone: (receivedCodeCoverage) => {
@@ -281,9 +302,6 @@ addHook({
   return Mocha
 })
 
-// Only used to set `mocha.options.delay` to true in serial mode. When the mocha CLI is used,
-// setting options.delay in Mocha#run is not enough to delay the execution.
-// TODO: modify this hook to grab the data in parallel mode, so that ITR and EFD can work.
 addHook({
   name: 'mocha',
   versions: ['>=5.2.0'],
@@ -293,15 +311,20 @@ addHook({
     if (!testStartCh.hasSubscribers) {
       return runMocha.apply(this, arguments)
     }
-
     const mocha = arguments[0]
+
     /**
      * This attaches `run` to the global context, which we'll call after
-     * our configuration and skippable suites requests
+     * our configuration and skippable suites requests.
+     * You need this both here and in Mocha#run hook: the programmatic API
+     * does not call `runMocha`, so it needs to be in Mocha#run. When using
+     * the CLI, modifying `options.delay` in Mocha#run is not enough (it's too late),
+     * so it also needs to be here.
      */
     if (!mocha.options.parallel) {
       mocha.options.delay = true
     }
+
     return runMocha.apply(this, arguments)
   })
   return run
@@ -318,18 +341,7 @@ addHook({
 
   patched.add(Runner)
 
-  shimmer.wrap(Runner.prototype, 'runTests', runTests => function (suite, fn) {
-    if (isEarlyFlakeDetectionEnabled) {
-      // by the time we reach `this.on('test')`, it is too late. We need to add retries here
-      suite.tests.forEach(test => {
-        if (!test.isPending() && isNewTest(test, knownTests)) {
-          test._ddIsNew = true
-          retryTest(test, earlyFlakeDetectionNumRetries)
-        }
-      })
-    }
-    return runTests.apply(this, arguments)
-  })
+  shimmer.wrap(Runner.prototype, 'runTests', runTests => getRunTestsWrapper(runTests, config))
 
   shimmer.wrap(Runner.prototype, 'run', run => function () {
     if (!testStartCh.hasSubscribers) {
@@ -513,10 +525,10 @@ addHook({
 // Used to start and finish test session and test module
 addHook({
   name: 'mocha',
-  versions: ['>=5.2.0'],
+  versions: ['>=8.0.0'],
   file: 'lib/nodejs/parallel-buffered-runner.js'
 }, (ParallelBufferedRunner, frameworkVersion) => {
-  shimmer.wrap(ParallelBufferedRunner.prototype, 'run', run => function () {
+  shimmer.wrap(ParallelBufferedRunner.prototype, 'run', run => function (cb, { files }) {
     if (!testStartCh.hasSubscribers) {
       return run.apply(this, arguments)
     }
@@ -524,8 +536,82 @@ addHook({
     this.once('start', getOnStartHandler(true, frameworkVersion))
     this.once('end', getOnEndHandler(true))
 
-    return run.apply(this, arguments)
+    getExecutionConfiguration(this, true, () => {
+      if (config.isEarlyFlakeDetectionEnabled) {
+        const testSuites = files.map(file => getTestSuitePath(file, process.cwd()))
+        const isFaulty = getIsFaultyEarlyFlakeDetection(
+          testSuites,
+          config.knownTests?.mocha || {},
+          config.earlyFlakeDetectionFaultyThreshold
+        )
+        if (isFaulty) {
+          config.isEarlyFlakeDetectionEnabled = false
+          config.isEarlyFlakeDetectionFaulty = true
+        }
+      }
+      run.apply(this, arguments)
+    })
+
+    return this
   })
 
   return ParallelBufferedRunner
 })
+
+// Only in parallel mode: BufferedWorkerPool#run is used to run a test file in a worker
+// If Early Flake Detection is enabled,
+// In this hook we pass the known tests to the worker and collect the new tests that run
+addHook({
+  name: 'mocha',
+  versions: ['>=8.0.0'],
+  file: 'lib/nodejs/buffered-worker-pool.js'
+}, (BufferedWorkerPoolPackage) => {
+  const { BufferedWorkerPool } = BufferedWorkerPoolPackage
+
+  shimmer.wrap(BufferedWorkerPool.prototype, 'run', run => async function (testSuiteAbsolutePath, workerArgs) {
+    if (!testStartCh.hasSubscribers || !config.isEarlyFlakeDetectionEnabled) {
+      return run.apply(this, arguments)
+    }
+
+    const testPath = getTestSuitePath(testSuiteAbsolutePath, process.cwd())
+    const testSuiteKnownTests = config.knownTests.mocha?.[testPath] || []
+
+    // We pass the known tests for the test file to the worker
+    const testFileResult = await run.apply(
+      this,
+      [
+        testSuiteAbsolutePath,
+        {
+          ...workerArgs,
+          _ddEfdNumRetries: config.earlyFlakeDetectionNumRetries,
+          _ddKnownTests: {
+            mocha: {
+              [testPath]: testSuiteKnownTests
+            }
+          }
+        }
+      ]
+    )
+    const tests = testFileResult
+      .events
+      .filter(event => event.eventName === 'test end')
+      .map(event => event.data)
+
+    // `newTests` is filled in the worker process, so we need to use the test results to fill it here too.
+    for (const test of tests) {
+      if (isNewTest(test, config.knownTests)) {
+        const testFullName = getTestFullName(test)
+        const tests = newTests[testFullName]
+
+        if (!tests) {
+          newTests[testFullName] = [test]
+        } else {
+          tests.push(test)
+        }
+      }
+    }
+    return testFileResult
+  })
+
+  return BufferedWorkerPoolPackage
+})

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')
@@ -25,6 +24,7 @@ const originalFns = new WeakMap()
 const testToStartLine = new WeakMap()
 const testFileToSuiteAr = new Map()
 const wrappedFunctions = new WeakSet()
+const newTests = {}
 
 function isNewTest (test, knownTests) {
   const testSuite = getTestSuitePath(test.file, process.cwd())
@@ -114,7 +114,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.
@@ -152,7 +152,7 @@ function runnableWrapper (RunnablePackage, libraryConfig) {
   return RunnablePackage
 }
 
-function getOnTestHandler (isMain, newTests) {
+function getOnTestHandler (isMain) {
   return function (test) {
     const testStartLine = testToStartLine.get(test)
     const asyncResource = new AsyncResource('bound-anonymous-fn')
@@ -180,22 +180,22 @@ function getOnTestHandler (isMain, newTests) {
       testStartLine
     }
 
-    if (isMain) {
-      testInfo.isNew = isNew
-      testInfo.isEfdRetry = isEfdRetry
-      // We want to store the result of the new tests
-      if (isNew) {
-        const testFullName = getTestFullName(test)
-        if (newTests[testFullName]) {
-          newTests[testFullName].push(test)
-        } else {
-          newTests[testFullName] = [test]
-        }
-      }
-    } else {
+    if (!isMain) {
       testInfo.isParallel = true
     }
 
+    testInfo.isNew = isNew
+    testInfo.isEfdRetry = isEfdRetry
+    // We want to store the result of the new tests
+    if (isNew) {
+      const testFullName = getTestFullName(test)
+      if (newTests[testFullName]) {
+        newTests[testFullName].push(test)
+      } else {
+        newTests[testFullName] = [test]
+      }
+    }
+
     asyncResource.runInAsyncScope(() => {
       testStartCh.publish(testInfo)
     })
@@ -328,6 +328,23 @@ function getOnPendingHandler () {
     }
   }
 }
+
+// Hook to add retries to tests if EFD is enabled
+function getRunTestsWrapper (runTests, config) {
+  return function (suite, fn) {
+    if (config.isEarlyFlakeDetectionEnabled) {
+      // by the time we reach `this.on('test')`, it is too late. We need to add retries here
+      suite.tests.forEach(test => {
+        if (!test.isPending() && isNewTest(test, config.knownTests)) {
+          test._ddIsNew = true
+          retryTest(test, config.earlyFlakeDetectionNumRetries)
+        }
+      })
+    }
+    return runTests.apply(this, arguments)
+  }
+}
+
 module.exports = {
   isNewTest,
   retryTest,
@@ -346,5 +363,7 @@ module.exports = {
   getOnHookEndHandler,
   getOnFailHandler,
   getOnPendingHandler,
-  testFileToSuiteAr
+  testFileToSuiteAr,
+  getRunTestsWrapper,
+  newTests
 }

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

@@ -9,19 +9,47 @@ const {
   getOnTestEndHandler,
   getOnHookEndHandler,
   getOnFailHandler,
-  getOnPendingHandler
+  getOnPendingHandler,
+  getRunTestsWrapper
 } = require('./utils')
 require('./common')
 
 const workerFinishCh = channel('ci:mocha:worker:finish')
 
+const config = {}
+
+addHook({
+  name: 'mocha',
+  versions: ['>=8.0.0'],
+  file: 'lib/mocha.js'
+}, (Mocha) => {
+  shimmer.wrap(Mocha.prototype, 'run', run => function () {
+    if (this.options._ddKnownTests) {
+      // EFD is enabled if there's a list of known tests
+      config.isEarlyFlakeDetectionEnabled = true
+      config.knownTests = this.options._ddKnownTests
+      config.earlyFlakeDetectionNumRetries = this.options._ddEfdNumRetries
+      delete this.options._ddKnownTests
+      delete this.options._ddEfdNumRetries
+    }
+    return run.apply(this, arguments)
+  })
+
+  return Mocha
+})
+
 // Runner is also hooked in mocha/main.js, but in here we only generate test events.
 addHook({
   name: 'mocha',
   versions: ['>=5.2.0'],
   file: 'lib/runner.js'
 }, function (Runner) {
+  shimmer.wrap(Runner.prototype, 'runTests', runTests => getRunTestsWrapper(runTests, config))
+
   shimmer.wrap(Runner.prototype, 'run', run => function () {
+    if (!workerFinishCh.hasSubscribers) {
+      return run.apply(this, arguments)
+    }
     // We flush when the worker ends with its test file (a mocha instance in a worker runs a single test file)
     this.on('end', () => {
       workerFinishCh.publish()

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()
           }
         })
-      }
+      })
     }
   }
 }

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

@@ -92,7 +92,7 @@ function wrapUnifiedCommand (command, operation, name) {
     }
     return instrument(operation, command, this, arguments, server, ns, ops, { name })
   }
-  return shimmer.wrap(command, wrapped)
+  return wrapped
 }
 
 function wrapConnectionCommand (command, operation, name, instrumentFn = instrument) {
@@ -109,7 +109,7 @@ function wrapConnectionCommand (command, operation, name, instrumentFn = instrum
     ns = `${ns.db}.${ns.collection}`
     return instrumentFn(operation, command, this, arguments, topology, ns, ops, { name })
   }
-  return shimmer.wrap(command, wrapped)
+  return wrapped
 }
 
 function wrapQuery (query, operation, name) {
@@ -123,7 +123,7 @@ function wrapQuery (query, operation, name) {
     return instrument(operation, query, this, arguments, pool, ns, ops)
   }
 
-  return shimmer.wrap(query, wrapped)
+  return wrapped
 }
 
 function wrapCursor (cursor, operation, name) {
@@ -135,7 +135,7 @@ function wrapCursor (cursor, operation, name) {
     const ns = this.ns
     return instrument(operation, cursor, this, arguments, pool, ns, {}, { name })
   }
-  return shimmer.wrap(cursor, wrapped)
+  return wrapped
 }
 
 function wrapCommand (command, operation, name) {
@@ -145,7 +145,7 @@ function wrapCommand (command, operation, name) {
     }
     return instrument(operation, command, this, arguments, this, ns, ops, { name })
   }
-  return shimmer.wrap(command, wrapped)
+  return wrapped
 }
 
 function instrument (operation, command, ctx, args, server, ns, ops, options = {}) {
@@ -164,7 +164,7 @@ function instrument (operation, command, ctx, args, server, ns, ops, options = {
   return asyncResource.runInAsyncScope(() => {
     startCh.publish({ ns, ops, options: serverInfo, name })
 
-    args[index] = asyncResource.bind(function (err, res) {
+    args[index] = shimmer.wrapFunction(callback, callback => asyncResource.bind(function (err, res) {
       if (err) {
         errorCh.publish(err)
       }
@@ -174,7 +174,7 @@ function instrument (operation, command, ctx, args, server, ns, ops, options = {
       if (callback) {
         return callback.apply(this, arguments)
       }
-    })
+    }))
 
     try {
       return command.apply(ctx, args)

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

@@ -128,22 +128,21 @@ addHook({
                     const resolve = arguments[0]
                     const reject = arguments[1]
 
-                    // not using shimmer here because resolve/reject could be empty
-                    arguments[0] = function wrappedResolve () {
+                    arguments[0] = shimmer.wrapFunction(resolve, resolve => function wrappedResolve () {
                       finish()
 
                       if (resolve) {
                         return resolve.apply(this, arguments)
                       }
-                    }
+                    })
 
-                    arguments[1] = function wrappedReject () {
+                    arguments[1] = shimmer.wrapFunction(reject, reject => function wrappedReject () {
                       finish()
 
                       if (reject) {
                         return reject.apply(this, arguments)
                       }
-                    }
+                    })
 
                     return originalThen.apply(this, arguments)
                   }
@@ -169,7 +168,7 @@ addHook({
   versions: ['6', '>=7'],
   file: 'lib/helpers/query/sanitizeFilter.js'
 }, sanitizeFilter => {
-  return shimmer.wrap(sanitizeFilter, function wrappedSanitizeFilter () {
+  return shimmer.wrapFunction(sanitizeFilter, sanitizeFilter => function wrappedSanitizeFilter () {
     const sanitizedObject = sanitizeFilter.apply(this, arguments)
 
     if (sanitizeFilterFinishCh.hasSubscribers) {

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

@@ -37,14 +37,14 @@ addHook({ name: 'mysql', file: 'lib/Connection.js', versions: ['>=2'] }, Connect
 
         if (res._callback) {
           const cb = callbackResource.bind(res._callback)
-          res._callback = asyncResource.bind(function (error, result) {
+          res._callback = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (error, result) {
             if (error) {
               errorCh.publish(error)
             }
             finishCh.publish(result)
 
             return cb.apply(this, arguments)
-          })
+          }))
         } else {
           const cb = asyncResource.bind(function () {
             finishCh.publish(undefined)
@@ -92,7 +92,7 @@ addHook({ name: 'mysql', file: 'lib/Pool.js', versions: ['>=2'] }, Pool => {
 
       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/mysql2.js

@@ -31,19 +31,19 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connec
   return Connection
 
   function bindExecute (cmd, execute, asyncResource) {
-    return asyncResource.bind(function executeWithTrace (packet, connection) {
+    return shimmer.wrapFunction(execute, execute => asyncResource.bind(function executeWithTrace (packet, connection) {
       if (this.onResult) {
         this.onResult = asyncResource.bind(this.onResult)
       }
 
       return execute.apply(this, arguments)
-    }, cmd)
+    }, cmd))
   }
 
   function wrapExecute (cmd, execute, asyncResource, config) {
     const callbackResource = new AsyncResource('bound-anonymous-fn')
 
-    return asyncResource.bind(function executeWithTrace (packet, connection) {
+    return shimmer.wrapFunction(execute, execute => asyncResource.bind(function executeWithTrace (packet, connection) {
       const sql = cmd.statement ? cmd.statement.query : cmd.sql
       const payload = { sql, conf: config }
       startCh.publish(payload)
@@ -57,13 +57,13 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connec
       if (this.onResult) {
         const onResult = callbackResource.bind(this.onResult)
 
-        this.onResult = asyncResource.bind(function (error) {
+        this.onResult = shimmer.wrapFunction(onResult, onResult => asyncResource.bind(function (error) {
           if (error) {
             errorCh.publish(error)
           }
           finishCh.publish(undefined)
           onResult.apply(this, arguments)
-        }, 'bound-anonymous-fn', this)
+        }, 'bound-anonymous-fn', this))
       } else {
         this.on('error', asyncResource.bind(error => errorCh.publish(error)))
         this.on('end', asyncResource.bind(() => finishCh.publish(undefined)))
@@ -76,6 +76,6 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connec
       } catch (err) {
         errorCh.publish(err)
       }
-    }, cmd)
+    }, cmd))
   }
 })