[test optimization] Add support for --workerThreads flag in jest (#7840)

Author: juan-fernandez

integration-tests/ci-visibility/jest-plugin-tests/jest-worker-message.js

@@ -0,0 +1,14 @@
+'use strict'
+
+const assert = require('node:assert/strict')
+const { parentPort } = require('node:worker_threads')
+
+if (parentPort) {
+  parentPort.postMessage({ source: 'jest-worker-message' })
+}
+
+describe('jest-worker-message', () => {
+  it('passes after sending a non-array worker message', () => {
+    assert.strictEqual(true, true)
+  })
+})

integration-tests/ci-visibility/run-jest.js

@@ -21,6 +21,12 @@ if (process.env.RUN_IN_PARALLEL) {
   options.maxWorkers = 2
 }
 
+if (process.env.USE_WORKER_THREADS) {
+  delete options.runInBand
+  options.maxWorkers = 2
+  options.workerThreads = true
+}
+
 if (process.env.OLD_RUNNER) {
   options.testRunner = 'jest-jasmine2'
 }

integration-tests/jest/jest.spec.js

@@ -998,6 +998,93 @@ describe(`jest@${JEST_VERSION} commonJS`, () => {
     })
   })
 
+  context('when jest is using worker threads', () => {
+    onlyLatestIt('ignores non-array worker-thread messages', (done) => {
+      childProcess = fork(testFile, {
+        cwd,
+        env: {
+          ...getCiVisAgentlessConfig(receiver.port),
+          TESTS_TO_RUN: 'jest-plugin-tests/jest-worker-message',
+          USE_WORKER_THREADS: 'true',
+        },
+        stdio: 'pipe',
+      })
+      childProcess.stdout?.on('data', (chunk) => {
+        testOutput += chunk.toString()
+      })
+      childProcess.stderr?.on('data', (chunk) => {
+        testOutput += chunk.toString()
+      })
+
+      Promise.all([
+        once(childProcess, 'message'),
+        receiver.gatherPayloads(({ url }) => url === '/api/v2/citestcycle', 5000),
+      ]).then(([, eventsRequests]) => {
+        const tests = eventsRequests.map(({ payload }) => payload)
+          .flatMap(({ events }) => events)
+          .filter(event => event.type === 'test')
+          .map(event => event.content)
+
+        assert.strictEqual(tests.length, 1)
+        assert.strictEqual(
+          tests[0].meta[TEST_NAME],
+          'jest-worker-message passes after sending a non-array worker message'
+        )
+        assert.strictEqual(tests[0].meta[TEST_STATUS], 'pass')
+        assert.doesNotMatch(testOutput, /TypeError/)
+        done()
+      }).catch(done)
+    })
+
+    onlyLatestIt('reports tests when using agentless', (done) => {
+      childProcess = fork(testFile, {
+        cwd,
+        env: {
+          ...getCiVisAgentlessConfig(receiver.port),
+          USE_WORKER_THREADS: 'true',
+        },
+        stdio: 'pipe',
+      })
+
+      receiver.gatherPayloads(({ url }) => url === '/api/v2/citestcycle', 5000).then(eventsRequests => {
+        const events = eventsRequests.map(({ payload }) => payload)
+          .flatMap(({ events }) => events)
+        const eventTypes = events.map(event => event.type)
+        assertObjectContains(eventTypes, ['test', 'test_suite_end', 'test_session_end', 'test_module_end'])
+
+        const tests = events.filter(event => event.type === 'test').map(event => event.content)
+        assert.ok(tests.length >= 2)
+        tests.forEach(testEvent => {
+          assert.strictEqual(testEvent.meta[TEST_STATUS], 'pass')
+        })
+
+        done()
+      }).catch(done)
+    })
+
+    onlyLatestIt('reports tests when using evp proxy', (done) => {
+      receiver.setInfoResponse({ endpoints: ['/evp_proxy/v2'] })
+      childProcess = fork(testFile, {
+        cwd,
+        env: {
+          ...getCiVisEvpProxyConfig(receiver.port),
+          USE_WORKER_THREADS: 'true',
+        },
+        stdio: 'pipe',
+      })
+
+      receiver.gatherPayloads(({ url }) => url === '/evp_proxy/v2/api/v2/citestcycle', 5000)
+        .then(eventsRequests => {
+          const eventTypes = eventsRequests.map(({ payload }) => payload)
+            .flatMap(({ events }) => events)
+            .map(event => event.type)
+
+          assertObjectContains(eventTypes, ['test', 'test_suite_end', 'test_session_end', 'test_module_end'])
+          done()
+        }).catch(done)
+    })
+  })
+
   it('reports timeout error message', (done) => {
     childProcess = fork(testFile, {
       cwd,

packages/datadog-instrumentations/src/jest.js

@@ -1834,7 +1834,12 @@ addHook({
 
 function onMessageWrapper (onMessage) {
   return function () {
-    const [code, data] = arguments[0]
+    const response = arguments[0]
+    if (!Array.isArray(response)) {
+      return onMessage.apply(this, arguments)
+    }
+
+    const [code, data] = response
     if (code === JEST_WORKER_TRACE_PAYLOAD_CODE) { // datadog trace payload
       collectDynamicNamesFromTraces(data, newTestsWithDynamicNames)
       workerReportTraceCh.publish(data)
@@ -1897,14 +1902,22 @@ function sendWrapper (send) {
   }
 }
 
+function wrapWorker (worker) {
+  // ChildProcessWorker uses _child (child_process), ExperimentalWorker uses _worker (worker_threads)
+  const workerChannel = worker._child || worker._worker
+  if (!workerChannel) return
+
+  shimmer.wrap(workerChannel, worker._child ? 'send' : 'postMessage', sendWrapper)
+  shimmer.wrap(worker, '_onMessage', onMessageWrapper)
+  workerChannel.removeAllListeners('message')
+  workerChannel.on('message', worker._onMessage.bind(worker))
+}
+
 function enqueueWrapper (enqueue) {
   return function () {
     shimmer.wrap(arguments[0], 'onStart', onStart => function (worker) {
       if (worker && !wrappedWorkers.has(worker)) {
-        shimmer.wrap(worker._child, 'send', sendWrapper)
-        shimmer.wrap(worker, '_onMessage', onMessageWrapper)
-        worker._child.removeAllListeners('message')
-        worker._child.on('message', worker._onMessage.bind(worker))
+        wrapWorker(worker)
         wrappedWorkers.add(worker)
       }
       return onStart.apply(this, arguments)
@@ -1934,6 +1947,21 @@ addHook({
   return childProcessWorker
 })
 
+addHook({
+  name: 'jest-worker',
+  versions: ['>=24.9.0 <30.0.0'],
+  file: 'build/workers/NodeThreadsWorker.js',
+}, (nodeThreadsWorker) => {
+  const ExperimentalWorker = nodeThreadsWorker.default
+  shimmer.wrap(ExperimentalWorker.prototype, 'send', sendWrapper)
+  if (ExperimentalWorker.prototype._onMessage) {
+    shimmer.wrap(ExperimentalWorker.prototype, '_onMessage', onMessageWrapper)
+  } else if (ExperimentalWorker.prototype.onMessage) {
+    shimmer.wrap(ExperimentalWorker.prototype, 'onMessage', onMessageWrapper)
+  }
+  return nodeThreadsWorker
+})
+
 addHook({
   name: 'jest-worker',
   versions: ['>=30.0.0'],

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

@@ -2,10 +2,6 @@
 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) {
@@ -29,12 +25,6 @@ class Writer {
   }
 
   _sendPayload (data, onDone = () => {}) {
-    // ## Jest
-    // Only available when `child_process` is used for the jest worker.
-    // If worker_threads is used, this will not work
-    // 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:
@@ -47,29 +37,28 @@ class Writer {
       ? { __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
-
+    // child_process workers (jest default, cucumber)
     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()
-      }
+      return
+    }
+
+    // worker_threads (jest --workerThreads, vitest)
+    const { isMainThread, parentPort } = require('node:worker_threads')
+    if (!isMainThread && parentPort) {
       try {
         parentPort.postMessage(payload)
       } catch (error) {
         log.error('Error posting message to parent port', error)
       } finally {
         onDone()
       }
-    } else {
-      onDone()
+      return
     }
+
+    onDone()
   }
 }