[test optimization] Add Dynamic Instrumentation support for Vitest (#4959)

Author: juan-fernandez

integration-tests/ci-visibility/vitest-tests/bad-sum.mjs

@@ -0,0 +1,7 @@
+export function sum (a, b) {
+  const localVar = 10
+  if (a > 10) {
+    throw new Error('a is too large')
+  }
+  return a + b + localVar - localVar
+}

integration-tests/ci-visibility/vitest-tests/breakpoint-not-hit.mjs

@@ -0,0 +1,18 @@
+import { describe, test, expect } from 'vitest'
+import { sum } from './bad-sum'
+
+let numAttempt = 0
+
+describe('dynamic instrumentation', () => {
+  test('can sum', () => {
+    const shouldFail = numAttempt++ === 0
+    if (shouldFail) {
+      expect(sum(11, 2)).to.equal(13)
+    } else {
+      expect(sum(1, 2)).to.equal(3)
+    }
+  })
+  test('is not retried', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})

integration-tests/ci-visibility/vitest-tests/dynamic-instrumentation.mjs

@@ -0,0 +1,11 @@
+import { describe, test, expect } from 'vitest'
+import { sum } from './bad-sum'
+
+describe('dynamic instrumentation', () => {
+  test('can sum', () => {
+    expect(sum(11, 2)).to.equal(13)
+  })
+  test('is not retried', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})

integration-tests/vitest/vitest.spec.js

@@ -24,7 +24,11 @@ const {
   TEST_NAME,
   TEST_EARLY_FLAKE_ENABLED,
   TEST_EARLY_FLAKE_ABORT_REASON,
-  TEST_SUITE
+  TEST_SUITE,
+  DI_ERROR_DEBUG_INFO_CAPTURED,
+  DI_DEBUG_ERROR_FILE,
+  DI_DEBUG_ERROR_LINE,
+  DI_DEBUG_ERROR_SNAPSHOT_ID
 } = require('../../packages/dd-trace/src/plugins/util/test')
 const { DD_HOST_CPU_COUNT } = require('../../packages/dd-trace/src/plugins/util/env')
 
@@ -896,5 +900,204 @@ versions.forEach((version) => {
         })
       })
     })
