[test optimization] Improve vitest performance

ci/init.js

@@ -8,6 +8,8 @@ const isJestWorker = !!getEnvironmentVariable('JEST_WORKER_ID')
 const isCucumberWorker = !!getEnvironmentVariable('CUCUMBER_WORKER_ID')
 const isMochaWorker = !!getEnvironmentVariable('MOCHA_WORKER_ID')
 
+// We can't use VITEST_WORKER_ID because it's set _after_ the worker is initialized
+const isVitestWorker = !!getEnvironmentVariable('TINYPOOL_WORKER_ID')
 const isPlaywrightWorker = !!getEnvironmentVariable('DD_PLAYWRIGHT_WORKER')
 
 const packageManagers = [
@@ -76,6 +78,12 @@ if (isPlaywrightWorker) {
   }
 }
 
+if (isVitestWorker) {
+  options.experimental = {
+    exporter: 'vitest_worker'
+  }
+}
+
 if (shouldInit) {
   tracer.init(options)
   tracer.use('fs', false)

ext/exporters.d.ts

@@ -6,7 +6,8 @@ declare const exporters: {
   JEST_WORKER: 'jest_worker',
   CUCUMBER_WORKER: 'cucumber_worker',
   MOCHA_WORKER: 'mocha_worker',
-  PLAYWRIGHT_WORKER: 'playwright_worker'
+  PLAYWRIGHT_WORKER: 'playwright_worker',
+  VITEST_WORKER: 'vitest_worker'
 }
 
 export = exporters

ext/exporters.js

@@ -7,5 +7,6 @@ module.exports = {
   JEST_WORKER: 'jest_worker',
   CUCUMBER_WORKER: 'cucumber_worker',
   MOCHA_WORKER: 'mocha_worker',
-  PLAYWRIGHT_WORKER: 'playwright_worker'
+  PLAYWRIGHT_WORKER: 'playwright_worker',
+  VITEST_WORKER: 'vitest_worker'
 }

integration-tests/ci-visibility/run-tinypool.mjs

@@ -0,0 +1,13 @@
+import Tinypool from 'tinypool'
+
+const pool = new Tinypool({
+  filename: new URL('./tinypool-worker.mjs', import.meta.url).href,
+})
+
+const result = await pool.run({ a: 4, b: 6 })
+// eslint-disable-next-line no-console
+console.log('result', result)
+
+// Make sure to destroy pool once it's not needed anymore
+// This terminates all pool's idle workers
+await pool.destroy()

integration-tests/ci-visibility/tinypool-worker.mjs

@@ -0,0 +1,3 @@
+export default ({ a, b }) => {
+  return a + b
+}

integration-tests/vitest/vitest.spec.js

@@ -67,7 +67,8 @@ versions.forEach((version) => {
       sandbox = await createSandbox([
         `vitest@${version}`,
         `@vitest/coverage-istanbul@${version}`,
-        `@vitest/coverage-v8@${version}`
+        `@vitest/coverage-v8@${version}`,
+        'tinypool'
       ], true)
       cwd = sandbox.folder
     })
@@ -2054,5 +2055,24 @@ versions.forEach((version) => {
         })
       })
     })
+
+    it('does not blow up when tinypool is used outside of a test', (done) => {
+      childProcess = exec('node ./ci-visibility/run-tinypool.mjs', {
+        cwd,
+        env: getCiVisAgentlessConfig(receiver.port),
+        stdio: 'pipe'
+      })
+      childProcess.stdout.on('data', (chunk) => {
+        testOutput += chunk.toString()
+      })
+      childProcess.stderr.on('data', (chunk) => {
+        testOutput += chunk.toString()
+      })
+      childProcess.on('exit', (code) => {
+        assert.include(testOutput, 'result 10')
+        assert.equal(code, 0)
+        done()
+      })
+    })
   })
 })

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -133,5 +133,6 @@ module.exports = {
   vm: () => require('../vm'),
   when: () => require('../when'),
   winston: () => require('../winston'),
-  workerpool: () => require('../mocha')
+  workerpool: () => require('../mocha'),
+  tinypool: { esmFirst: true, fn: () => require('../vitest') }
 }

