dd-trace 5.103.0 → 5.105.0

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

@@ -13,12 +13,12 @@ const tracers = new WeakSet()
 const wrappedModels = new WeakSet()
 
 /**
- * Publishes already-converted AI guard style messages to the AIGuard channel.
+ * Publishes already-converted AI-style messages to the AI Guard evaluation channel.
  *
- * @param {Array<object>} messages - AI guard style messages to evaluate
+ * @param {Array<object>} messages - AI-style messages to evaluate.
  * @returns {Promise<void>}
  */
-function publishToAIGuard (messages) {
+function publishEvaluation (messages) {
   return new Promise((resolve, reject) => {
     aiguardChannel.publish({ messages, integration: 'ai', resolve, reject })
   })
@@ -47,10 +47,11 @@ function wrapModelWithAIGuard (model) {
 
         // Run AI Guard input evaluation and LLM call in parallel.
         // The LLM has no side effects so it is safe to discard its result if AI Guard blocks.
-        return Promise.all([publishToAIGuard(inputMessages), originalResult])
+        return Promise.all([publishEvaluation(inputMessages), originalResult])
           .then(([, result]) => {
             if (!result.content?.length) return result
-            return publishToAIGuard(buildOutputMessages(inputMessages, result.content))
+            const outputMessages = buildOutputMessages(inputMessages, result.content)
+            return publishEvaluation(outputMessages)
               .then(() => result)
           })
       }
@@ -70,7 +71,7 @@ function wrapModelWithAIGuard (model) {
 
         // Run AI Guard input evaluation and LLM call in parallel.
         // The LLM has no side effects so it is safe to discard its result if AI Guard blocks.
-        return Promise.all([publishToAIGuard(inputMessages), originalResult])
+        return Promise.all([publishEvaluation(inputMessages), originalResult])
           .then(([, result]) => {
             const chunks = []
             const reader = result.stream.getReader()
@@ -89,7 +90,7 @@ function wrapModelWithAIGuard (model) {
               const content = toolCalls.length ? toolCalls : text ? [{ type: 'text', text }] : []
 
               const evaluate = content.length
-                ? publishToAIGuard(buildOutputMessages(inputMessages, content))
+                ? publishEvaluation(buildOutputMessages(inputMessages, content))
                 : Promise.resolve()
 
               return evaluate.then(() => {

package/packages/datadog-instrumentations/src/aws-sdk.js

@@ -198,7 +198,7 @@ function wrapSmithySend (send) {
       })
 
       if (typeof cb === 'function') {
-        args[args.length - 1] = shimmer.wrapFunction(cb, cb => function (err, result) {
+        args[args.length - 1] = shimmer.wrapCallback(cb, cb => function (err, result) {
           addResponse(ctx, err, result)
 
           handleCompletion(result, ctx, channels)
@@ -270,7 +270,7 @@ function handleCompletion (result, ctx, channels) {
 
 function wrapCb (cb, channels, ctx) {
   // eslint-disable-next-line n/handle-callback-err
-  return shimmer.wrapFunction(cb, cb => function wrappedCb (err, response) {
+  return shimmer.wrapCallback(cb, cb => function wrappedCb (err, response) {
     ctx = { request: ctx.request, response }
     return channels.responseStart.runStores(ctx, () => {
       try {
@@ -324,6 +324,19 @@ addHook({ name: '@aws-sdk/smithy-client', versions: ['>=3'] }, smithy => {
   return smithy
 })
 
+// `@aws-sdk/client-*` >= 3.1046.0 dropped `@smithy/smithy-client` and now
+// extends from `@smithy/core/client` directly. The `Client.send` contract is
+// unchanged, but the host module moved -- patch the new home so the v3 hooks
+// keep firing.
+addHook({
+  name: '@smithy/core',
+  file: 'dist-cjs/submodules/client/index.js',
+  versions: ['>=3.24.0'],
+}, smithyCoreClient => {
+  shimmer.wrap(smithyCoreClient.Client.prototype, 'send', wrapSmithySend)
+  return smithyCoreClient
+})
+
 addHook({ name: 'aws-sdk', versions: ['>=2.3.0'] }, AWS => {
   shimmer.wrap(AWS.config, 'setPromisesDependency', setPromisesDependency => {
     return function wrappedSetPromisesDependency (dep) {

package/packages/datadog-instrumentations/src/azure-cosmos.js

@@ -0,0 +1,7 @@
+'use strict'
+
+const { addHook, getHooks } = require('./helpers/instrument')
+
+for (const hook of getHooks('@azure/cosmos')) {
+  addHook(hook, exports => exports)
+}

package/packages/datadog-instrumentations/src/azure-functions.js

@@ -26,6 +26,9 @@ addHook({ name: '@azure/functions', versions: ['>=4'], patchDefault: false }, (a
   // Event Hub triggers
   shimmer.wrap(app, 'eventHub', wrapHandler)
 
+  // CosmosDB triggers
+  shimmer.wrap(app, 'cosmosDB', wrapHandler)
+
   return azureFunction
 })
 

package/packages/datadog-instrumentations/src/cassandra-driver.js

@@ -29,9 +29,12 @@ addHook({ name: 'cassandra-driver', versions: ['>=3.0.0'] }, cassandra => {
 
       try {
         const res = batch.apply(this, arguments)
-        if (typeof res === 'function' || !res) {
+        if (typeof res === 'function') {
           return wrapCallback(finishCh, errorCh, startCtx, res)
         }
+        if (!res) {
+          return res
+        }
         return res.then(
           () => finish(finishCh, errorCh, startCtx),
           err => finish(finishCh, errorCh, startCtx, err)
@@ -162,7 +165,7 @@ function finish (finishCh, errorCh, ctx, error) {
 }
 
 function wrapCallback (finishCh, errorCh, ctx, callback) {
-  return shimmer.wrapFunction(callback, callback => function (err) {
+  return shimmer.wrapCallback(callback, callback => function (err) {
     if (err) {
       ctx.error = err
       errorCh.publish(ctx)

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

@@ -7,19 +7,26 @@ const shimmer = require('../../datadog-shimmer')
 const log = require('../../dd-trace/src/log')
 const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
 const {
-  getCoveredFilenamesFromCoverage,
+  getCoveredFilesFromCoverage,
+  getExecutableFilesFromCoverage,
   resetCoverage,
   mergeCoverage,
   fromCoverageMapToCoverage,
   getTestSuitePath,
+  getRelativeCoverageFiles,
   CUCUMBER_WORKER_TRACE_PAYLOAD_CODE,
   getIsFaultyEarlyFlakeDetection,
   getEfdRetryCount,
+  getMaxEfdRetryCount,
+  applySkippedCoverageToCoverage,
+  getTestCoverageLinesPercentage,
   recordAttemptToFixExecution,
   collectAttemptToFixExecutionsFromTraces,
   logAttemptToFixTestExecution,
   logTestOptimizationSummary,
+  getTestOptimizationRequestResults,
 } = require('../../dd-trace/src/plugins/util/test')
+const { writeCoverageBackfillToCache } = require('../../dd-trace/src/ci-visibility/test-optimization-cache')
 const satisfies = require('../../../vendor/dist/semifies')
 const { addHook, channel } = require('./helpers/instrument')
 
@@ -51,6 +58,8 @@ const itrSkippedSuitesCh = channel('ci:cucumber:itr:skipped-suites')
 
 const getCodeCoverageCh = channel('ci:nyc:get-coverage')
 
+const DD_EFD_RETRY_COUNT_MESSAGE = '_ddEfdRetryCount'
+
 const isMarkedAsUnskippable = (pickle) => {
   return pickle.tags.some(tag => tag.name === '@datadog:unskippable')
 }
@@ -67,7 +76,7 @@ const atrStatusesByScenarioKey = new Map()
 const numRetriesByPickleId = new Map()
 const efdRetryCountByPickleId = new Map()
 const efdSlowAbortedPickleIds = new Set()
-const testCaseStartedTimesById = new Map()
+const finishedParallelSuites = new Set()
 const numAttemptToCtx = new Map()
 const newTestsByTestFullname = new Map()
 const attemptToFixTestsByTestFullname = new Map()
@@ -82,10 +91,14 @@ let pickleByFile = {}
 const pickleResultByFile = {}
 
 let skippableSuites = []
+let skippableSuitesCoverage = {}
+let skippedSuitesCoverage = {}
 let itrCorrelationId = ''
 let isForcedToRun = false
 let isUnskippable = false
+let isItrEnabled = false
 let isSuitesSkippingEnabled = false
+let isCoverageReportUploadEnabled = false
 let isEarlyFlakeDetectionEnabled = false
 let earlyFlakeDetectionNumRetries = 0
 let earlyFlakeDetectionSlowTestRetries = {}
@@ -102,11 +115,55 @@ let numTestRetries = 0
 let knownTests = {}
 let skippedSuites = []
 let isSuitesSkipped = false
+let repositoryRoot
 
 function isValidKnownTests (receivedKnownTests) {
   return !!receivedKnownTests.cucumber
 }
 
+function hasSkippableSuitesCoverage () {
+  return skippableSuitesCoverage &&
+    typeof skippableSuitesCoverage === 'object' &&
+    Object.keys(skippableSuitesCoverage).length > 0
+}
+
+function isTiaCoverageBackfillEnabled () {
+  return isItrEnabled && isCoverageReportUploadEnabled
+}
+
+function getCoverageRootDir () {
+  return repositoryRoot || process.cwd()
+}
+
+function shouldReportCodeCoverageLinesPct (hasBackfilledCoverage) {
+  return !isSuitesSkipped || hasBackfilledCoverage
+}
+
+function getSkippedSuitesCoverageForRun () {
+  return isSuitesSkipped && isTiaCoverageBackfillEnabled() && hasSkippableSuitesCoverage()
+    ? skippableSuitesCoverage
+    : {}
+}
+
+function applySkippedCoverageToCucumberCoverageMap () {
+  if (!isTiaCoverageBackfillEnabled()) return false
+  return applySkippedCoverageToCoverage(originalCoverageMap, skippedSuitesCoverage, getCoverageRootDir())
+}
+
+function getCucumberTestSessionCoverageFiles () {
+  return getRelativeCoverageFiles(getExecutableFilesFromCoverage(originalCoverageMap), getCoverageRootDir())
+}
+
+function resetSuiteSkippingRunState () {
+  skippableSuites = []
+  skippableSuitesCoverage = {}
+  skippedSuitesCoverage = {}
+  skippedSuites = []
+  isSuitesSkipped = false
+  repositoryRoot = undefined
+  writeCoverageBackfillToCache({})
+}
+
 function getSuiteStatusFromTestStatuses (testStatuses) {
   if (testStatuses.includes('fail')) {
     return 'fail'
@@ -117,6 +174,68 @@ function getSuiteStatusFromTestStatuses (testStatuses) {
   return 'pass'
 }
 
+function getConfiguredEfdRetryCount () {
+  const maxSlowTestRetryCount = getMaxEfdRetryCount(earlyFlakeDetectionSlowTestRetries)
+  return maxSlowTestRetryCount || earlyFlakeDetectionNumRetries
+}
+
+function publishWorkerEfdRetryCount (pickle, retryCount) {
+  if (typeof process.send !== 'function') return
+
+  try {
+    process.send({
+      [DD_EFD_RETRY_COUNT_MESSAGE]: {
+        pickleId: pickle.id,
+        retryCount,
+        testFileAbsolutePath: pickle.uri,
+        testName: pickle.name,
+      },
+    })
+  } catch {
+    // ignore IPC errors
+  }
+}
+
+function finishParallelSuiteIfDone (testFileAbsolutePath) {
+  const finished = pickleResultByFile[testFileAbsolutePath]
+  const expectedPickles = pickleByFile[testFileAbsolutePath]
+
+  if (!finished || !expectedPickles || finished.length !== expectedPickles.length) return
+  if (finishedParallelSuites.has(testFileAbsolutePath)) return
+
+  finishedParallelSuites.add(testFileAbsolutePath)
+  testSuiteFinishCh.publish({
+    status: getSuiteStatusFromTestStatuses(finished),
+    testSuitePath: getTestSuitePath(testFileAbsolutePath, process.cwd()),
+  })
+}
+
+function maybeRecordFinalParallelEfdStatus ({ pickleId, testFileAbsolutePath, testFullname }) {
+  const efdRetryCount = efdRetryCountByPickleId.get(pickleId)
+  const testStatuses = newTestsByTestFullname.get(testFullname)
+  const finished = pickleResultByFile[testFileAbsolutePath]
+
+  if (efdRetryCount === undefined || !testStatuses || !finished) return
+  if (testStatuses.length !== efdRetryCount + 1) return
+
+  finished.push(getTestStatusFromRetries(testStatuses))
+  newTestsByTestFullname.delete(testFullname)
+  finishParallelSuiteIfDone(testFileAbsolutePath)
+}
+
+function handleEfdRetryCountMessage (message) {
+  const { pickleId, retryCount, testFileAbsolutePath, testName } = message
+
+  if (!pickleId || typeof retryCount !== 'number' || !testFileAbsolutePath || !testName) return
+
+  efdRetryCountByPickleId.set(pickleId, retryCount)
+  maybeRecordFinalParallelEfdStatus({
+    pickleId,
+    testFileAbsolutePath,
+    testFullname: `${testFileAbsolutePath}:${testName}`,
+  })
+}
+
 function getStatusFromResult (result) {
   if (result.status === 1) {
     return { status: 'pass' }
@@ -617,6 +736,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     if (!libraryConfigurationCh.hasSubscribers) {
       return start.apply(this, arguments)
     }
+    resetSuiteSkippingRunState()
     const options = getCucumberOptions(this)
 
     if (!isParallel && this.adapter?.options) {
@@ -626,11 +746,14 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
 
     const configurationResponse = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
 
+    repositoryRoot = configurationResponse.repositoryRoot
+    isItrEnabled = configurationResponse.libraryConfig?.isItrEnabled
     isEarlyFlakeDetectionEnabled = configurationResponse.libraryConfig?.isEarlyFlakeDetectionEnabled
     earlyFlakeDetectionNumRetries = configurationResponse.libraryConfig?.earlyFlakeDetectionNumRetries
     earlyFlakeDetectionSlowTestRetries = configurationResponse.libraryConfig?.earlyFlakeDetectionSlowTestRetries ?? {}
     earlyFlakeDetectionFaultyThreshold = configurationResponse.libraryConfig?.earlyFlakeDetectionFaultyThreshold
-    isSuitesSkippingEnabled = configurationResponse.libraryConfig?.isSuitesSkippingEnabled
+    isSuitesSkippingEnabled = isItrEnabled && configurationResponse.libraryConfig?.isSuitesSkippingEnabled
+    isCoverageReportUploadEnabled = configurationResponse.libraryConfig?.isCoverageReportUploadEnabled
     isFlakyTestRetriesEnabled = configurationResponse.libraryConfig?.isFlakyTestRetriesEnabled
     const configRetryCount = configurationResponse.libraryConfig?.flakyTestRetriesCount
     numTestRetries = (typeof configRetryCount === 'number' && configRetryCount > 0) ? configRetryCount : 0
@@ -639,21 +762,35 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     testManagementAttemptToFixRetries = configurationResponse.libraryConfig?.testManagementAttemptToFixRetries
     isImpactedTestsEnabled = configurationResponse.libraryConfig?.isImpactedTestsEnabled
 
+    const {
+      knownTestsResponse,
+      testManagementTestsResponse,
+      skippableSuitesResponse,
+    } = await getTestOptimizationRequestResults({
+      isKnownTestsEnabled,
+      isTestManagementTestsEnabled,
+      isSuitesSkippingEnabled,
+      getKnownTests: () => getChannelPromise(knownTestsCh),
+      getTestManagementTests: () => getChannelPromise(testManagementTestsCh),
+      getSkippableSuites: () => getChannelPromise(skippableSuitesCh),
+    })
+
     if (isKnownTestsEnabled) {
-      const knownTestsResponse = await getChannelPromise(knownTestsCh)
-      if (knownTestsResponse.err) {
+      const currentKnownTestsResponse = knownTestsResponse || await getChannelPromise(knownTestsCh)
+      if (currentKnownTestsResponse.err) {
         isEarlyFlakeDetectionEnabled = false
         isKnownTestsEnabled = false
       } else {
-        knownTests = knownTestsResponse.knownTests
+        knownTests = currentKnownTestsResponse.knownTests
       }
     }
 
     if (isSuitesSkippingEnabled) {
-      const skippableResponse = await getChannelPromise(skippableSuitesCh)
+      const skippableResponse = skippableSuitesResponse || await getChannelPromise(skippableSuitesCh)
 
       errorSkippableRequest = skippableResponse.err
       skippableSuites = skippableResponse.skippableSuites ?? []
+      skippableSuitesCoverage = skippableResponse.skippableSuitesCoverage ?? {}
 
       if (!errorSkippableRequest) {
         const filteredPickles = isCoordinator
@@ -674,6 +811,8 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
         }
 
         skippedSuites = [...filteredPickles.skippedSuites]
+        skippedSuitesCoverage = getSkippedSuitesCoverageForRun()
+        writeCoverageBackfillToCache(skippedSuitesCoverage, getCoverageRootDir())
         itrCorrelationId = skippableResponse.itrCorrelationId
       }
     }
@@ -694,11 +833,12 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     }
 
     if (isTestManagementTestsEnabled) {
-      const testManagementTestsResponse = await getChannelPromise(testManagementTestsCh)
-      if (testManagementTestsResponse.err) {
+      const currentTestManagementTestsResponse =
+        testManagementTestsResponse || await getChannelPromise(testManagementTestsCh)
+      if (currentTestManagementTestsResponse.err) {
         isTestManagementTestsEnabled = false
       } else {
-        testManagementTests = testManagementTestsResponse.testManagementTests
+        testManagementTests = currentTestManagementTestsResponse.testManagementTests
       }
     }
 
@@ -720,7 +860,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     attemptToFixTestsByTestFullname.clear()
     efdRetryCountByPickleId.clear()
     efdSlowAbortedPickleIds.clear()
-    testCaseStartedTimesById.clear()
+    finishedParallelSuites.clear()
     newTestsByTestFullname.clear()
     sessionStartCh.publish({ command, frameworkVersion })
 
@@ -736,13 +876,25 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     }
 
     let testCodeCoverageLinesTotal
+    let testSessionCoverageFiles
 
-    if (global.__coverage__) {
+    if (global.__coverage__ || untestedCoverage) {
       try {
+        let hasBackfilledCoverage = false
         if (untestedCoverage) {
           originalCoverageMap.merge(fromCoverageMapToCoverage(untestedCoverage))
         }
-        testCodeCoverageLinesTotal = originalCoverageMap.getCoverageSummary().lines.pct
+        hasBackfilledCoverage = applySkippedCoverageToCucumberCoverageMap()
+        if (shouldReportCodeCoverageLinesPct(hasBackfilledCoverage)) {
+          testCodeCoverageLinesTotal = getTestCoverageLinesPercentage(
+            originalCoverageMap,
+            undefined,
+            getCoverageRootDir()
+          )
+        }
+        if (isTiaCoverageBackfillEnabled()) {
+          testSessionCoverageFiles = getCucumberTestSessionCoverageFiles()
+        }
       } catch {
         // ignore errors
       }
@@ -754,6 +906,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
       status: success ? 'pass' : 'fail',
       isSuitesSkipped,
       testCodeCoverageLinesTotal,
+      testSessionCoverageFiles,
       numSkippedSuites: skippedSuites.length,
       hasUnskippableSuites: isUnskippable,
       hasForcedToRunSuites: isForcedToRun,
@@ -875,6 +1028,9 @@ function getWrappedRunTestCase (runTestCaseFunction, isNewerCucumberVersion = fa
           efdSlowAbortedPickleIds.add(pickle.id)
         }
       }
+      if (isWorker) {
+        publishWorkerEfdRetryCount(pickle, efdRetryCount)
+      }
       for (let retryIndex = 0; retryIndex < efdRetryCount; retryIndex++) {
         numRetriesByPickleId.set(pickle.id, retryIndex + 1)
         // eslint-disable-next-line no-await-in-loop
@@ -925,7 +1081,7 @@ function getWrappedRunTestCase (runTestCaseFunction, isNewerCucumberVersion = fa
       // last test in suite
       const testSuiteStatus = getSuiteStatusFromTestStatuses(pickleResultByFile[testFileAbsolutePath])
       if (global.__coverage__) {
-        const coverageFiles = getCoveredFilenamesFromCoverage(global.__coverage__)
+        const coverageFiles = getCoveredFilesFromCoverage(global.__coverage__)
 
         testSuiteCodeCoverageCh.publish({
           coverageFiles,
@@ -973,6 +1129,10 @@ function getWrappedParseWorkerMessage (parseWorkerMessageFunction, isNewVersion)
         return
       }
     }
+    if (message[DD_EFD_RETRY_COUNT_MESSAGE]) {
+      handleEfdRetryCountMessage(message[DD_EFD_RETRY_COUNT_MESSAGE])
+      return
+    }
 
     const envelope = isNewVersion ? message.envelope : message.jsonEnvelope
 
@@ -989,12 +1149,13 @@ function getWrappedParseWorkerMessage (parseWorkerMessageFunction, isNewVersion)
         return parseWorkerMessageFunction.apply(this, arguments)
       }
     }
+    if (parsed[DD_EFD_RETRY_COUNT_MESSAGE]) {
+      handleEfdRetryCountMessage(parsed[DD_EFD_RETRY_COUNT_MESSAGE])
+      return
+    }
     let pickle
 
     if (parsed.testCaseStarted) {
-      if (parsed.testCaseStarted.id) {
-        testCaseStartedTimesById.set(parsed.testCaseStarted.id, performance.now())
-      }
       if (isNewVersion) {
         pickle = this.inProgress[worker.id].pickle
       } else {
@@ -1016,10 +1177,6 @@ function getWrappedParseWorkerMessage (parseWorkerMessageFunction, isNewVersion)
 
     // after calling `parseWorkerMessageFunction`, the test status can already be read
     if (parsed.testCaseFinished) {
-      const testCaseStartedId = parsed.testCaseFinished.testCaseStartedId
-      const testCaseStartedAt = testCaseStartedTimesById.get(testCaseStartedId)
-      testCaseStartedTimesById.delete(testCaseStartedId)
-
       let worstTestStepResult
       if (isNewVersion && eventDataCollector) {
         pickle = this.inProgress[worker.id].pickle
@@ -1052,21 +1209,15 @@ function getWrappedParseWorkerMessage (parseWorkerMessageFunction, isNewVersion)
         }
         let efdRetryCount = efdRetryCountByPickleId.get(pickle.id)
         if (efdRetryCount === undefined) {
-          const firstExecutionDurationMs = testCaseStartedAt === undefined ? 0 : performance.now() - testCaseStartedAt
           efdRetryCount = status === 'skip'
             ? 0
-            : getEfdRetryCount(firstExecutionDurationMs, earlyFlakeDetectionSlowTestRetries)
+            : getConfiguredEfdRetryCount()
           efdRetryCountByPickleId.set(pickle.id, efdRetryCount)
           if (efdRetryCount === 0 && status !== 'skip') {
             efdSlowAbortedPickleIds.add(pickle.id)
           }
         }
-        // We have finished all retries
-        if (testStatuses.length === efdRetryCount + 1) {
-          const newTestFinalStatus = getTestStatusFromRetries(testStatuses)
-          // we only push to `finished` if the retries have finished
-          finished.push(newTestFinalStatus)
-        }
+        maybeRecordFinalParallelEfdStatus({ pickleId: pickle.id, testFileAbsolutePath, testFullname })
       } else if (
         isTestManagementTestsEnabled &&
         getTestProperties(getTestSuitePath(testFileAbsolutePath, process.cwd()), pickle.name).attemptToFix
@@ -1090,12 +1241,7 @@ function getWrappedParseWorkerMessage (parseWorkerMessageFunction, isNewVersion)
         finished.push(status)
       }
 
-      if (finished.length === pickleByFile[testFileAbsolutePath].length) {
-        testSuiteFinishCh.publish({
-          status: getSuiteStatusFromTestStatuses(finished),
-          testSuitePath: getTestSuitePath(testFileAbsolutePath, process.cwd()),
-        })
-      }
+      finishParallelSuiteIfDone(testFileAbsolutePath)
     }
 
     return parseWorkerResponse

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

@@ -3,6 +3,7 @@
 const shimmer = require('../../datadog-shimmer')
 const { addHook } = require('./helpers/instrument')
 const { createCallbackInstrumentor } = require('./helpers/callback-instrumentor')
+const { createPromiseInstrumentor } = require('./helpers/promise-instrumentor')
 
 const rrtypes = {
   resolveAny: 'ANY',
@@ -18,30 +19,55 @@ const rrtypes = {
   resolveSoa: 'SOA',
 }
 
-addHook({ name: 'dns' }, dns => {
-  const lookup = createCallbackInstrumentor('apm:dns:lookup', { captureResult: true })
-  const lookupService = createCallbackInstrumentor('apm:dns:lookup_service', { captureResult: true })
-  const resolve = createCallbackInstrumentor('apm:dns:resolve', { captureResult: true })
-  const reverse = createCallbackInstrumentor('apm:dns:reverse', { captureResult: true })
-
-  shimmer.wrap(dns, 'lookup', lookup(buildArgsContext()))
-  shimmer.wrap(dns, 'lookupService', lookupService(buildArgsContext()))
-  shimmer.wrap(dns, 'resolve', resolve(buildArgsContext()))
-  shimmer.wrap(dns, 'reverse', reverse(buildArgsContext()))
-
-  patchResolveShorthands(dns, resolve)
+// `dns.promises` and `require('dns/promises')` resolve to the same exports object. Both
+// access paths register a hook, so without a guard the second hook to fire would stack a
+// second wrap layer on top and publish every `apm:dns:*` event twice per call. The WeakSet
+// collapses the two hooks to one wrap regardless of which one runs first.
+const wrappedPromiseApis = new WeakSet()
 
-  if (dns.Resolver) {
-    shimmer.wrap(dns.Resolver.prototype, 'resolve', resolve(buildArgsContext()))
-    shimmer.wrap(dns.Resolver.prototype, 'reverse', reverse(buildArgsContext()))
+addHook({ name: 'dns' }, dns => {
+  patchApi(dns, createCallbackInstrumentor, buildCallbackArgsContext)
 
-    patchResolveShorthands(dns.Resolver.prototype, resolve)
+  if (dns.promises) {
+    patchPromiseApi(dns.promises)
   }
 
   return dns
 })
 
-function patchResolveShorthands (prototype, resolve) {
+addHook({ name: 'dns/promises' }, dnsPromises => {
+  patchPromiseApi(dnsPromises)
+  return dnsPromises
+})
+
+function patchPromiseApi (api) {
+  if (wrappedPromiseApis.has(api)) return
+  wrappedPromiseApis.add(api)
+  patchApi(api, createPromiseInstrumentor, buildPromiseArgsContext)
+}
+
+function patchApi (api, instrumentorFactory, buildArgsContext) {
+  const lookup = instrumentorFactory('apm:dns:lookup', { captureResult: true })
+  const lookupService = instrumentorFactory('apm:dns:lookup_service', { captureResult: true })
+  const resolve = instrumentorFactory('apm:dns:resolve', { captureResult: true })
+  const reverse = instrumentorFactory('apm:dns:reverse', { captureResult: true })
+
+  shimmer.wrap(api, 'lookup', lookup(buildArgsContext()))
+  shimmer.wrap(api, 'lookupService', lookupService(buildArgsContext()))
+  shimmer.wrap(api, 'resolve', resolve(buildArgsContext()))
+  shimmer.wrap(api, 'reverse', reverse(buildArgsContext()))
+
+  patchResolveShorthands(api, resolve, buildArgsContext)
+
+  if (api.Resolver) {
+    shimmer.wrap(api.Resolver.prototype, 'resolve', resolve(buildArgsContext()))
+    shimmer.wrap(api.Resolver.prototype, 'reverse', reverse(buildArgsContext()))
+
+    patchResolveShorthands(api.Resolver.prototype, resolve, buildArgsContext)
+  }
+}
+
+function patchResolveShorthands (prototype, resolve, buildArgsContext) {
   for (const method of Object.keys(rrtypes)) {
     if (prototype[method]) {
       shimmer.wrap(prototype, method, resolve(buildArgsContext(rrtypes[method])))
@@ -49,7 +75,7 @@ function patchResolveShorthands (prototype, resolve) {
   }
 }
 
-function buildArgsContext (rrtype) {
+function buildCallbackArgsContext (rrtype) {
   return function (_, args) {
     if (args.length < 2) return
     const captured = [...args]
@@ -60,3 +86,13 @@ function buildArgsContext (rrtype) {
     return { args: captured }
   }
 }
+
+function buildPromiseArgsContext (rrtype) {
+  return function (_, args) {
+    const captured = [...args]
+    if (rrtype) {
+      captured.push(rrtype)
+    }
+    return { args: captured }
+  }
+}

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

@@ -51,9 +51,9 @@ function createWrapSelect () {
   const connectCh = channel('apm:elasticsearch:query:connect')
   return function wrapRequest (request) {
     return function (...args) {
-      if (args.length === 1) {
-        const cb = args[0]
-        args[0] = shimmer.wrapFunction(cb, cb => function (err, connection) {
+      const cb = args[0]
+      if (args.length === 1 && typeof cb === 'function') {
+        args[0] = shimmer.wrapCallback(cb, cb => function (err, connection) {
           if (connectCh.hasSubscribers && connection && connection.host) {
             connectCh.publish({ hostname: connection.host.host, port: connection.host.port })
           }
@@ -85,7 +85,7 @@ function createWrapRequest (name) {
           cb = arguments[lastIndex]
 
           if (typeof cb === 'function') {
-            arguments[lastIndex] = shimmer.wrapFunction(cb, cb => function (error) {
+            arguments[lastIndex] = shimmer.wrapCallback(cb, cb => function (error) {
               if (error) {
                 ctx.error = error
                 errorCh.publish(ctx)