Skip to content

Commit 51a58bc

Browse files
juan-fernandezjuan-fernandez
authoredFeb 3, 2025
[test optimization][SDTEST-1355] Fix ATR + DI (#5176)
1 parent ee6dbec commit 51a58bc

File tree

8 files changed

+75
-41
lines changed

8 files changed

+75
-41
lines changed
 

‎packages/datadog-instrumentations/src/jest.js

+3-7
Original file line numberDiff line numberDiff line change
@@ -313,10 +313,11 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
313313
const asyncResource = asyncResources.get(event.test)
314314

315315
if (status === 'fail') {
316+
const shouldSetProbe = this.isDiEnabled && willBeRetried && numTestExecutions === 1
316317
asyncResource.runInAsyncScope(() => {
317318
testErrCh.publish({
318319
error: formatJestError(event.test.errors[0]),
319-
shouldSetProbe: this.isDiEnabled && willBeRetried && numTestExecutions === 1,
320+
shouldSetProbe,
320321
promises
321322
})
322323
})
@@ -336,18 +337,13 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
336337
testFinishCh.publish({
337338
status,
338339
testStartLine: getTestLineStart(event.test.asyncError, this.testSuite),
339-
promises,
340-
shouldRemoveProbe: this.isDiEnabled && !willBeRetried
340+
promises
341341
})
342342
})
343343

