fix performance

Author: juan-fernandez

packages/datadog-instrumentations/src/vitest.js

@@ -369,7 +369,6 @@ 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
@@ -386,29 +385,31 @@ addHook({
   return suitePackage
 })
 
-const processToHandler = new WeakMap()
+const processToHandler = new WeakSet()
+
+function threadHandler (thread) {
+  if (processToHandler.has(thread.process)) {
+    return
+  }
+  processToHandler.add(thread.process)
+  thread.process.on('message', (message) => {
+    if (message.__tinypool_worker_message__ && message.data) {
+      workerReporterCh.publish(message.data)
+    }
+  })
+}
 
 addHook({
   name: 'tinypool',
   versions: ['>=1.0.0'],
   file: 'dist/index.js'
 }, (TinyPool) => {
-  // 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 => {
-      if (processToHandler.has(thread.process)) {
-        return
-      }
-      processToHandler.set(thread.process, true)
-      thread.process.on('message', (message) => {
-        if (message.__tinypool_worker_message__ && message.data) {
-          workerReporterCh.publish(message.data)
-        }
-      })
-    })
-
+  shimmer.wrap(TinyPool.prototype, 'run', run => async function () {
+    // we need to do this before and after because the threads list gets recycled
+    // (the processes are re-created)
+    this.threads.forEach(threadHandler)
+    const res = await run.apply(this, arguments)
+    this.threads.forEach(threadHandler)
     return res
   })
 

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

@@ -1,5 +1,8 @@
 'use strict'
 const { JSONEncoder } = require('../../encode/json-encoder')
+const { getEnvironmentVariable } = require('../../../config-helper')
+
+const isVitestWorker = !!getEnvironmentVariable('TINYPOOL_WORKER_ID')
 
 class Writer {
   constructor (interprocessCode) {
@@ -34,7 +37,7 @@ 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
-      if (process.env.TINYPOOL_WORKER_ID) {
+      if (isVitestWorker) {
         // 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 }, () => {
@@ -45,8 +48,9 @@ class Writer {
           onDone()
         })
       }
+    } else {
+      onDone()
     }
-    // onDone()
   }
 }