+
+    // dynamic instrumentation only supported from >=2.0.0
+    if (version === 'latest') {
+      context('dynamic instrumentation', () => {
+        it('does not activate it if DD_TEST_DYNAMIC_INSTRUMENTATION_ENABLED is not set', (done) => {
+          receiver.setSettings({
+            itr_enabled: false,
+            code_coverage: false,
+            tests_skipping: false,
+            flaky_test_retries_enabled: false
+          })
+
+          const eventsPromise = receiver
+            .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/citestcycle'), (payloads) => {
+              const events = payloads.flatMap(({ payload }) => payload.events)
+
+              const tests = events.filter(event => event.type === 'test').map(event => event.content)
+              const retriedTests = tests.filter(test => test.meta[TEST_IS_RETRY] === 'true')
+
+              assert.equal(retriedTests.length, 1)
+              const [retriedTest] = retriedTests
+
+              assert.notProperty(retriedTest.meta, DI_ERROR_DEBUG_INFO_CAPTURED)
+              assert.notProperty(retriedTest.meta, DI_DEBUG_ERROR_FILE)
+              assert.notProperty(retriedTest.metrics, DI_DEBUG_ERROR_LINE)
+              assert.notProperty(retriedTest.meta, DI_DEBUG_ERROR_SNAPSHOT_ID)
+            })
+
+          const logsPromise = receiver
+            .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/logs'), (payloads) => {
+              if (payloads.length > 0) {
+                throw new Error('Unexpected logs')
+              }
+            }, 5000)
+
+          childProcess = exec(
+            './node_modules/.bin/vitest run --retry=1',
+            {
+              cwd,
+              env: {
+                ...getCiVisAgentlessConfig(receiver.port),
+                TEST_DIR: 'ci-visibility/vitest-tests/dynamic-instrumentation*',
+                NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init'
+              },
+              stdio: 'pipe'
+            }
+          )
+
+          childProcess.on('exit', () => {
+            Promise.all([eventsPromise, logsPromise]).then(() => {
+              done()
+            }).catch(done)
+          })
+        })
+
+        it('runs retries with dynamic instrumentation', (done) => {
+          receiver.setSettings({
+            itr_enabled: false,
+            code_coverage: false,
+            tests_skipping: false,
+            flaky_test_retries_enabled: false,
+            early_flake_detection: {
+              enabled: false
+            }
+            // di_enabled: true // TODO
+          })
+
+          let snapshotIdByTest, snapshotIdByLog
+          let spanIdByTest, spanIdByLog, traceIdByTest, traceIdByLog
+
+          const eventsPromise = receiver
+            .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/citestcycle'), (payloads) => {
+              const events = payloads.flatMap(({ payload }) => payload.events)
+
+              const tests = events.filter(event => event.type === 'test').map(event => event.content)
+              const retriedTests = tests.filter(test => test.meta[TEST_IS_RETRY] === 'true')
+
+              assert.equal(retriedTests.length, 1)
+              const [retriedTest] = retriedTests
+
+              assert.propertyVal(retriedTest.meta, DI_ERROR_DEBUG_INFO_CAPTURED, 'true')
+              assert.propertyVal(
+                retriedTest.meta,
+                DI_DEBUG_ERROR_FILE,
+                'ci-visibility/vitest-tests/bad-sum.mjs'
+              )
+              assert.equal(retriedTest.metrics[DI_DEBUG_ERROR_LINE], 4)
+              assert.exists(retriedTest.meta[DI_DEBUG_ERROR_SNAPSHOT_ID])
+
+              snapshotIdByTest = retriedTest.meta[DI_DEBUG_ERROR_SNAPSHOT_ID]
+              spanIdByTest = retriedTest.span_id.toString()
+              traceIdByTest = retriedTest.trace_id.toString()
+
+              const notRetriedTest = tests.find(test => test.meta[TEST_NAME].includes('is not retried'))
+
+              assert.notProperty(notRetriedTest.meta, DI_ERROR_DEBUG_INFO_CAPTURED)
+            })
+
+          const logsPromise = receiver
+            .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/logs'), (payloads) => {
+              const [{ logMessage: [diLog] }] = payloads
+              assert.deepInclude(diLog, {
+                ddsource: 'dd_debugger',
+                level: 'error'
+              })
+              assert.equal(diLog.debugger.snapshot.language, 'javascript')
+              assert.deepInclude(diLog.debugger.snapshot.captures.lines['4'].locals, {
+                a: {
+                  type: 'number',
+                  value: '11'
+                },
+                b: {
+                  type: 'number',
+                  value: '2'
+                },
+                localVar: {
+                  type: 'number',
+                  value: '10'
+                }
+              })
+              spanIdByLog = diLog.dd.span_id
+              traceIdByLog = diLog.dd.trace_id
+              snapshotIdByLog = diLog.debugger.snapshot.id
+            }, 5000)
+
+          childProcess = exec(
+            './node_modules/.bin/vitest run --retry=1',
+            {
+              cwd,
+              env: {
+                ...getCiVisAgentlessConfig(receiver.port),
+                TEST_DIR: 'ci-visibility/vitest-tests/dynamic-instrumentation*',
+                NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init',
+                DD_TEST_DYNAMIC_INSTRUMENTATION_ENABLED: '1'
+              },
+              stdio: 'pipe'
+            }
+          )
+
+          childProcess.on('exit', () => {
+            Promise.all([eventsPromise, logsPromise]).then(() => {
+              assert.equal(snapshotIdByTest, snapshotIdByLog)
+              assert.equal(spanIdByTest, spanIdByLog)
+              assert.equal(traceIdByTest, traceIdByLog)
+              done()
+            }).catch(done)
+          })
+        })
+
+        it('does not crash if the retry does not hit the breakpoint', (done) => {
+          const eventsPromise = receiver
+            .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/citestcycle'), (payloads) => {
+              const events = payloads.flatMap(({ payload }) => payload.events)
+
+              const tests = events.filter(event => event.type === 'test').map(event => event.content)
+              const retriedTests = tests.filter(test => test.meta[TEST_IS_RETRY] === 'true')
+
+              assert.equal(retriedTests.length, 1)
+              const [retriedTest] = retriedTests
+
+              assert.propertyVal(retriedTest.meta, DI_ERROR_DEBUG_INFO_CAPTURED, 'true')
+              assert.propertyVal(
+                retriedTest.meta,
+                DI_DEBUG_ERROR_FILE,
+                'ci-visibility/vitest-tests/bad-sum.mjs'
+              )
+              assert.equal(retriedTest.metrics[DI_DEBUG_ERROR_LINE], 4)
+              assert.exists(retriedTest.meta[DI_DEBUG_ERROR_SNAPSHOT_ID])
+            })
+
+          const logsPromise = receiver
+            .gatherPayloadsMaxTimeout(({ url }) => url.endsWith('/api/v2/logs'), (payloads) => {
+              if (payloads.length > 0) {
+                throw new Error('Unexpected logs')
+              }
+            }, 5000)
+
+          childProcess = exec(
+            './node_modules/.bin/vitest run --retry=1',
+            {
+              cwd,
+              env: {
+                ...getCiVisAgentlessConfig(receiver.port),
+                TEST_DIR: 'ci-visibility/vitest-tests/breakpoint-not-hit*',
+                NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init',
+                DD_TEST_DYNAMIC_INSTRUMENTATION_ENABLED: '1'
+              },
+              stdio: 'pipe'
+            }
+          )
+
+          childProcess.on('exit', () => {
+            Promise.all([eventsPromise, logsPromise]).then(() => {
+              done()
+            }).catch(done)
+          })
+        })
+      })
+    }
   })
 })

