dd-trace 6.5.0 → 6.6.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "dd-trace",
-  "version": "6.5.0",
+  "version": "6.6.0",
   "description": "Datadog APM tracing client for JavaScript",
   "main": "index.js",
   "typings": "index.d.ts",
@@ -179,7 +179,7 @@
     "@datadog/native-iast-taint-tracking": "4.2.0",
     "@datadog/native-metrics": "3.1.2",
     "@datadog/openfeature-node-server": "2.0.0",
-    "@datadog/pprof": "5.15.1",
+    "@datadog/pprof": "5.17.0",
     "@datadog/wasm-js-rewriter": "5.0.1",
     "@opentelemetry/api": ">=1.0.0 <1.10.0",
     "@opentelemetry/api-logs": "<1.0.0",

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

@@ -28,6 +28,7 @@ const {
 } = 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 { getChannelPromise } = require('./helpers/channel')
 const { addHook, channel } = require('./helpers/instrument')
 
 const cucumberWorkerThreadsPatchModule = require.resolve('./cucumber-worker-threads')
@@ -563,12 +564,6 @@ function getErrorFromCucumberResult (cucumberResult) {
   return error
 }
 
-function getChannelPromise (channelToPublishTo, frameworkVersion = null) {
-  return new Promise(resolve => {
-    channelToPublishTo.publish({ onDone: resolve, frameworkVersion })
-  })
-}
-
 function getShouldBeSkippedSuite (pickle, suitesToSkip) {
   const testSuitePath = getTestSuitePath(pickle.uri, process.cwd())
   const isUnskippable = isMarkedAsUnskippable(pickle)
@@ -1056,7 +1051,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     }
     let errorSkippableRequest
 
-    const configurationResponse = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
+    const configurationResponse = await getChannelPromise(libraryConfigurationCh, { frameworkVersion }) || {}
 
     repositoryRoot = configurationResponse.repositoryRoot
     isItrEnabled = configurationResponse.libraryConfig?.isItrEnabled
