better vitest support

packages/datadog-instrumentations/src/vitest.js

@@ -111,18 +111,18 @@ addHook({
     })
 
     // TODO: probably no need to flush so often - maybe we can use the test suite finish
-    shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterRunFiles', onAfterRunFiles => async function () {
-      let onFinish = null
-      const onFinishPromise = new Promise(resolve => {
-        onFinish = resolve
-      })
+    // shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterRunFiles', onAfterRunFiles => async function () {
+    //   let onFinish = null
+    //   const onFinishPromise = new Promise(resolve => {
+    //     onFinish = resolve
+    //   })
 
-      onAfterRunFilesCh.publish(onFinish)
+    //   onAfterRunFilesCh.publish(onFinish)
 
-      await onFinishPromise
+    //   await onFinishPromise
 
-      return await onAfterRunFiles.apply(this, arguments)
-    })
+    //   return await onAfterRunFiles.apply(this, arguments)
+    // })
   }
   return vitestPackage
 })
@@ -144,9 +144,20 @@ addHook({
         testSuiteStartCh.publish(testPath[0])
       })
       const startTestsResponse = await startTests.apply(this, arguments)
+
+      let onFinish = null
+      const onFinishPromise = new Promise(resolve => {
+        onFinish = resolve
+      })
+
       asyncResource.runInAsyncScope(() => {
-        testSuiteFinishCh.publish(startTestsResponse[0].result.state)
+        testSuiteFinishCh.publish({ status: startTestsResponse[0].result.state, onFinish })
       })
+
+      console.log('waiting flush')
+      await onFinishPromise
+      console.log('flushed')
+
       return startTestsResponse
     })
 
@@ -164,11 +175,80 @@ addHook({
   return vitestPackage
 })
 
-addHook({
-  name: '@vitest',
-  versions: ['>=0.0.0']
-}, vitestPackage => {
-  debugger
-  console.log('@vitest', { isMainThread })
-  return vitestPackage
-})
+let privateSet
+
+const isPrivateTools = (tinypoolPackage) => {
+  return tinypoolPackage.__privateGet
+}
+
+const isTinyPoolClass = (tinypoolPackage) => {
+  return tinypoolPackage.name === 'Tinypool'
+}
+
+const isThreadPool = (value) => {
+  return value.constructor.name === 'ThreadPool'
+}
+
+let threadPool
+
+// MAYBE NOT NEEDED IF I CAN PASS SESSION AND MODULE ID TO THE WORKERS
+// addHook({
+//   name: 'tinypool',
+//   versions: ['>=0.0.0']
+// }, tinypoolPackage => {
+//   debugger
+//   if (isPrivateTools(tinypoolPackage)) {
+//     // maybe get private set??
+//     const oldPrivateSet = tinypoolPackage.__privateSet
+
+//     tinypoolPackage.__privateSet = function (obj, member, value, setter) {
+//       if (isThreadPool(value)) {
+//         threadPool = value
+//         for (const worker of threadPool.workers.readyItems.values()) {
+//           debugger
+//           const oldWorkerMessage = worker.onMessage
+//           worker.onMessage = function (message) {
+//             debugger
+//             // we have to stop our message from being handled
+//             if (message.type === 'ci:vitest:worker:ready') {
+//               return
+//             }
+//             return oldWorkerMessage.apply(this, arguments)
+//           }
+//           const oldMessageHandler = worker.worker.thread._events.message
+
+//           worker.worker.thread.removeListener('message', oldMessageHandler)
+
+//           worker.worker.thread.on('message', function (message) {
+//             debugger
+//             // we have to stop our message from being handled
+//             if (message.type === 'ci:vitest:worker:ready') {
+//               return
+//             }
+//             return oldMessageHandler.apply(this, arguments)
+//           })
+
+//           // debugger
+//           // worker.worker.thread._events.message = function (message) {
+//           //   debugger
+//           //   // we have to stop our message from being handled
+//           //   if (message.type === 'ci:vitest:worker:ready') {
+//           //     return
+//           //   }
+//           //   return oldMessageHandler.apply(this, arguments)
+//           // }
+//         }
+//       }
+//       return oldPrivateSet.apply(this, arguments)
+//     }
+//   }
+//   if (isTinyPoolClass(tinypoolPackage)) {
+//     shimmer.wrap(tinypoolPackage.prototype, 'run', run => async function (task) {
+//       debugger
+//       const res = await run.apply(this, arguments)
+//       debugger
+//       return res
+//     })
+//   }
+//   return tinypoolPackage
+// })

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

@@ -37,15 +37,20 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addSub('ci:vitest:test-suite:start', (testSuiteAbsolutePath) => {
+      const testSessionSpanContext = this.tracer.extract('text_map', {
+        'x-datadog-trace-id': process.env.DD_CIVISIBILITY_TEST_SESSION_ID,
+        'x-datadog-parent-id': process.env.DD_CIVISIBILITY_TEST_MODULE_ID
+      })
+
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
       const testSuiteMetadata = getTestSuiteCommonTags(
         this.command,
         this.frameworkVersion,
         testSuite,
         'vitest'
       )
-      const testSuiteSpan = this.tracer.startSpan('mocha.test_suite', {
-        childOf: this.testModuleSpan,
+      const testSuiteSpan = this.tracer.startSpan('vitest.test_suite', {
+        childOf: testSessionSpanContext,
         tags: {
           [COMPONENT]: this.constructor.id,
           ...this.testEnvironmentMetadata,
@@ -57,28 +62,30 @@ class VitestPlugin extends CiPlugin {
       this.testSuiteSpan = testSuiteSpan
     })
 
-    this.addSub('ci:vitest:test-suite:finish', status => {
+    this.addSub('ci:vitest:test-suite:finish', ({ status, onFinish }) => {
       const store = storage.getStore()
       if (store && store.span) {
         const span = store.span
         span.setTag(TEST_STATUS, status)
         span.finish()
+        finishAllTraceSpans(span)
       }
+      this.tracer._exporter.flush(onFinish)
     })
 
-    // TODO: do we need to flush?
+    // TODO: do we need to flush? - probably not because it's just two spans in the main process
     this.addSub('ci:vitest:session:finish', (status) => {
-      console.log('session finish!', status)
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.finish()
       this.testSessionSpan.finish()
       finishAllTraceSpans(this.testSessionSpan)
     })
 
-    this.addSub('ci:vitest:run-files', onFinish => {
-      this.tracer._exporter.flush(onFinish)
-    })
+    // this.addSub('ci:vitest:run-files', onFinish => {
+    //   console.log('flushing in run files')
+    //   this.tracer._exporter.flush(onFinish)
+    // })
   }
 }
 

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

@@ -91,6 +91,10 @@ module.exports = class CiPlugin extends Plugin {
           ...testModuleSpanMetadata
         }
       })
+      // only for vitest
+      process.env.DD_CIVISIBILITY_TEST_SESSION_ID = this.testSessionSpan.context().toTraceId()
+      process.env.DD_CIVISIBILITY_TEST_MODULE_ID = this.testModuleSpan.context().toSpanId()
+
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_CREATED, 'module')
     })