packages/datadog-instrumentations/src/vitest.js

@@ -316,12 +316,17 @@ addHook({
 
     // We finish the previous test here because we know it has failed already
     if (numAttempt > 0) {
+      const probe = {}
       const asyncResource = taskToAsync.get(task)
       const testError = task.result?.errors?.[0]
       if (asyncResource) {
         asyncResource.runInAsyncScope(() => {
-          testErrorCh.publish({ error: testError })
+          testErrorCh.publish({ error: testError, willBeRetried: true, probe })
         })
+        // We wait for the probe to be set
+        if (probe.setProbePromise) {
+          await probe.setProbePromise
+        }
       }
     }
 

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

@@ -17,7 +17,12 @@ const {
   TEST_SOURCE_START,
   TEST_IS_NEW,
   TEST_EARLY_FLAKE_ENABLED,
-  TEST_EARLY_FLAKE_ABORT_REASON
+  TEST_EARLY_FLAKE_ABORT_REASON,
+  TEST_NAME,
+  DI_ERROR_DEBUG_INFO_CAPTURED,
+  DI_DEBUG_ERROR_SNAPSHOT_ID,
+  DI_DEBUG_ERROR_FILE,
+  DI_DEBUG_ERROR_LINE
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const {
@@ -31,6 +36,8 @@ const {
 // This is because there's some loss of resolution.
 const MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION = 5
 
+const debuggerParameterPerTest = new Map()
+
 class VitestPlugin extends CiPlugin {
   static get id () {
     return 'vitest'
@@ -81,6 +88,26 @@ class VitestPlugin extends CiPlugin {
         extraTags
       )
 
+      const debuggerParameters = debuggerParameterPerTest.get(testName)
+
+      if (debuggerParameters) {
+        const spanContext = span.context()
+
+        // TODO: handle race conditions with this.retriedTestIds
+        this.retriedTestIds = {
+          spanId: spanContext.toSpanId(),
+          traceId: spanContext.toTraceId()
+        }
+        const { snapshotId, file, line } = debuggerParameters
+
+        // TODO: should these be added on test:end if and only if the probe is hit?
+        // Sync issues: `hitProbePromise` might be resolved after the test ends
+        span.setTag(DI_ERROR_DEBUG_INFO_CAPTURED, 'true')
+        span.setTag(DI_DEBUG_ERROR_SNAPSHOT_ID, snapshotId)
+        span.setTag(DI_DEBUG_ERROR_FILE, file)
+        span.setTag(DI_DEBUG_ERROR_LINE, line)
+      }
+
       this.enter(span, store)
     })
 
@@ -110,11 +137,16 @@ class VitestPlugin extends CiPlugin {
       }
     })
 
-    this.addSub('ci:vitest:test:error', ({ duration, error }) => {
+    this.addSub('ci:vitest:test:error', ({ duration, error, willBeRetried, probe }) => {
       const store = storage.getStore()
       const span = store?.span
 
       if (span) {
+        if (willBeRetried && this.di) {
+          const testName = span.context()._tags[TEST_NAME]
+          const debuggerParameters = this.addDiProbe(error, probe)
+          debuggerParameterPerTest.set(testName, debuggerParameters)
+        }
         this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', {
           hasCodeowners: !!span.context()._tags[TEST_CODE_OWNERS]
         })

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

@@ -75,18 +75,20 @@ async function addBreakpoint (snapshotId, probe) {
 
   log.debug(`Adding breakpoint at ${path}:${line}`)
 
-  let generatedPosition = { line }
-  let hasSourceMap = false
+  let lineNumber = line
 
   if (sourceMapURL && sourceMapURL.startsWith('data:')) {
-    hasSourceMap = true
-    generatedPosition = await processScriptWithInlineSourceMap({ file, line, sourceMapURL })
+    try {
+      lineNumber = await processScriptWithInlineSourceMap({ file, line, sourceMapURL })
+    } catch (err) {
+      log.error(err)
+    }
   }
 
   const { breakpointId } = await session.post('Debugger.setBreakpoint', {
     location: {
       scriptId,
-      lineNumber: hasSourceMap ? generatedPosition.line : generatedPosition.line - 1
+      lineNumber: lineNumber - 1
     }
   })
 
@@ -120,5 +122,5 @@ async function processScriptWithInlineSourceMap (params) {
 
   consumer.destroy()
 
-  return generatedPosition
+  return generatedPosition.line
 }