344344
if (promises.isProbeReady) {
345345
await promises.isProbeReady
346346
}
347-
348-
if (promises.isProbeRemoved) {
349-
await promises.isProbeRemoved
350-
}
351347
}
352348
if (event.name === 'test_skip' || event.name === 'test_todo') {
353349
const asyncResource = new AsyncResource('bound-anonymous-fn')

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

+6-6
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,7 @@ class CucumberPlugin extends CiPlugin {
231231

232232
this.activeTestSpan = testSpan
233233
// Time we give the breakpoint to be hit
234-
if (promises && this.runningTestProbeId) {
234+
if (promises && this.runningTestProbe) {
235235
promises.hitBreakpointPromise = new Promise((resolve) => {
236236
setTimeout(resolve, BREAKPOINT_HIT_GRACE_PERIOD_MS)
237237
})
@@ -248,8 +248,8 @@ class CucumberPlugin extends CiPlugin {
248248
if (isFirstAttempt && this.di && error && this.libraryConfig?.isDiEnabled) {
249249
const probeInformation = this.addDiProbe(error)
250250
if (probeInformation) {
251-
const { probeId, stackIndex } = probeInformation
252-
this.runningTestProbeId = probeId
251+
const { file, line, stackIndex } = probeInformation
252+
this.runningTestProbe = { file, line }
253253
this.testErrorStackIndex = stackIndex
254254
// TODO: we're not waiting for setProbePromise to be resolved, so there might be race conditions
255255
}
@@ -359,9 +359,9 @@ class CucumberPlugin extends CiPlugin {
359359
this.tracer._exporter.flush()
360360
}
361361
this.activeTestSpan = null
362-
if (this.runningTestProbeId) {
363-
this.removeDiProbe(this.runningTestProbeId)
364-
this.runningTestProbeId = null
362+
if (this.runningTestProbe) {
363+
this.removeDiProbe(this.runningTestProbe)
364+
this.runningTestProbe = null
365365
}
366366
}
367367
})

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

+3-8
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,7 @@ class JestPlugin extends CiPlugin {
291291
if (isJestWorker) {
292292
this.tracer._exporter.flush()
293293
}
294+
this.removeAllDiProbes()
294295
})
295296

296297
/**
@@ -324,7 +325,7 @@ class JestPlugin extends CiPlugin {
324325
this.activeTestSpan = span
325326
})
326327

327-
this.addSub('ci:jest:test:finish', ({ status, testStartLine, promises, shouldRemoveProbe }) => {
328+
this.addSub('ci:jest:test:finish', ({ status, testStartLine }) => {
328329
const span = storage.getStore().span
329330
span.setTag(TEST_STATUS, status)
330331
if (testStartLine) {
@@ -346,10 +347,6 @@ class JestPlugin extends CiPlugin {
346347
span.finish()
347348
finishAllTraceSpans(span)
348349
this.activeTestSpan = null
349-
if (shouldRemoveProbe && this.runningTestProbeId) {
350-
promises.isProbeRemoved = withTimeout(this.removeDiProbe(this.runningTestProbeId), 2000)
351-
this.runningTestProbeId = null
352-
}
353350
})
354351

355352
this.addSub('ci:jest:test:err', ({ error, shouldSetProbe, promises }) => {
@@ -362,9 +359,7 @@ class JestPlugin extends CiPlugin {
362359
if (shouldSetProbe) {
363360
const probeInformation = this.addDiProbe(error)
364361
if (probeInformation) {
365-
const { probeId, setProbePromise, stackIndex } = probeInformation
366-
this.runningTestProbeId = probeId
367-
this.testErrorStackIndex = stackIndex
362+
const { setProbePromise } = probeInformation
368363
promises.isProbeReady = withTimeout(setProbePromise, 2000)
369364
}
370365
}

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

+5-5
Original file line numberDiff line numberDiff line change
@@ -219,9 +219,9 @@ class MochaPlugin extends CiPlugin {
219219
span.finish()
220220
finishAllTraceSpans(span)
221221
this.activeTestSpan = null
222-
if (this.di && this.libraryConfig?.isDiEnabled && this.runningTestProbeId && isLastRetry) {
223-
this.removeDiProbe(this.runningTestProbeId)
224-
this.runningTestProbeId = null
222+
if (this.di && this.libraryConfig?.isDiEnabled && this.runningTestProbe && isLastRetry) {
223+
this.removeDiProbe(this.runningTestProbe)
224+
this.runningTestProbe = null
225225
}
226226
}
227227
})
@@ -275,8 +275,8 @@ class MochaPlugin extends CiPlugin {
275275
if (isFirstAttempt && willBeRetried && this.di && this.libraryConfig?.isDiEnabled) {
276276
const probeInformation = this.addDiProbe(err)
277277
if (probeInformation) {
278-
const { probeId, stackIndex } = probeInformation
279-
this.runningTestProbeId = probeId
278+
const { file, line, stackIndex } = probeInformation
279+
this.runningTestProbe = { file, line }
280280
this.testErrorStackIndex = stackIndex
281281
test._ddShouldWaitForHitProbe = true
282282
// TODO: we're not waiting for setProbePromise to be resolved, so there might be race conditions

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

+4-4
Original file line numberDiff line numberDiff line change
@@ -135,8 +135,8 @@ class VitestPlugin extends CiPlugin {
135135
if (shouldSetProbe && this.di) {
136136
const probeInformation = this.addDiProbe(error)
137137
if (probeInformation) {
138-
const { probeId, stackIndex, setProbePromise } = probeInformation
139-
this.runningTestProbeId = probeId
138+
const { file, line, stackIndex, setProbePromise } = probeInformation
139+
this.runningTestProbe = { file, line }
140140
this.testErrorStackIndex = stackIndex
141141
promises.setProbePromise = setProbePromise
142142
}
@@ -237,8 +237,8 @@ class VitestPlugin extends CiPlugin {
237237
this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
238238
// TODO: too frequent flush - find for method in worker to decrease frequency
239239
this.tracer._exporter.flush(onFinish)
240-
if (this.runningTestProbeId) {
241-
this.removeDiProbe(this.runningTestProbeId)
240+
if (this.runningTestProbe) {
241+
this.removeDiProbe(this.runningTestProbe)
242242
}
243243
})
244244

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

+2
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,8 @@ class TestVisDynamicInstrumentation {
119119
const onHit = this.onHitBreakpointByProbeId.get(probeId)
120120
if (onHit) {
121121
onHit({ snapshot })
122+
} else {
123+
log.warn('Received a breakpoint hit for an unknown probe')
122124
}
123125
}).unref()
124126

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

+17-10
Original file line numberDiff line numberDiff line change
@@ -93,11 +93,14 @@ async function addBreakpoint (probe) {
9393
probe.location = { file, lines: [String(line)] }
9494

9595
const script = findScriptFromPartialPath(file)
96-
if (!script) throw new Error(`No loaded script found for ${file}`)
96+
if (!script) {
97+
log.error(`No loaded script found for ${file}`)
98+
throw new Error(`No loaded script found for ${file}`)
99+
}
97100

98101
const [path, scriptId, sourceMapURL] = script
99102

100-
log.debug(`Adding breakpoint at ${path}:${line}`)
103+
log.warn(`Adding breakpoint at ${path}:${line}`)
101104

102105
let lineNumber = line
103106

@@ -109,15 +112,19 @@ async function addBreakpoint (probe) {
109112
}
110113
}
111114

112-
const { breakpointId } = await session.post('Debugger.setBreakpoint', {
113-
location: {
114-
scriptId,
115-
lineNumber: lineNumber - 1
116-
}
117-
})
115+
try {
116+
const { breakpointId } = await session.post('Debugger.setBreakpoint', {
117+
location: {
118+
scriptId,
119+
lineNumber: lineNumber - 1
120+
}
121+
})
118122

119-
breakpointIdToProbe.set(breakpointId, probe)
120-
probeIdToBreakpointId.set(probe.id, breakpointId)
123+
breakpointIdToProbe.set(breakpointId, probe)
124+
probeIdToBreakpointId.set(probe.id, breakpointId)
125+
} catch (e) {
126+
log.error(`Error setting breakpoint at ${path}:${line}:`, e)
127+
}
121128
}
122129

123130
function start () {

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

+35-1
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ module.exports = class CiPlugin extends Plugin {
4545
constructor (...args) {
4646
super(...args)
4747

48+
this.fileLineToProbeId = new Map()
4849
this.rootDir = process.cwd() // fallback in case :session:start events are not emitted
4950

5051
this.addSub(`ci:${this.constructor.id}:library-configuration`, ({ onDone }) => {
@@ -335,7 +336,22 @@ module.exports = class CiPlugin extends Plugin {
335336
})
336337
}
337338

338-
removeDiProbe (probeId) {
339+
removeAllDiProbes () {
340+
if (this.fileLineToProbeId.size === 0) {
341+
return Promise.resolve()
342+
}
343+
log.debug('Removing all Dynamic Instrumentation probes')
344+
return Promise.all(Array.from(this.fileLineToProbeId.keys())
345+
.map((fileLine) => {
346+
const [file, line] = fileLine.split(':')
347+
return this.removeDiProbe({ file, line })
348+
}))
349+
}
350+
351+
removeDiProbe ({ file, line }) {
352+
const probeId = this.fileLineToProbeId.get(`${file}:${line}`)
353+
log.warn(`Removing probe from ${file}:${line}, with id: ${probeId}`)
354+
this.fileLineToProbeId.delete(probeId)
339355
return this.di.removeProbe(probeId)
340356
}
341357

@@ -346,9 +362,27 @@ module.exports = class CiPlugin extends Plugin {
346362
log.warn('Could not add breakpoint for dynamic instrumentation')
347363
return
348364
}
365+
log.debug('Adding breakpoint for Dynamic Instrumentation')
366+
367+
this.testErrorStackIndex = stackIndex
368+
const activeProbeKey = `${file}:${line}`
369+
370+
if (this.fileLineToProbeId.has(activeProbeKey)) {
371+
log.warn('Probe already set for this line')
372+
const oldProbeId = this.fileLineToProbeId.get(activeProbeKey)
373+
return {
374+
probeId: oldProbeId,
375+
setProbePromise: Promise.resolve(),
376+
stackIndex,
377+
file,
378+
line
379+
}
380+
}
349381

350382
const [probeId, setProbePromise] = this.di.addLineProbe({ file, line }, this.onDiBreakpointHit.bind(this))
351383

384+
this.fileLineToProbeId.set(activeProbeKey, probeId)
385+
352386
return {
353387
probeId,
354388
setProbePromise,

0 commit comments

Comments
 (0)