@@ -1089,7 +1084,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
 
     if (isKnownTestsEnabled) {
       const currentKnownTestsResponse = knownTestsResponse || await getChannelPromise(knownTestsCh)
-      if (currentKnownTestsResponse.err) {
+      if (!currentKnownTestsResponse || currentKnownTestsResponse.err) {
         isEarlyFlakeDetectionEnabled = false
         isKnownTestsEnabled = false
       } else {
@@ -1098,7 +1093,9 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     }
 
     if (isSuitesSkippingEnabled) {
-      const skippableResponse = skippableSuitesResponse || await getChannelPromise(skippableSuitesCh)
+      const skippableResponse = skippableSuitesResponse ||
+        await getChannelPromise(skippableSuitesCh) ||
+        { err: true }
 
       errorSkippableRequest = skippableResponse.err
       skippableSuites = skippableResponse.skippableSuites ?? []
@@ -1147,7 +1144,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     if (isTestManagementTestsEnabled) {
       const currentTestManagementTestsResponse =
         testManagementTestsResponse || await getChannelPromise(testManagementTestsCh)
-      if (currentTestManagementTestsResponse.err) {
+      if (!currentTestManagementTestsResponse || currentTestManagementTestsResponse.err) {
         isTestManagementTestsEnabled = false
       } else {
         testManagementTests = currentTestManagementTestsResponse.testManagementTests
@@ -1156,7 +1153,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
 
     if (isImpactedTestsEnabled) {
       const impactedTestsResponse = await getChannelPromise(modifiedFilesCh)
-      if (!impactedTestsResponse.err) {
+      if (impactedTestsResponse && !impactedTestsResponse.err) {
         modifiedFiles = impactedTestsResponse.modifiedFiles
       }
     }
@@ -1215,7 +1212,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
       global.__coverage__ = fromCoverageMapToCoverage(originalCoverageMap)
     }
 
-    sessionFinishCh.publish({
+    const flushPromise = getChannelPromise(sessionFinishCh, {
       status: success ? 'pass' : 'fail',
       isSuitesSkipped,
       testCodeCoverageLinesTotal,
@@ -1228,9 +1225,11 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
       isTestManagementTestsEnabled,
       isParallel,
     })
+
     logTestOptimizationSummary({ attemptToFixExecutions })
     loggedAttemptToFixTests.clear()
     eventDataCollector = null
+    await flushPromise
     return result
   }
 }

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

@@ -0,0 +1,74 @@
+'use strict'
+
+/**
+ * @template {unknown[]} T
+ * @param {(...args: T) => void} onDone
+ * @returns {(...args: T) => void}
+ */
+function getCompletion (onDone) {
+  let hasCompleted = false
+  return (...args) => {
+    if (hasCompleted) return
+
+    hasCompleted = true
+    onDone(...args)
+  }
+}
+
+/**
+ * @template {unknown[]} T
+ * @param {import('node:diagnostics_channel').Channel} channel
+ * @param {Record<string, unknown>} payload
+ * @param {(...args: T) => void} onDone
+ * @returns {void}
+ */
+function publishWithCompletion (channel, payload, onDone) {
+  const complete = getCompletion(onDone)
+  channel.publish({ ...payload, onDone: complete })
+  if (!channel.hasSubscribers) complete()
+}
+
+/**
+ * @template {unknown[]} T
+ * @param {import('node:diagnostics_channel').Channel} channel
+ * @param {Record<string, unknown> & { onDone?: (...args: T) => void }} payload
+ * @param {(...args: T) => void} onDone
+ * @returns {void}
+ */
+function runStoresWithCompletion (channel, payload, onDone) {
+  const complete = getCompletion(onDone)
+  payload.onDone = complete
+  channel.runStores(payload, () => {})
+  if (!channel.hasSubscribers) complete()
+}
+
+/**
+ * @template T
+ * @param {import('node:diagnostics_channel').Channel} channel
+ * @param {Record<string, unknown>} [payload]
+ * @returns {Promise<T>}
+ */
+function getChannelPromise (channel, payload = {}) {
+  return new Promise(resolve => {
+    publishWithCompletion(channel, payload, resolve)
+  })
+}
+
+/**
+ * @template T
+ * @param {import('node:diagnostics_channel').Channel} channel
+ * @param {Record<string, unknown>} [payload]
+ * @returns {Promise<T>}
+ */
+function getRunStoresPromise (channel, payload = {}) {
+  return new Promise(resolve => {
+    runStoresWithCompletion(channel, { ...payload }, resolve)
+  })
+}
+
+module.exports = {
+  getChannelPromise,
+  getRunStoresPromise,
+  publishWithCompletion,
+  runStoresWithCompletion,
+}

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

@@ -46,6 +46,10 @@ const {
   addCoverageBackfillUntestedFiles,
   getCoverageBackfillFiles,
 } = require('./jest/coverage-backfill')
+const {
+  getChannelPromise,
+  publishWithCompletion,
+} = require('./helpers/channel')
 const { addHook, channel } = require('./helpers/instrument')
 
 const testSessionStartCh = channel('ci:jest:session:start')
@@ -1971,6 +1975,9 @@ function shouldFinishBailTestSession (globalConfig, results) {
   return !!globalConfig?.bail && getNumBailFailures(results) >= globalConfig.bail
 }
 
+/**
+ * @param {Record<string, unknown> & { onDone?: () => void }} payload
+ */
 async function waitForTestSessionFinish (payload) {
   if (!testSessionFinishCh.hasSubscribers || hasFinishedTestSession) return
 
@@ -1978,8 +1985,9 @@ async function waitForTestSessionFinish (payload) {
 
   let timeoutId
 
+  let onDone
   const flushPromise = new Promise((resolve) => {
-    payload.onDone = () => {
+    onDone = () => {
       clearTimeout(timeoutId)
       resolve()
     }
@@ -1992,7 +2000,7 @@ async function waitForTestSessionFinish (payload) {
     timeoutId.unref?.()
   })
 
-  testSessionFinishCh.publish(payload)
+  publishWithCompletion(testSessionFinishCh, payload, onDone)
 
   const waitingResult = await Promise.race([flushPromise, timeoutPromise])
 
@@ -2157,12 +2165,6 @@ function getWrappedScheduleTests (scheduleTests, frameworkVersion) {
   }
 }
 
-function getChannelPromise (channelToPublishTo, payload = {}) {
-  return new Promise(resolve => {
-    channelToPublishTo.publish({ ...payload, onDone: resolve })
-  })
-}
-
 function searchSourceWrapper (searchSourcePackage, frameworkVersion) {
   const SearchSource = searchSourcePackage.default ?? searchSourcePackage
 
@@ -2517,9 +2519,7 @@ function getCliWrapper (isNewJestVersion) {
 
       if (codeCoverageReportCh.hasSubscribers) {
         const rootDir = result.globalConfig?.rootDir || process.cwd()
-        await new Promise((resolve) => {
-          codeCoverageReportCh.publish({ rootDir, onDone: resolve })
-        })
+        await getChannelPromise(codeCoverageReportCh, { rootDir })
       }
 
       logSessionSummary(ignoredFailuresSummary, getAttemptToFixExecutionsFromJestResults(result))
@@ -2537,6 +2537,10 @@ function shouldWaitForTestSuiteFinish (environment) {
   return isJestWorker && environment.globalConfig?.workerIdleMemoryLimit !== undefined
 }
 
+/**
+ * @param {Record<string, unknown>} payload
+ * @param {boolean} waitForFinish
+ */
 function publishTestSuiteFinish (payload, waitForFinish) {
   if (!testSuiteFinishCh.hasSubscribers) return
 
@@ -2545,12 +2549,9 @@ function publishTestSuiteFinish (payload, waitForFinish) {
     return
   }
 
-  return new Promise(resolve => {
-    testSuiteFinishCh.publish({
-      ...payload,
-      waitForFinish,
-      onDone: resolve,
-    })
+  return getChannelPromise(testSuiteFinishCh, {
+    ...payload,
+    waitForFinish,
   })
 }
 

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

@@ -3,6 +3,11 @@
 const { createCoverageMap } = require('../../../../vendor/dist/istanbul-lib-coverage')
 const satisfies = require('../../../../vendor/dist/semifies')
 const { DD_MAJOR } = require('../../../../version')
+const {
+  getRunStoresPromise,
+  publishWithCompletion,
+  runStoresWithCompletion,
+} = require('../helpers/channel')
 const { addHook, channel } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
 const { isMarkedAsUnskippable } = require('../../../datadog-plugin-jest/src/util')
@@ -240,6 +245,41 @@ function resetSuiteSkippingRunState () {
   writeCoverageBackfillToCache({})
 }
 
+/**
+ * @param {((failures: number) => void) | undefined} callback
+ * @returns {{ onRunDone: (failures: number) => void, onFlushDone: () => void }}
+ */
+function getRunCompletionCallbacks (callback) {
+  let failures
+  let hasRunFinished = false
+  let hasFlushFinished = false
+  let hasCompleted = false
+  const onDone = callback || (() => {})
+
+  const completeIfReady = () => {
+    if (hasCompleted || !hasRunFinished || !hasFlushFinished) return
+
+    hasCompleted = true
+    onDone(failures)
+  }
+
+  return {
+    onRunDone: (runFailures) => {
+      if (hasRunFinished) return
+
+      failures = runFailures
+      hasRunFinished = true
+      completeIfReady()
+    },
+    onFlushDone: () => {
+      if (hasFlushFinished) return
+
+      hasFlushFinished = true
+      completeIfReady()
+    },
+  }
+}
+
 function getOnStartHandler (frameworkVersion) {
   return function () {
     const processArgv = process.argv.slice(2).join(' ')
@@ -251,7 +291,12 @@ function getOnStartHandler (frameworkVersion) {
   }
 }
 
-function getOnEndHandler (isParallel) {
+/**
+ * @param {boolean} isParallel
+ * @param {() => void} onDone
+ * @returns {() => void}
+ */
+function getOnEndHandler (isParallel, onDone) {
   return function () {
     let status = 'pass'
     let error
@@ -337,7 +382,7 @@ function getOnEndHandler (isParallel) {
       global.__coverage__ = fromCoverageMapToCoverage(originalCoverageMap)
     }
 
-    testSessionFinishCh.publish({
+    publishWithCompletion(testSessionFinishCh, {
       status,
       isSuitesSkipped,
       testCodeCoverageLinesTotal,
@@ -350,19 +395,13 @@ function getOnEndHandler (isParallel) {
       isEarlyFlakeDetectionFaulty: config.isEarlyFlakeDetectionFaulty,
       isTestManagementEnabled: config.isTestManagementTestsEnabled,
       isParallel,
-    })
+    }, onDone)
 
     logTestOptimizationSummary({ attemptToFixExecutions, newTestsWithDynamicNames })
     loggedAttemptToFixTests.clear()
   }
 }
 
-function getRunStoresPromise (channelToPublishTo, ctx) {
-  return new Promise(resolve => {
-    channelToPublishTo.runStores({ ...ctx, onDone: resolve }, () => {})
-  })
-}
-
 function applyKnownTestsResponse ({ err, knownTests }) {
   if (err) {
     config.knownTests = []
@@ -395,13 +434,14 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
   let skippableSuitesResponse
   resetSuiteSkippingRunState()
 
-  const onReceivedSkippableSuites = ({
-    err,
-    skippableSuites,
-    itrCorrelationId: responseItrCorrelationId,
-    skippableSuitesCoverage: responseSkippableSuitesCoverage,
-  }) => {
-    if (err) {
+  const onReceivedSkippableSuites = (response) => {
+    const {
+      err,
+      skippableSuites,
+      itrCorrelationId: responseItrCorrelationId,
+      skippableSuitesCoverage: responseSkippableSuitesCoverage,
+    } = response || {}
+    if (!response || err) {
       suitesToSkip = []
       skippableSuitesCoverage = {}
     } else {
@@ -443,12 +483,12 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
       return
     }
 
-    ctx.onDone = onReceivedSkippableSuites
-    skippableSuitesCh.runStores(ctx, () => {})
+    runStoresWithCompletion(skippableSuitesCh, ctx, onReceivedSkippableSuites)
   }
 
-  const onReceivedImpactedTests = ({ err, modifiedFiles: receivedModifiedFiles }) => {
-    if (err) {
+  const onReceivedImpactedTests = (response) => {
+    const { err, modifiedFiles: receivedModifiedFiles } = response || {}
+    if (!response || err) {
       config.modifiedFiles = []
       config.isImpactedTestsEnabled = false
     } else {
@@ -465,8 +505,7 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
 
   const continueAfterTestRequests = () => {
     if (config.isImpactedTestsEnabled) {
-      ctx.onDone = onReceivedImpactedTests
-      modifiedFilesCh.runStores(ctx, () => {})
+      runStoresWithCompletion(modifiedFilesCh, ctx, onReceivedImpactedTests)
     } else if (config.isSuitesSkippingEnabled) {
       requestSkippableSuites()
     } else {
@@ -476,8 +515,14 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     }
   }
 
-  const onReceivedConfiguration = ({ err, isTestDynamicInstrumentationEnabled, libraryConfig, repositoryRoot }) => {
-    if (err || !skippableSuitesCh.hasSubscribers || !knownTestsCh.hasSubscribers) {
+  const onReceivedConfiguration = (response) => {
+    const {
+      err,
+      isTestDynamicInstrumentationEnabled,
+      libraryConfig,
+      repositoryRoot,
+    } = response || {}
+    if (!response || err || !skippableSuitesCh.hasSubscribers || !knownTestsCh.hasSubscribers) {
       return mochaGlobalRunCh.runStores(ctx, () => {
         onFinishRequest()
       })
@@ -526,9 +571,7 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     })
   }
 
-  ctx.onDone = onReceivedConfiguration
-
-  libraryConfigurationCh.runStores(ctx, () => {})
+  runStoresWithCompletion(libraryConfigurationCh, ctx, onReceivedConfiguration)
 }
 
 // In this hook we delay the execution with options.delay to grab library configuration,
@@ -641,6 +684,9 @@ addHook({
       return run.apply(this, args)
     }
 
+    const { onRunDone, onFlushDone } = getRunCompletionCallbacks(args[0])
+    args[0] = onRunDone
+
     const { suitesByTestFile, numSuitesByTestFile } = getSuitesByTestFile(this.suite)
     // Root-level tests (direct children of root, no describe wrapper) keyed by file.
     // Populated during the root 'suite' event so the normal finish path can include them
@@ -752,7 +798,7 @@ addHook({
       finishRootSuiteForFile(test.file)
     }
 
-    const onEnd = getOnEndHandler(false)
+    const onEnd = getOnEndHandler(false, onFlushDone)
 
     this.once('start', getOnStartHandler(frameworkVersion))
 
@@ -1041,8 +1087,11 @@ addHook({
       return run.apply(this, arguments)
     }
 
+    const { onRunDone, onFlushDone } = getRunCompletionCallbacks(cb)
+    arguments[0] = onRunDone
+
     this.once('start', getOnStartHandler(frameworkVersion))
-    this.once('end', getOnEndHandler(true))
+    this.once('end', getOnEndHandler(true, onFlushDone))
 
     // Populate unskippable suites before config is fetched (matches serial mode at Mocha.prototype.run)
     for (const filePath of files) {
@@ -1082,7 +1131,7 @@ addHook({
         skippedSuites = skippedFiles
         skippedSuitesCoverage = getSkippedSuitesCoverageForRun()
         writeCoverageBackfillToCache(skippedSuitesCoverage, getCoverageRootDir())
-        run.apply(this, [cb, { files: filteredFiles }])
+        run.apply(this, [onRunDone, { files: filteredFiles }])
       } else {
         run.apply(this, arguments)
       }

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

@@ -29,6 +29,7 @@ const {
   RUM_TEST_EXECUTION_ID_COOKIE_NAME: RUM_COOKIE_NAME,
 } = require('../../dd-trace/src/ci-visibility/rum')
 const { DD_MAJOR } = require('../../../version')
+const { getChannelPromise } = require('./helpers/channel')
 const { addHook, channel, tracingChannel } = require('./helpers/instrument')
 
 const testStartCh = channel('ci:playwright:test:start')
@@ -614,12 +615,6 @@ function getTestByTestId (dispatcher, testId) {
   }
 }
 
-function getChannelPromise (channelToPublishTo, params) {
-  return new Promise(resolve => {
-    channelToPublishTo.publish({ onDone: resolve, ...params })
-  })
-}
-
 // Inspired by https://github.com/microsoft/playwright/blob/2b77ed4d7aafa85a600caa0b0d101b72c8437eeb/packages/playwright/src/reporters/base.ts#L293
 // We can't use test.outcome() directly because it's set on follow up handlers:
 // our `testEndHandler` is called before the outcome is set.
@@ -1206,8 +1201,6 @@ function dispatcherHookNew (dispatcherExport, runWrapper) {
 function runAllTestsWrapper (runAllTests, playwrightVersion) {
   // Config parameter is only available from >=1.55.0
   return async function (config) {
-    let onDone
-
     rootDir = getRootDir(this, config)
     const projects = getProjectsFromRunner(this, config)
     const isFailureScreenshotEnabled = isFailureScreenshotCaptureEnabled(projects)
@@ -1379,17 +1372,12 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
     logTestOptimizationSummary({ attemptToFixExecutions, newTestsWithDynamicNames })
     loggedAttemptToFixTests.clear()
 
-    const flushWait = new Promise(resolve => {
-      onDone = resolve
-    })
-    testSessionFinishCh.publish({
+    await getChannelPromise(testSessionFinishCh, {
       status: preventedToFail ? 'pass' : STATUS_TO_TEST_STATUS[sessionStatus],
       isEarlyFlakeDetectionEnabled,
       isEarlyFlakeDetectionFaulty,
       isTestManagementTestsEnabled,
-      onDone,
     })
-    await flushWait
 
     startedSuites = []
     remainingTestsByFile = {}
@@ -2112,12 +2100,6 @@ function instrumentWorkerMainMethods (workerMain) {
       annotationTags = parseAnnotations(annotations)
     }
 
-    let onDone
-
-    const flushPromise = new Promise(resolve => {
-      onDone = resolve
-    })
-
     // Wait for the properties to be received, but do not block the worker forever if IPC fails.
     const ddPropertiesTimeoutPromise = new Promise(resolve => {
       const ddPropertiesTimeout = realSetTimeout(() => {
@@ -2143,7 +2125,7 @@ function instrumentWorkerMainMethods (workerMain) {
       testStatus: STATUS_TO_TEST_STATUS[status],
     })
 
-    testFinishCh.publish({
+    await getChannelPromise(testFinishCh, {
       testStatus: STATUS_TO_TEST_STATUS[status],
       steps: steps.filter(step => step.testId === testId),
       error,
@@ -2161,14 +2143,11 @@ function instrumentWorkerMainMethods (workerMain) {
       hasFailedAttemptToFixRetries: test._ddHasFailedAttemptToFixRetries,
       isAtrRetry: test._ddIsAtrRetry,
       isModified: test._ddIsModified,
-      onDone,
       finalStatus,
       earlyFlakeAbortReason: test._ddEarlyFlakeAbortReason,
       ...testCtx.currentStore,
     })
 
-    await flushPromise
-
     return res
   })
 

package/packages/datadog-instrumentations/src/vitest-main-no-worker-init.js

@@ -419,7 +419,7 @@ function createMainProcessReporter (reporterState) {
     testSuiteFinishCh.publish({
       status: getDatadogStatus(testSuiteResult),
       deferFlush: true,
-      onFinish: noop,
+      onDone: noop,
       ...testSuiteCtx.currentStore,
     })
     testSuiteContexts.delete(testModuleId)

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

@@ -16,6 +16,7 @@ const {
   getTestSuitePath,
   isModifiedTest,
 } = require('../../dd-trace/src/plugins/util/test')
+const { getChannelPromise } = require('./helpers/channel')
 const { addHook } = require('./helpers/instrument')
 const noWorkerInit = require('./vitest-main-no-worker-init')
 const {
@@ -38,7 +39,6 @@ const {
   workerReportLogsCh,
   codeCoverageReportCh,
   findExportByName,
-  getChannelPromise,
   getTypeTasks,
   getWorkspaceProject,
   setProvidedContext,
@@ -433,7 +433,8 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
       err,
       libraryConfig,
       requestErrorTags: receivedRequestErrorTags = {},
-    } = await getChannelPromise(libraryConfigurationCh, frameworkVersion, {
+    } = await getChannelPromise(libraryConfigurationCh, {
+      frameworkVersion,
       isVitestNoWorkerInitActive: shouldInstallNoWorkerInit,
     })
     requestErrorTags = receivedRequestErrorTags
@@ -450,7 +451,7 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
   resetMainProcessProvidedContext(ctx)
 
   if (testSessionConfigurationCh.hasSubscribers) {
-    testSessionConfiguration = await getChannelPromise(testSessionConfigurationCh, frameworkVersion)
+    testSessionConfiguration = await getChannelPromise(testSessionConfigurationCh, { frameworkVersion }) || {}
     const {
       testSessionId,
       testModuleId,
@@ -494,7 +495,7 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
 
   if (isKnownTestsEnabled) {
     const currentKnownTestsResponse = knownTestsResponse || await getChannelPromise(knownTestsCh)
-    if (currentKnownTestsResponse.err) {
+    if (!currentKnownTestsResponse || currentKnownTestsResponse.err) {
       isEarlyFlakeDetectionEnabled = false
     } else {
       knownTests = currentKnownTestsResponse.knownTests
@@ -538,12 +539,13 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
   }
 
   if (isTestManagementTestsEnabled) {
-    const { err, testManagementTests: receivedTestManagementTests } =
+    const testManagementResponse =
       testManagementTestsResponse || await getChannelPromise(testManagementTestsCh)
-    if (err) {
+    if (!testManagementResponse || testManagementResponse.err) {
       isTestManagementTestsEnabled = false
       log.error('Could not get test management tests.')
     } else {
+      const { testManagementTests: receivedTestManagementTests } = testManagementResponse
       testManagementTests = receivedTestManagementTests
       testManagementTestsBySuite = getTestManagementTestsBySuite(receivedTestManagementTests)
       shouldSendTestProperties = true
@@ -558,8 +560,7 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
 
   if (isImpactedTestsEnabled) {
     const modifiedFilesResponse = await getChannelPromise(modifiedFilesCh)
-    const { err } = modifiedFilesResponse
-    if (err) {
+    if (!modifiedFilesResponse || modifiedFilesResponse.err) {
       log.error('Could not get modified tests.')
     } else {
       modifiedFiles = modifiedFilesResponse.modifiedFiles
@@ -812,18 +813,13 @@ function getFinishWrapper (exitOrClose) {
       return exitOrClose.apply(this, arguments)
     }
 
-    let onFinish
-
-    const flushPromise = new Promise(resolve => {
-      onFinish = resolve
-    })
     const failedSuites = this.state.getFailedFilepaths()
     let error
     if (failedSuites.length) {
       error = new Error(`Test suites failed: ${failedSuites.length}.`)
     }
 
-    testSessionFinishCh.publish({
+    const flushPromise = getChannelPromise(testSessionFinishCh, {
       status: getSessionStatus(this.state),
       testCodeCoverageLinesTotal,
       error,
@@ -833,7 +829,6 @@ function getFinishWrapper (exitOrClose) {
       requestErrorTags,
       vitestPool,
       isVitestNoWorkerInitActive,
-      onFinish,
     })
 
     logTestOptimizationSummary({ attemptToFixExecutions, newTestsWithDynamicNames })
@@ -842,9 +837,7 @@ function getFinishWrapper (exitOrClose) {
 
     // If coverage was generated, publish coverage report channel for upload
     if (coverageRootDir && codeCoverageReportCh.hasSubscribers) {
-      await new Promise((resolve) => {
-        codeCoverageReportCh.publish({ rootDir: coverageRootDir, onDone: resolve })
-      })
+      await getChannelPromise(codeCoverageReportCh, { rootDir: coverageRootDir })
     }
 
     return exitOrClose.apply(this, arguments)
@@ -1213,16 +1206,11 @@ async function reportTypecheckFile (file, sessionConfiguration, frameworkVersion
     testSuiteErrorCh.runStores(testSuiteCtx, () => {})
   }
 
-  let onFinish
-  const onFinishPromise = new Promise(resolve => {
-    onFinish = resolve
-  })
-  testSuiteFinishCh.publish({
+  await getChannelPromise(testSuiteFinishCh, {
+    frameworkVersion,
     status: getTypecheckTaskStatus(file),
-    onFinish,
     ...testSuiteCtx.currentStore,
   })
-  await onFinishPromise
 }
 
 async function reportTypecheckResults (result, frameworkVersion, ctx, typechecker) {
@@ -1234,7 +1222,7 @@ async function reportTypecheckResults (result, frameworkVersion, ctx, typechecke
   }
   const providedContext = getMainProcessProvidedContext(ctx)
   const sessionConfiguration = testSessionConfigurationCh.hasSubscribers
-    ? await getChannelPromise(testSessionConfigurationCh, frameworkVersion)
+    ? await getChannelPromise(testSessionConfigurationCh, { frameworkVersion }) || {}
     : {}
 
   await Promise.all(result.files.map(file => reportTypecheckFile(

package/packages/datadog-instrumentations/src/vitest-util.js

@@ -39,12 +39,6 @@ function findExportByName (pkg, name) {
   }
 }
 
-function getChannelPromise (channelToPublishTo, frameworkVersion, payload) {
-  return new Promise(resolve => {
-    channelToPublishTo.publish({ ...payload, onDone: resolve, frameworkVersion })
-  })
-}
-
 function getTestRunnerExport (testPackage) {
   return findExportByName(testPackage, 'VitestTestRunner') || findExportByName(testPackage, 'TestRunner')
 }
@@ -239,7 +233,6 @@ module.exports = {
   workerReportLogsCh,
   codeCoverageReportCh,
   findExportByName,
-  getChannelPromise,
   getTestRunnerExport,
   getTypeTasks,
   getTestName,

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

@@ -9,6 +9,7 @@ const {
   recordAttemptToFixExecution,
   logAttemptToFixTestExecution,
 } = require('../../dd-trace/src/plugins/util/test')
+const { getChannelPromise } = require('./helpers/channel')
 const { addHook } = require('./helpers/instrument')
 const {
   testStartCh,
@@ -661,11 +662,6 @@ addHook({
     testSuiteStartCh.runStores(testSuiteCtx, () => {})
     const startTestsResponse = await startTests.apply(this, arguments)
 
-    let onFinish = null
-    const onFinishPromise = new Promise(resolve => {
-      onFinish = resolve
-    })
-
     const testTasks = getTypeTasks(startTestsResponse[0].tasks)
     const testEventPromises = []
 
@@ -811,9 +807,10 @@ addHook({
       testSuiteErrorCh.runStores(testSuiteCtx, () => {})
     }
 
-    testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish, ...testSuiteCtx.currentStore })
-
-    await onFinishPromise
+    await getChannelPromise(testSuiteFinishCh, {
+      status: testSuiteResult.state,
+      ...testSuiteCtx.currentStore,
+    })
 
     return startTestsResponse
   })

package/packages/datadog-plugin-cucumber/src/index.js

@@ -72,6 +72,7 @@ class CucumberPlugin extends CiPlugin {
       isEarlyFlakeDetectionFaulty,
       isTestManagementTestsEnabled,
       isParallel,
+      onDone,
     }) => {
       this._exportPendingWorkerTraces()
       const {
@@ -127,7 +128,7 @@ class CucumberPlugin extends CiPlugin {
       })
 
       this.libraryConfig = null
-      this.tracer._exporter.flush()
+      this.tracer._exporter.flush(onDone)
     })
 
     this.addSub('ci:cucumber:test-suite:start', ({

package/packages/datadog-plugin-mocha/src/index.js

@@ -390,6 +390,7 @@ class MochaPlugin extends CiPlugin {
       isEarlyFlakeDetectionFaulty,
       isTestManagementEnabled,
       isParallel,
+      onDone,
     }) => {
       this._exportPendingWorkerTraces()
       if (this.testSessionSpan) {
@@ -456,7 +457,7 @@ class MochaPlugin extends CiPlugin {
         })
       }
       this.libraryConfig = null
-      this.tracer._exporter.flush()
+      this.tracer._exporter.flush(onDone)
     })
 
     this.addBind('ci:mocha:global:run', (ctx) => {

package/packages/datadog-plugin-playwright/src/index.js

@@ -388,7 +388,10 @@ class PlaywrightPlugin extends CiPlugin {
       earlyFlakeAbortReason,
       onDone,
     }) => {
-      if (!span) return
+      if (!span) {
+        onDone?.()
+        return
+      }
 
       const isRUMActive = span.context().getTag(TEST_IS_RUM_ACTIVE)
 

package/packages/datadog-plugin-vitest/src/index.js

@@ -400,7 +400,7 @@ class VitestPlugin extends CiPlugin {
       return ctx.currentStore
     })
 
-    this.addSub('ci:vitest:test-suite:finish', ({ testSuiteSpan, status, deferFlush, onFinish }) => {
+    this.addSub('ci:vitest:test-suite:finish', ({ testSuiteSpan, status, deferFlush, onDone }) => {
       if (testSuiteSpan) {
         testSuiteSpan.setTag(TEST_STATUS, status)
         testSuiteSpan.finish()
@@ -408,10 +408,10 @@ class VitestPlugin extends CiPlugin {
       }
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
       if (deferFlush) {
-        onFinish()
+        onDone()
         return
       }
-      this.tracer._exporter.flush(onFinish)
+      this.tracer._exporter.flush(onDone)
       if (this.runningTestProbe) {
         this.removeDiProbe(this.runningTestProbe)
       }
@@ -442,7 +442,7 @@ class VitestPlugin extends CiPlugin {
       requestErrorTags,
       vitestPool,
       isVitestNoWorkerInitActive,
-      onFinish,
+      onDone,
     }) => {
       for (const [tag, value] of Object.entries(requestErrorTags || {})) {
         this.testSessionSpan.setTag(tag, value)
@@ -481,7 +481,7 @@ class VitestPlugin extends CiPlugin {
         provider: this.ciProviderName,
         autoInjected: !!this._tracerConfig.testOptimization.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
       })
-      this.tracer._exporter.flush(onFinish)
+      this.tracer._exporter.flush(onDone)
     })
 
     this.addSub('ci:vitest:coverage-report', ({ rootDir, onDone }) => {

package/packages/dd-trace/src/profiling/profilers/space.js

@@ -59,7 +59,7 @@ class NativeSpaceProfiler {
   }
 
   profile (restart) {
-    const profile = this.#pprof.heap.profile(undefined, this.#mapper, getThreadLabels)
+    const profile = this.#pprof.heap.profile(undefined, this.#mapper, getThreadLabels, 'pack')
     if (!restart) {
       this.stop()
     }

package/packages/dd-trace/src/profiling/profilers/wall.js

@@ -187,6 +187,7 @@ class NativeWallProfiler {
 
     this.#pprof.time.start({
       collectCpuTime: this.#cpuProfilingEnabled,
+      columnNumbers: 'pack',
       durationMillis: this.#flushIntervalMillis,
       intervalMicros: this.#samplingIntervalMicros,
       lineNumbers: false,