clean up and add test

Author: juan-fernandez

integration-tests/vitest/vitest.spec.js

@@ -2054,5 +2054,13 @@ versions.forEach((version) => {
         })
       })
     })
+
+    it.skip('does not blow up when tinypool is used outside of a test', (done) => {
+      // childProcess = exec('node ./ci-visibility/run-workerpool.js', {
+      //   cwd,
+      //   env: getCiVisAgentlessConfig(receiver.port),
+      //   stdio: 'pipe'
+      // })
+    })
   })
 })

packages/datadog-instrumentations/src/vitest.js

@@ -1,6 +1,13 @@
+'use strict'
+
 const { addHook, channel } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 const log = require('../../dd-trace/src/log')
+const {
+  VITEST_WORKER_TRACE_PAYLOAD_CODE,
+  VITEST_WORKER_LOGS_PAYLOAD_CODE
+} = require('../../dd-trace/src/plugins/util/test')
+
 // test hooks
 const testStartCh = channel('ci:vitest:test:start')
 const testFinishTimeCh = channel('ci:vitest:test:finish-time')
@@ -27,7 +34,8 @@ const isEarlyFlakeDetectionFaultyCh = channel('ci:vitest:is-early-flake-detectio
 const testManagementTestsCh = channel('ci:vitest:test-management-tests')
 const impactedTestsCh = channel('ci:vitest:modified-tests')
 
-const workerReporterCh = channel('ci:vitest:worker-report:trace')
+const workerReportTraceCh = channel('ci:vitest:worker-report:trace')
+const workerReportLogsCh = channel('ci:vitest:worker-report:logs')
 
 const taskToCtx = new WeakMap()
 const taskToStatuses = new WeakMap()
@@ -39,6 +47,7 @@ const modifiedTasks = new WeakSet()
 let isRetryReasonEfd = false
 let isRetryReasonAttemptToFix = false
 const switchedStatuses = new WeakSet()
+const workerProcesses = new WeakSet()
 
 const BREAKPOINT_HIT_GRACE_PERIOD_MS = 400
 
@@ -66,8 +75,6 @@ function getProvidedContext () {
       _ddModifiedTests: modifiedTests
     } = globalThis.__vitest_worker__.providedContext
 
-    // console.log('globalThis.__vitest_worker__', globalThis.__vitest_worker__)
-
     return {
       isDiEnabled: _ddIsDiEnabled,
       isEarlyFlakeDetectionEnabled: _ddIsEarlyFlakeDetectionEnabled,
@@ -195,11 +202,6 @@ function getSortWrapper (sort, frameworkVersion) {
     let testManagementAttemptToFixRetries = 0
     let isDiEnabled = false
 
-    // console.log('this.ctx.vitest', this.ctx.vitest)
-    // console.log('this.ctx.rpc', this.ctx.rpc)
-    console.log('__vitest_worker__', global.__vitest_worker__)
-    // console.log('this.ctx.getCoreWorkspaceProject()', this.ctx.getCoreWorkspaceProject())
-    // console.log('this.ctx', this)
     try {
       const { err, libraryConfig } = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
       if (!err) {
@@ -275,7 +277,6 @@ function getSortWrapper (sort, frameworkVersion) {
         log.warn('Could not send Dynamic Instrumentation configuration to workers.')
       }
     }
-    debugger
 
     if (isTestManagementTestsEnabled) {
       const { err, testManagementTests: receivedTestManagementTests } = await getChannelPromise(testManagementTestsCh)
@@ -369,32 +370,18 @@ function getCreateCliWrapper (vitestPackage, frameworkVersion) {
   return vitestPackage
 }
 
-// UNUSED RIGHT NOW, but we can use it to bind the async resource to the test fn
-// getFn is what's used to get the test fn to run it with vitest:
-// https://github.com/vitest-dev/vitest/blob/0cbad1b0d0d56f1ec60f8496678d1435f8bb8977/packages/runner/src/run.ts#L315-L321
-let getFn = null
-
-// run in workers only
-addHook({
-  name: '@vitest/runner',
-  versions: ['>=1.6.0'],
-  file: 'dist/index.js'
-}, (suitePackage) => {
-  getFn = suitePackage.getFn
-
-  return suitePackage
-})
-
-const processToHandler = new WeakSet()
-
 function threadHandler (thread) {
-  if (processToHandler.has(thread.process)) {
+  if (workerProcesses.has(thread.process)) {
     return
   }
-  processToHandler.add(thread.process)
+  workerProcesses.add(thread.process)
   thread.process.on('message', (message) => {
     if (message.__tinypool_worker_message__ && message.data) {
-      workerReporterCh.publish(message.data)
+      if (message.interprocessCode === VITEST_WORKER_TRACE_PAYLOAD_CODE) {
+        workerReportTraceCh.publish(message.data)
+      } else if (message.interprocessCode === VITEST_WORKER_LOGS_PAYLOAD_CODE) {
+        workerReportLogsCh.publish(message.data)
+      }
     }
   })
 }
@@ -405,12 +392,11 @@ addHook({
   file: 'dist/index.js'
 }, (TinyPool) => {
   shimmer.wrap(TinyPool.prototype, 'run', run => async function () {
-    // we need to do this before and after because the threads list gets recycled
-    // (the processes are re-created)
+    // We have to do this before and after because the threads list gets recycled, that is, the processes are re-created
     this.threads.forEach(threadHandler)
-    const res = await run.apply(this, arguments)
+    const runResult = await run.apply(this, arguments)
     this.threads.forEach(threadHandler)
-    return res
+    return runResult
   })
 
   return TinyPool
@@ -426,7 +412,6 @@ addHook({
   // `onBeforeRunTask` is run before any repetition or attempt is run
   // `onBeforeRunTask` is an async function
   shimmer.wrap(VitestTestRunner.prototype, 'onBeforeRunTask', onBeforeRunTask => function (task) {
-    // console.log('on before run task', process.env)
     const testName = getTestName(task)
 
     const {
@@ -512,7 +497,6 @@ addHook({
   // `onAfterRunTask` is run after all repetitions or attempts are run
   // `onAfterRunTask` is an async function
   shimmer.wrap(VitestTestRunner.prototype, 'onAfterRunTask', onAfterRunTask => function (task) {
-    // console.log('task', task)
     const { isEarlyFlakeDetectionEnabled, isTestManagementTestsEnabled } = getProvidedContext()
 
     if (isTestManagementTestsEnabled) {
@@ -913,8 +897,6 @@ addHook({
 
     testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish, ...testSuiteCtx.currentStore })
 
-    // console.log('test usite finish!')
-    // TODO: fix too frequent flushes
     await onFinishPromise
 
     return startTestsResponse

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

@@ -404,7 +404,6 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addSub('ci:vitest:worker-report:trace', (traces) => {
-      // it has no test session or test module id so there are hanging
       const formattedTraces = JSON.parse(traces).map(trace => {
         return trace.map(span => ({
           ...span,
@@ -418,6 +417,12 @@ class VitestPlugin extends CiPlugin {
         this.tracer._exporter.export(trace)
       })
     })
+
+    this.addSub('ci:vitest:worker-report:logs', (logsPayloads) => {
+      JSON.parse(logsPayloads).forEach(({ testConfiguration, logMessage }) => {
+        this.tracer._exporter.exportDiLogs(testConfiguration, logMessage)
+      })
+    })
   }
 
   getTestProperties (testManagementTests, testSuite, testName) {

packages/dd-trace/src/ci-visibility/exporters/test-worker/index.js

@@ -7,7 +7,9 @@ const {
   CUCUMBER_WORKER_TRACE_PAYLOAD_CODE,
   MOCHA_WORKER_TRACE_PAYLOAD_CODE,
   JEST_WORKER_LOGS_PAYLOAD_CODE,
-  PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE
+  PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE,
+  VITEST_WORKER_TRACE_PAYLOAD_CODE,
+  VITEST_WORKER_LOGS_PAYLOAD_CODE
 } = require('../../../plugins/util/test')
 const { getEnvironmentVariable } = require('../../../config-helper')
 
@@ -24,6 +26,9 @@ function getInterprocessTraceCode () {
   if (getEnvironmentVariable('DD_PLAYWRIGHT_WORKER')) {
     return PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE
   }
+  if (getEnvironmentVariable('TINYPOOL_WORKER_ID')) {
+    return VITEST_WORKER_TRACE_PAYLOAD_CODE
+  }
   return null
 }
 
@@ -39,6 +44,9 @@ function getInterprocessLogsCode () {
   if (getEnvironmentVariable('JEST_WORKER_ID')) {
     return JEST_WORKER_LOGS_PAYLOAD_CODE
   }
+  if (getEnvironmentVariable('TINYPOOL_WORKER_ID')) {
+    return VITEST_WORKER_LOGS_PAYLOAD_CODE
+  }
   return null
 }
 

packages/dd-trace/src/ci-visibility/exporters/test-worker/writer.js

@@ -2,8 +2,6 @@
 const { JSONEncoder } = require('../../encode/json-encoder')
 const { getEnvironmentVariable } = require('../../../config-helper')
 
-const isVitestWorker = !!getEnvironmentVariable('TINYPOOL_WORKER_ID')
-
 class Writer {
   constructor (interprocessCode) {
     this._encoder = new JSONEncoder()
@@ -37,17 +35,15 @@ class Writer {
     // See cucumber code:
     // https://github.com/cucumber/cucumber-js/blob/5ce371870b677fe3d1a14915dc535688946f734c/src/runtime/parallel/run_worker.ts#L13
     if (process.send) { // it only works if process.send is available
-      if (isVitestWorker) {
-        // in vitest we have to trick the main process into thinking these are messages from
-        // tinypool so they are not rejected
-        process.send({ __tinypool_worker_message__: true, data }, () => {
-          onDone()
-        })
-      } else {
-        process.send([this._interprocessCode, data], () => {
-          onDone()
-        })
-      }
+      const isVitestWorker = !!getEnvironmentVariable('TINYPOOL_WORKER_ID')
+
+      const payload = isVitestWorker
+        ? { __tinypool_worker_message__: true, interprocessCode: this._interprocessCode, data }
+        : [this._interprocessCode, data]
+
+      process.send(payload, () => {
+        onDone()
+      })
     } else {
       onDone()
     }

packages/dd-trace/src/plugins/util/test.js

@@ -116,6 +116,10 @@ const MOCHA_WORKER_TRACE_PAYLOAD_CODE = 80
 // playwright worker variables
 const PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE = 90
 
+// vitest worker variables
+const VITEST_WORKER_TRACE_PAYLOAD_CODE = 100
+const VITEST_WORKER_LOGS_PAYLOAD_CODE = 102
+
 // Early flake detection util strings
 const EFD_STRING = "Retried by Datadog's Early Flake Detection"
 const EFD_TEST_NAME_REGEX = new RegExp(EFD_STRING + String.raw` \(#\d+\): `, 'g')
@@ -211,6 +215,8 @@ module.exports = {
   CUCUMBER_WORKER_TRACE_PAYLOAD_CODE,
   MOCHA_WORKER_TRACE_PAYLOAD_CODE,
   PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE,
+  VITEST_WORKER_TRACE_PAYLOAD_CODE,
+  VITEST_WORKER_LOGS_PAYLOAD_CODE,
   TEST_SOURCE_START,
   TEST_SKIPPED_BY_ITR,
   TEST_IS_NEW,