[test optimization] Fix threads config in vitest (#6824)

Author: juan-fernandez

integration-tests/vitest/vitest.spec.js

@@ -52,7 +52,8 @@ const {
   TEST_RETRY_REASON_TYPES,
   TEST_IS_MODIFIED,
   DD_CAPABILITIES_IMPACTED_TESTS,
-  VITEST_POOL
+  VITEST_POOL,
+  TEST_IS_TEST_FRAMEWORK_WORKER
 } = require('../../packages/dd-trace/src/plugins/util/test')
 const { DD_HOST_CPU_COUNT } = require('../../packages/dd-trace/src/plugins/util/env')
 const { NODE_MAJOR } = require('../../version')
@@ -209,12 +210,20 @@ versions.forEach((version) => {
             )
 
             testEvents.forEach(test => {
+              // `threads` config will report directly. TODO: update this once we're testing vitest@>=4
+              if (poolConfig === 'forks') {
+                assert.equal(test.content.meta[TEST_IS_TEST_FRAMEWORK_WORKER], 'true')
+              }
               assert.equal(test.content.meta[TEST_COMMAND], 'vitest run')
               assert.exists(test.content.metrics[DD_HOST_CPU_COUNT])
               assert.equal(test.content.meta[DD_TEST_IS_USER_PROVIDED_SERVICE], 'false')
             })
 
             testSuiteEvents.forEach(testSuite => {
+              // `threads` config will report directly. TODO: update this once we're testing vitest@>=4
+              if (poolConfig === 'forks') {
+                assert.equal(testSuite.content.meta[TEST_IS_TEST_FRAMEWORK_WORKER], 'true')
+              }
               assert.equal(testSuite.content.meta[TEST_COMMAND], 'vitest run')
               assert.isTrue(
                 testSuite.content.meta[TEST_SOURCE_FILE].startsWith('ci-visibility/vitest-tests/test-visibility')

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

@@ -345,7 +345,6 @@ class VitestPlugin extends CiPlugin {
         finishAllTraceSpans(testSuiteSpan)
       }
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
-      // TODO: too frequent flush - find for method in worker to decrease frequency
       this.tracer._exporter.flush(onFinish)
       if (this.runningTestProbe) {
         this.removeDiProbe(this.runningTestProbe)

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

@@ -1,6 +1,11 @@
 'use strict'
 const { JSONEncoder } = require('../../encode/json-encoder')
 const { getEnvironmentVariable } = require('../../../config-helper')
+const log = require('../../../log')
+const {
+  VITEST_WORKER_TRACE_PAYLOAD_CODE,
+  VITEST_WORKER_LOGS_PAYLOAD_CODE
+} = require('../../../plugins/util/test')
 
 class Writer {
   constructor (interprocessCode) {
@@ -26,25 +31,42 @@ class Writer {
   _sendPayload (data, onDone = () => {}) {
     // ## Jest
     // Only available when `child_process` is used for the jest worker.
-    // https://github.com/facebook/jest/blob/bb39cb2c617a3334bf18daeca66bd87b7ccab28b/packages/jest-worker/README.md#experimental-worker
     // If worker_threads is used, this will not work
-    // TODO: make it compatible with worker_threads
+    // TODO: make `jest` instrumentation compatible with worker_threads
+    // https://github.com/facebook/jest/blob/bb39cb2c617a3334bf18daeca66bd87b7ccab28b/packages/jest-worker/README.md#experimental-worker
 
     // ## Cucumber
     // This reports to the test's main process the same way test data is reported by Cucumber
     // 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
-      // Old because vitest@>=4 use DD_VITEST_WORKER and report arrays just like other frameworks
-      const isVitestWorkerOld = !!getEnvironmentVariable('TINYPOOL_WORKER_ID')
 
-      const payload = isVitestWorkerOld
-        ? { __tinypool_worker_message__: true, interprocessCode: this._interprocessCode, data }
-        : [this._interprocessCode, data]
+    // Old because vitest@>=4 uses `DD_VITEST_WORKER` and reports arrays just like other frameworks
+    // Before vitest@>=4, we need the `__tinypool_worker_message__` property, or tinypool will crash
+    const isVitestWorkerOld = !!getEnvironmentVariable('TINYPOOL_WORKER_ID')
+    const payload = isVitestWorkerOld
+      ? { __tinypool_worker_message__: true, interprocessCode: this._interprocessCode, data }
+      : [this._interprocessCode, data]
 
+    const isVitestTestWorker =
+      this._interprocessCode === VITEST_WORKER_TRACE_PAYLOAD_CODE ||
+      this._interprocessCode === VITEST_WORKER_LOGS_PAYLOAD_CODE
+
+    if (process.send) {
       process.send(payload, () => {
         onDone()
       })
+    } else if (isVitestTestWorker) { // TODO: worker_threads are only supported in vitest right now
+      const { isMainThread, parentPort } = require('worker_threads')
+      if (isMainThread) {
+        return onDone()
+      }
+      try {
+        parentPort.postMessage(payload)
+      } catch (error) {
+        log.error('Error posting message to parent port', error)
+      } finally {
+        onDone()
+      }
     } else {
       onDone()
     }

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

@@ -34,7 +34,8 @@ const {
   getLibraryCapabilitiesTags,
   getPullRequestDiff,
   getModifiedFilesFromDiff,
-  getPullRequestBaseBranch
+  getPullRequestBaseBranch,
+  TEST_IS_TEST_FRAMEWORK_WORKER
 } = require('./util/test')
 const { getRepositoryRoot } = require('./util/git')
 const Plugin = require('./plugin')
@@ -311,6 +312,7 @@ module.exports = class CiPlugin extends Plugin {
           span.parent_id = id(span.parent_id)
 
           if (span.name?.startsWith(`${this.constructor.id}.`)) {
+            span.meta[TEST_IS_TEST_FRAMEWORK_WORKER] = 'true'
             // augment with git information (since it will not be available in the worker)
             for (const key in this.testEnvironmentMetadata) {
               // CAREFUL: this bypasses the metadata/metrics distinction

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

@@ -132,6 +132,8 @@ const PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE = 90
 const VITEST_WORKER_TRACE_PAYLOAD_CODE = 100
 const VITEST_WORKER_LOGS_PAYLOAD_CODE = 102
 
+const TEST_IS_TEST_FRAMEWORK_WORKER = 'test.is_test_framework_worker'
+
 // Library Capabilities Tagging
 const DD_CAPABILITIES_TEST_IMPACT_ANALYSIS = '_dd.library_capabilities.test_impact_analysis'
 const DD_CAPABILITIES_EARLY_FLAKE_DETECTION = '_dd.library_capabilities.early_flake_detection'
@@ -226,6 +228,7 @@ module.exports = {
   PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE,
   VITEST_WORKER_TRACE_PAYLOAD_CODE,
   VITEST_WORKER_LOGS_PAYLOAD_CODE,
+  TEST_IS_TEST_FRAMEWORK_WORKER,
   TEST_SOURCE_START,
   TEST_SKIPPED_BY_ITR,
   TEST_IS_NEW,