[test optimization][SDTEST-1355] Fix ATR + DI

packages/datadog-instrumentations/src/jest.js

@@ -313,10 +313,11 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
         const asyncResource = asyncResources.get(event.test)
 
         if (status === 'fail') {
+          const shouldSetProbe = this.isDiEnabled && willBeRetried && numTestExecutions === 1
           asyncResource.runInAsyncScope(() => {
             testErrCh.publish({
               error: formatJestError(event.test.errors[0]),
-              shouldSetProbe: this.isDiEnabled && willBeRetried && numTestExecutions === 1,
+              shouldSetProbe,
               promises
             })
           })
@@ -336,18 +337,13 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
           testFinishCh.publish({
             status,
             testStartLine: getTestLineStart(event.test.asyncError, this.testSuite),
-            promises,
-            shouldRemoveProbe: this.isDiEnabled && !willBeRetried
+            promises
           })
         })
 
         if (promises.isProbeReady) {
           await promises.isProbeReady
         }
-
-        if (promises.isProbeRemoved) {
-          await promises.isProbeRemoved
-        }
       }
       if (event.name === 'test_skip' || event.name === 'test_todo') {
         const asyncResource = new AsyncResource('bound-anonymous-fn')

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

@@ -231,7 +231,7 @@ class CucumberPlugin extends CiPlugin {
 
       this.activeTestSpan = testSpan
       // Time we give the breakpoint to be hit
-      if (promises && this.runningTestProbeId) {
+      if (promises && this.runningTestProbe) {
         promises.hitBreakpointPromise = new Promise((resolve) => {
           setTimeout(resolve, BREAKPOINT_HIT_GRACE_PERIOD_MS)
         })
@@ -248,8 +248,8 @@ class CucumberPlugin extends CiPlugin {
       if (isFirstAttempt && this.di && error && this.libraryConfig?.isDiEnabled) {
         const probeInformation = this.addDiProbe(error)
         if (probeInformation) {
-          const { probeId, stackIndex } = probeInformation
-          this.runningTestProbeId = probeId
+          const { file, line, stackIndex } = probeInformation
+          this.runningTestProbe = { file, line }
           this.testErrorStackIndex = stackIndex
           // TODO: we're not waiting for setProbePromise to be resolved, so there might be race conditions
         }
@@ -359,9 +359,9 @@ class CucumberPlugin extends CiPlugin {
           this.tracer._exporter.flush()
         }
         this.activeTestSpan = null
-        if (this.runningTestProbeId) {
-          this.removeDiProbe(this.runningTestProbeId)
-          this.runningTestProbeId = null
+        if (this.runningTestProbe) {
+          this.removeDiProbe(this.runningTestProbe)
+          this.runningTestProbe = null
         }
       }
     })

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

@@ -291,6 +291,7 @@ class JestPlugin extends CiPlugin {
       if (isJestWorker) {
         this.tracer._exporter.flush()
       }
+      this.removeAllDiProbes()
     })
 
     /**
@@ -324,7 +325,7 @@ class JestPlugin extends CiPlugin {
       this.activeTestSpan = span
     })
 
-    this.addSub('ci:jest:test:finish', ({ status, testStartLine, promises, shouldRemoveProbe }) => {
+    this.addSub('ci:jest:test:finish', ({ status, testStartLine }) => {
       const span = storage.getStore().span
       span.setTag(TEST_STATUS, status)
       if (testStartLine) {
@@ -346,10 +347,6 @@ class JestPlugin extends CiPlugin {
       span.finish()
       finishAllTraceSpans(span)
       this.activeTestSpan = null
-      if (shouldRemoveProbe && this.runningTestProbeId) {
-        promises.isProbeRemoved = withTimeout(this.removeDiProbe(this.runningTestProbeId), 2000)
-        this.runningTestProbeId = null
-      }
     })
 
     this.addSub('ci:jest:test:err', ({ error, shouldSetProbe, promises }) => {
@@ -362,9 +359,7 @@ class JestPlugin extends CiPlugin {
           if (shouldSetProbe) {
             const probeInformation = this.addDiProbe(error)
             if (probeInformation) {
-              const { probeId, setProbePromise, stackIndex } = probeInformation
-              this.runningTestProbeId = probeId
-              this.testErrorStackIndex = stackIndex
+              const { setProbePromise } = probeInformation
               promises.isProbeReady = withTimeout(setProbePromise, 2000)
             }
           }

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

@@ -219,9 +219,9 @@ class MochaPlugin extends CiPlugin {
         span.finish()
         finishAllTraceSpans(span)
         this.activeTestSpan = null
-        if (this.di && this.libraryConfig?.isDiEnabled && this.runningTestProbeId && isLastRetry) {
-          this.removeDiProbe(this.runningTestProbeId)
-          this.runningTestProbeId = null
+        if (this.di && this.libraryConfig?.isDiEnabled && this.runningTestProbe && isLastRetry) {
+          this.removeDiProbe(this.runningTestProbe)
+          this.runningTestProbe = null
         }
       }
     })
@@ -275,8 +275,8 @@ class MochaPlugin extends CiPlugin {
         if (isFirstAttempt && willBeRetried && this.di && this.libraryConfig?.isDiEnabled) {
           const probeInformation = this.addDiProbe(err)
           if (probeInformation) {
-            const { probeId, stackIndex } = probeInformation
-            this.runningTestProbeId = probeId
+            const { file, line, stackIndex } = probeInformation
+            this.runningTestProbe = { file, line }
             this.testErrorStackIndex = stackIndex
             test._ddShouldWaitForHitProbe = true
             // TODO: we're not waiting for setProbePromise to be resolved, so there might be race conditions

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

@@ -135,8 +135,8 @@ class VitestPlugin extends CiPlugin {
         if (shouldSetProbe && this.di) {
           const probeInformation = this.addDiProbe(error)
           if (probeInformation) {
-            const { probeId, stackIndex, setProbePromise } = probeInformation
-            this.runningTestProbeId = probeId
+            const { file, line, stackIndex, setProbePromise } = probeInformation
+            this.runningTestProbe = { file, line }
             this.testErrorStackIndex = stackIndex
             promises.setProbePromise = setProbePromise
           }
@@ -237,8 +237,8 @@ class VitestPlugin extends CiPlugin {
       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.runningTestProbeId) {
-        this.removeDiProbe(this.runningTestProbeId)
+      if (this.runningTestProbe) {
+        this.removeDiProbe(this.runningTestProbe)
       }
     })
 

packages/dd-trace/src/ci-visibility/dynamic-instrumentation/index.js

@@ -119,6 +119,8 @@ class TestVisDynamicInstrumentation {
       const onHit = this.onHitBreakpointByProbeId.get(probeId)
       if (onHit) {
         onHit({ snapshot })
+      } else {
+        log.warn('Received a breakpoint hit for an unknown probe')
       }
     }).unref()
 

packages/dd-trace/src/ci-visibility/dynamic-instrumentation/worker/index.js

@@ -93,11 +93,14 @@ async function addBreakpoint (probe) {
   probe.location = { file, lines: [String(line)] }
 
   const script = findScriptFromPartialPath(file)
-  if (!script) throw new Error(`No loaded script found for ${file}`)
+  if (!script) {
+    log.error(`No loaded script found for ${file}`)
+    throw new Error(`No loaded script found for ${file}`)
+  }
 
   const [path, scriptId, sourceMapURL] = script
 
-  log.debug(`Adding breakpoint at ${path}:${line}`)
+  log.warn(`Adding breakpoint at ${path}:${line}`)
 
   let lineNumber = line
 
@@ -109,15 +112,19 @@ async function addBreakpoint (probe) {
     }
   }
 
-  const { breakpointId } = await session.post('Debugger.setBreakpoint', {
-    location: {
-      scriptId,
-      lineNumber: lineNumber - 1
-    }
-  })
+  try {
+    const { breakpointId } = await session.post('Debugger.setBreakpoint', {
+      location: {
+        scriptId,
+        lineNumber: lineNumber - 1
+      }
+    })
 
-  breakpointIdToProbe.set(breakpointId, probe)
-  probeIdToBreakpointId.set(probe.id, breakpointId)
+    breakpointIdToProbe.set(breakpointId, probe)
+    probeIdToBreakpointId.set(probe.id, breakpointId)
+  } catch (e) {
+    log.error(`Error setting breakpoint at ${path}:${line}:`, e)
+  }
 }
 
 function start () {

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

@@ -45,6 +45,7 @@ module.exports = class CiPlugin extends Plugin {
   constructor (...args) {
     super(...args)
 
+    this.fileLineToProbeId = new Map()
     this.rootDir = process.cwd() // fallback in case :session:start events are not emitted
 
     this.addSub(`ci:${this.constructor.id}:library-configuration`, ({ onDone }) => {
@@ -335,7 +336,22 @@ module.exports = class CiPlugin extends Plugin {
     })
   }
 
-  removeDiProbe (probeId) {
+  removeAllDiProbes () {
+    if (this.fileLineToProbeId.size === 0) {
+      return Promise.resolve()
+    }
+    log.debug('Removing all Dynamic Instrumentation probes')
+    return Promise.all(Array.from(this.fileLineToProbeId.keys())
+      .map((fileLine) => {
+        const [file, line] = fileLine.split(':')
+        return this.removeDiProbe({ file, line })
+      }))
+  }
+
+  removeDiProbe ({ file, line }) {
+    const probeId = this.fileLineToProbeId.get(`${file}:${line}`)
+    log.warn(`Removing probe from ${file}:${line}, with id: ${probeId}`)
+    this.fileLineToProbeId.delete(probeId)
     return this.di.removeProbe(probeId)
   }
 
@@ -346,9 +362,27 @@ module.exports = class CiPlugin extends Plugin {
       log.warn('Could not add breakpoint for dynamic instrumentation')
       return
     }
+    log.debug('Adding breakpoint for Dynamic Instrumentation')
+
+    this.testErrorStackIndex = stackIndex
+    const activeProbeKey = `${file}:${line}`
+
+    if (this.fileLineToProbeId.has(activeProbeKey)) {
+      log.warn('Probe already set for this line')
+      const oldProbeId = this.fileLineToProbeId.get(activeProbeKey)
+      return {
+        probeId: oldProbeId,
+        setProbePromise: Promise.resolve(),
+        stackIndex,
+        file,
+        line
+      }
+    }
 
     const [probeId, setProbePromise] = this.di.addLineProbe({ file, line }, this.onDiBreakpointHit.bind(this))
 
+    this.fileLineToProbeId.set(activeProbeKey, probeId)
+
     return {
       probeId,
       setProbePromise,