wip

DataDog · Jun 18, 2024 · 689d9aa · 689d9aa
1 parent 4b03ee1
commit 689d9aa
Show file tree

Hide file tree

Showing 8 changed files with 24 additions and 5 deletions.
diff --git a/ci/init.js b/ci/init.js
@@ -1,10 +1,14 @@
 /* eslint-disable no-console */
 const tracer = require('../packages/dd-trace')
 const { isTrue } = require('../packages/dd-trace/src/util')
+const { isMainThread, parentPort } = require('node:worker_threads')
 
 const isJestWorker = !!process.env.JEST_WORKER_ID
 const isCucumberWorker = !!process.env.CUCUMBER_WORKER_ID
 const isMochaWorker = !!process.env.MOCHA_WORKER_ID
+// eslint-disable-next-line
+// https://github.com/vitest-dev/vitest/blob/f969fb0f9f0247a7daa2afee8f70de25ea5e123f/packages/vitest/src/node/pool.ts#L110-L111
+const isVitestWorker = !isMainThread && process.env.VITEST === 'true' && process.env.TEST === 'true'
 
 const options = {
   startupLogs: false,
@@ -51,6 +55,12 @@ if (isMochaWorker) {
   }
 }
 
+if (isVitestWorker) {
+  // if (parentPort?.postMessage) {
+  //   parentPort.postMessage({ type: 'ci:vitest:worker:ready' })
+  // }
+}
+
 if (shouldInit) {
   tracer.init(options)
   tracer.use('fs', false)

diff --git a/ext/exporters.d.ts b/ext/exporters.d.ts
@@ -4,8 +4,9 @@ declare const exporters: {
   DATADOG: 'datadog',
   AGENT_PROXY: 'agent_proxy',
   JEST_WORKER: 'jest_worker',
-  CUCUMBER_WORKER: 'cucumber_worker'
-  MOCHA_WORKER: 'mocha_worker'
+  CUCUMBER_WORKER: 'cucumber_worker',
+  MOCHA_WORKER: 'mocha_worker',
+  VITEST_WORKER: 'vitest_worker'
 }
 
 export = exporters
diff --git a/ext/exporters.js b/ext/exporters.js
@@ -6,5 +6,6 @@ module.exports = {
   AGENT_PROXY: 'agent_proxy',
   JEST_WORKER: 'jest_worker',
   CUCUMBER_WORKER: 'cucumber_worker',
-  MOCHA_WORKER: 'mocha_worker'
+  MOCHA_WORKER: 'mocha_worker',
+  VITEST_WORKER: 'vitest_worker'
 }
diff --git a/packages/datadog-instrumentations/src/helpers/hook.js b/packages/datadog-instrumentations/src/helpers/hook.js
@@ -20,7 +20,8 @@ function Hook (modules, onrequire) {
     const parts = [moduleBaseDir, moduleName].filter(v => v)
     const filename = path.join(...parts)
 
-    if (this._patched[filename]) return moduleExports
+    // this does not seem to work with IITM: the same file seems to import different stuff
+    // if (this._patched[filename]) return moduleExports
 
     this._patched[filename] = true
 
@@ -30,7 +31,7 @@ function Hook (modules, onrequire) {
   this._ritmHook = ritm(modules, {}, safeHook)
   this._iitmHook = iitm(modules, {}, (moduleExports, moduleName, moduleBaseDir) => {
     // TODO: Move this logic to import-in-the-middle and only do it for CommonJS
-    // modules and not ESM. In the meantime, all the modules we instrument are
+    // modules and not ESM. In the meantime, all thpe modules we instrument are
     // CommonJS modules for which the default export is always moved to
     // `default` anyway.
     if (moduleExports && moduleExports.default) {

diff --git a/packages/datadog-instrumentations/src/helpers/hooks.js b/packages/datadog-instrumentations/src/helpers/hooks.js
@@ -74,6 +74,7 @@ module.exports = {
   vitest: () => require('../vitest'),
   '@vitest': () => require('../vitest'),
   '@vitest/runner': () => require('../vitest'),
+  tinypool: () => require('../vitest'),
   workerpool: () => require('../mocha'),
   moleculer: () => require('../moleculer'),
   mongodb: () => require('../mongodb'),

diff --git a/packages/dd-trace/src/exporter.js b/packages/dd-trace/src/exporter.js
@@ -20,6 +20,7 @@ module.exports = name => {
     case exporters.JEST_WORKER:
     case exporters.CUCUMBER_WORKER:
     case exporters.MOCHA_WORKER:
+    case exporters.VITEST_WORKER:
       return require('./ci-visibility/exporters/test-worker')
     default:
       return inAWSLambda && !usingLambdaExtension ? require('./exporters/log') : require('./exporters/agent')

diff --git a/packages/dd-trace/src/iitm.js b/packages/dd-trace/src/iitm.js
@@ -8,6 +8,9 @@ const dc = require('dc-polyfill')
 if (semver.satisfies(process.versions.node, '>=14.13.1')) {
   const moduleLoadStartChannel = dc.channel('dd-trace:moduleLoadStart')
   addHook((name, namespace) => {
+    // if (name.includes('/vitest/') || name.includes('@vitest')) {
+    //   debugger
+    // }
     if (moduleLoadStartChannel.hasSubscribers) {
       moduleLoadStartChannel.publish({
         filename: name,

diff --git a/packages/dd-trace/src/plugins/index.js b/packages/dd-trace/src/plugins/index.js
@@ -57,6 +57,7 @@ module.exports = {
   get vitest () { return require('../../../datadog-plugin-vitest/src') },
   get '@vitest' () { return require('../../../datadog-plugin-vitest/src') },
   get '@vitest/runner' () { return require('../../../datadog-plugin-vitest/src') },
+  get tinypool () { return require('../../../datadog-plugin-vitest/src') },
   get workerpool () { return require('../../../datadog-plugin-mocha/src') },
   get moleculer () { return require('../../../datadog-plugin-moleculer/src') },
   get mongodb () { return require('../../../datadog-plugin-mongodb-core/src') },