new exporter

Author: juan-fernandez

ci/init.js

@@ -8,13 +8,18 @@ const isCucumberWorker = !!process.env.CUCUMBER_WORKER_ID
 const isMochaWorker = !!process.env.MOCHA_WORKER_ID
 
 const isPlaywrightWorker = !!process.env.DD_PLAYWRIGHT_WORKER
+// we can't use VITEST_WORKER_ID because it's set _after_ the worker is initialized
+// maybe we can intercept tinypool to inject it to be sure
+const isVitestWorker = !!process.env.TINYPOOL_WORKER_ID
 
 const packageManagers = [
   'npm',
   'yarn',
   'pnpm'
 ]
 
+console.log('init', isVitestWorker)
+
 const isPackageManager = () => {
   return packageManagers.some(packageManager => process.argv[1]?.includes(`bin/${packageManager}`))
 }
@@ -75,6 +80,13 @@ if (isPlaywrightWorker) {
   }
 }
 
+if (isVitestWorker) {
+  console.log('initializing as vites tworker')
+  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'
 }

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

@@ -132,5 +132,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,7 +1,8 @@
+const v8 = require('node:v8')
+
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 const log = require('../../dd-trace/src/log')
-
 // test hooks
 const testStartCh = channel('ci:vitest:test:start')
 const testFinishTimeCh = channel('ci:vitest:test:finish-time')
@@ -26,6 +27,8 @@ const knownTestsCh = channel('ci:vitest:known-tests')
 const isEarlyFlakeDetectionFaultyCh = channel('ci:vitest:is-early-flake-detection-faulty')
 const testManagementTestsCh = channel('ci:vitest:test-management-tests')
 
+const workerReporterCh = channel('ci:vitest:worker-report:trace')
+
 const taskToCtx = new WeakMap()
 const taskToStatuses = new WeakMap()
 const newTasks = new WeakSet()
@@ -61,6 +64,8 @@ function getProvidedContext () {
       _ddIsFlakyTestRetriesEnabled: isFlakyTestRetriesEnabled
     } = globalThis.__vitest_worker__.providedContext
 
+    // console.log('globalThis.__vitest_worker__', globalThis.__vitest_worker__)
+
     return {
       isDiEnabled: _ddIsDiEnabled,
       isEarlyFlakeDetectionEnabled: _ddIsEarlyFlakeDetectionEnabled,
@@ -187,6 +192,11 @@ function getSortWrapper (sort) {
     let knownTests = {}
     let testManagementTests = {}
 
+    // 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)
       if (!err) {
@@ -262,24 +272,25 @@ function getSortWrapper (sort) {
         log.warn('Could not send Dynamic Instrumentation configuration to workers.')
       }
     }
-
-    if (isTestManagementTestsEnabled) {
-      const { err, testManagementTests: receivedTestManagementTests } = await getChannelPromise(testManagementTestsCh)
-      if (!err) {
-        testManagementTests = receivedTestManagementTests
-        try {
-          const workspaceProject = this.ctx.getCoreWorkspaceProject()
-          workspaceProject._provided._ddIsTestManagementTestsEnabled = isTestManagementTestsEnabled
-          workspaceProject._provided._ddTestManagementAttemptToFixRetries = testManagementAttemptToFixRetries
-          workspaceProject._provided._ddTestManagementTests = testManagementTests
-        } catch {
-          log.warn('Could not send test management tests to workers so Test Management will not work.')
-        }
-      } else {
-        isTestManagementTestsEnabled = false
-        log.error('Could not get test management tests.')
-      }
-    }
+    debugger
+
+    // if (isTestManagementTestsEnabled) {
+    //   const { err, testManagementTests: receivedTestManagementTests } = await getChannelPromise(testManagementTestsCh)
+    //   if (!err) {
+    //     testManagementTests = receivedTestManagementTests
+    //     try {
+    //       const workspaceProject = this.ctx.getCoreWorkspaceProject()
+    //       workspaceProject._provided._ddIsTestManagementTestsEnabled = isTestManagementTestsEnabled
+    //       workspaceProject._provided._ddTestManagementAttemptToFixRetries = testManagementAttemptToFixRetries
+    //       workspaceProject._provided._ddTestManagementTests = testManagementTests
+    //     } catch {
+    //       log.warn('Could not send test management tests to workers so Test Management will not work.')
+    //     }
+    //   } else {
+    //     isTestManagementTestsEnabled = false
+    //     log.error('Could not get test management tests.')
+    //   }
+    // }
 
     let testCodeCoverageLinesTotal
 
@@ -344,6 +355,47 @@ 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
+})
+
+addHook({
+  name: 'tinypool',
+  versions: ['>=1.0.0'],
+  file: 'dist/index.js'
+}, (TinyPool) => {
+  debugger
+  // we can pass handle here to the worker, and then use it to send messages to the main process
+  shimmer.wrap(TinyPool.prototype, 'run', run => function (_, { channel }) {
+    const res = run.apply(this, arguments)
+
+    this.threads.forEach(thread => {
+      thread.process.on('message', (message) => {
+        if (message.__tinypool_worker_message__ && message.data) {
+          workerReporterCh.publish(message.data)
+        }
+      })
+    })
+
+    return res
+  })
+
+  return TinyPool
+})
+
 addHook({
   name: 'vitest',
   versions: ['>=1.6.0'],
@@ -354,6 +406,7 @@ 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 {
@@ -420,6 +473,7 @@ 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) {
@@ -819,6 +873,7 @@ addHook({
 
     testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish, ...testSuiteCtx.currentStore })
 
+    // console.log('test usite finish!')
     // TODO: fix too frequent flushes
     await onFinishPromise
 

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

@@ -35,6 +35,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
@@ -386,6 +387,22 @@ class VitestPlugin extends CiPlugin {
       })
       this.tracer._exporter.flush(onFinish)
     })
+
+    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,
+          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)
+      })
+    })
   }
 
   getTestProperties (testManagementTests, testSuite, testName) {

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

@@ -7,7 +7,8 @@ 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
 } = require('../../../plugins/util/test')
 
 function getInterprocessTraceCode () {
@@ -23,6 +24,9 @@ function getInterprocessTraceCode () {
   if (process.env.DD_PLAYWRIGHT_WORKER) {
     return PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE
   }
+  if (process.env.VITEST_WORKER_ID) {
+    return VITEST_WORKER_TRACE_PAYLOAD_CODE
+  }
   return null
 }
 

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

@@ -34,10 +34,19 @@ 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], () => {
-        onDone()
-      })
+      if (process.env.TINYPOOL_WORKER_ID) {
+        // 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()
+        })
+      }
     }
+    // onDone()
   }
 }
 

packages/dd-trace/src/exporter.js

@@ -21,6 +21,7 @@ module.exports = 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:
       return inAWSLambda && !usingLambdaExtension ? require('./exporters/log') : require('./exporters/agent')

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

@@ -69,6 +69,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') },