packages/datadog-instrumentations/src/vitest.js

@@ -1,6 +1,12 @@
+'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')
@@ -28,6 +34,9 @@ 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 workerReportTraceCh = channel('ci:vitest:worker-report:trace')
+const workerReportLogsCh = channel('ci:vitest:worker-report:logs')
+
 const taskToCtx = new WeakMap()
 const taskToStatuses = new WeakMap()
 const newTasks = new WeakSet()
@@ -38,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
 
@@ -360,6 +370,38 @@ function getCreateCliWrapper (vitestPackage, frameworkVersion) {
   return vitestPackage
 }
 
+function threadHandler (thread) {
+  if (workerProcesses.has(thread.process)) {
+    return
+  }
+  workerProcesses.add(thread.process)
+  thread.process.on('message', (message) => {
+    if (message.__tinypool_worker_message__ && 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)
+      }
+    }
+  })
+}
+
+addHook({
+  name: 'tinypool',
+  versions: ['>=1.0.0'],
+  file: 'dist/index.js'
+}, (TinyPool) => {
+  shimmer.wrap(TinyPool.prototype, 'run', run => async function () {
+    // 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 runResult = await run.apply(this, arguments)
+    this.threads.forEach(threadHandler)
+    return runResult
+  })
+
+  return TinyPool
+})
+
 addHook({
   name: 'vitest',
   versions: ['>=1.6.0'],
@@ -855,7 +897,6 @@ addHook({
 
     testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish, ...testSuiteCtx.currentStore })
 
-    // TODO: fix too frequent flushes
     await onFinishPromise
 
     return startTestsResponse

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

@@ -38,6 +38,7 @@ const {
   TELEMETRY_TEST_SESSION
 } = require('../../dd-trace/src/ci-visibility/telemetry')
 const { DD_MAJOR } = require('../../../version')
+const id = require('../../dd-trace/src/id')
 
 // Milliseconds that we subtract from the error test duration
 // so that they do not overlap with the following test
@@ -401,6 +402,27 @@ class VitestPlugin extends CiPlugin {
       })
       this.tracer._exporter.flush(onFinish)
     })
+
+    this.addSub('ci:vitest:worker-report:trace', (traces) => {
+      const formattedTraces = JSON.parse(traces).map(trace => {
+        return trace.map(span => ({
+          ...span,
+          span_id: id(span.span_id),
+          trace_id: id(span.trace_id),
+          parent_id: id(span.parent_id)
+        }))
+      })
+
+      formattedTraces.forEach(trace => {
+        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

@@ -1,5 +1,6 @@
 'use strict'
 const { JSONEncoder } = require('../../encode/json-encoder')
+const { getEnvironmentVariable } = require('../../../config-helper')
 
 class Writer {
   constructor (interprocessCode) {
@@ -34,9 +35,17 @@ 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
-      process.send([this._interprocessCode, data], () => {
+      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/exporter.js

@@ -19,6 +19,7 @@ module.exports = function getExporter (name) {
     case exporters.CUCUMBER_WORKER:
     case exporters.MOCHA_WORKER:
     case exporters.PLAYWRIGHT_WORKER:
+    case exporters.VITEST_WORKER:
       return require('./ci-visibility/exporters/test-worker')
     default: {
       const inAWSLambda = getEnvironmentVariable('AWS_LAMBDA_FUNCTION_NAME') !== undefined

packages/dd-trace/src/plugins/index.js

@@ -70,6 +70,7 @@ module.exports = {
   get 'mocha-each' () { return require('../../../datadog-plugin-mocha/src') },
   get vitest () { return require('../../../datadog-plugin-vitest/src') },
   get workerpool () { return require('../../../datadog-plugin-mocha/src') },
+  get tinypool () { return require('../../../datadog-plugin-vitest/src') },
   get moleculer () { return require('../../../datadog-plugin-moleculer/src') },
   get mongodb () { return require('../../../datadog-plugin-mongodb-core/src') },
   get 'mongodb-core' () { return require('../../../datadog-plugin-mongodb-core/src') },

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,