fix(test-optimization): complete test runs after subscriber loss (#9381)

* fix(test-optimization): complete test runs after subscriber loss

Finish subscribers can disable themselves while handling a diagnostic-channel publication, leaving no callback owner and hanging the framework run indefinitely. Treat losing the last subscriber as completion while preserving exporter-backed waits and Jest's timeout.

* test(test-optimization): cover vitest suite finish and cucumber config-loss

The vitest worker suite-finish path (the one that awaits the flush through the
shared channel helper) had no dedicated instrumentation spec and no CI job, so
its subscriber-loss behavior rode only on end-to-end coverage. Add a real-path
spec driving the `@vitest/runner` `startTests` hook, plus `instrumentation-vitest`
and `instrumentation-cucumber` jobs alongside the other frameworks.

The cucumber spec now pins the `|| {}` guard in getWrappedStart: a
library-configuration subscriber that disables itself without responding resolves
the request to `undefined`, and the run must still finish with every remote
feature treated as disabled.

Drive-by fix:

* Drop `frameworkVersion` from the vitest worker `testSuiteFinishCh` payload; the
  suite-finish subscriber never reads it and it is already carried in currentStore.

* fix(test-optimization): finish runs after plugin lifecycle errors

Mocha configuration requests publish through runStores, so an internal handler error can disable the plugin and leave delayed startup waiting forever. Re-enabling Playwright after test start can leave an active finish subscriber without a span, whose early return similarly never releases the worker.

Author: BridgeAR

integration-tests/ci-visibility/cucumber-programmatic/features/pass.feature

@@ -0,0 +1,3 @@
+Feature: Programmatic Cucumber run
+  Scenario: Pass
+    Then the scenario passes

integration-tests/ci-visibility/cucumber-programmatic/features/support/steps.js

@@ -0,0 +1,9 @@
+'use strict'
+
+const assert = require('node:assert/strict')
+
+const { Then } = require('@cucumber/cucumber')
+
+Then('the scenario passes', function () {
+  assert.ok(true)
+})

integration-tests/ci-visibility/playwright-plugin-lifecycle/plugin-lifecycle-test.js

@@ -0,0 +1,12 @@
+'use strict'
+
+const tracer = require('dd-trace')
+const { expect, test } = require('@playwright/test')
+
+tracer.use('playwright', false)
+
+test('finishes after the plugin is re-enabled during the test', () => {
+  tracer.use('playwright', true)
+
+  expect(1 + 2).toBe(3)
+})

integration-tests/ci-visibility/run-cucumber-programmatic.js

@@ -0,0 +1,24 @@
+'use strict'
+
+const { loadConfiguration, runCucumber } = require('@cucumber/cucumber/api')
+
+const completedMessage = 'programmatic Cucumber run completed'
+
+async function main () {
+  const { runConfiguration } = await loadConfiguration({
+    file: false,
+    provided: {
+      paths: ['ci-visibility/cucumber-programmatic/features/pass.feature'],
+      require: ['ci-visibility/cucumber-programmatic/features/support/steps.js'],
+    },
+  })
+  const { success } = await runCucumber(runConfiguration)
+
+  process.stdout.write(`${completedMessage}\n`)
+  if (!success) process.exitCode = 1
+}
+
+main().catch((error) => {
+  process.stderr.write(`${error.stack || error}\n`)
+  process.exitCode = 1
+})

integration-tests/ci-visibility/run-mocha.js

@@ -21,6 +21,9 @@ mocha.run((failures) => {
   if (process.send) {
     process.send('finished')
   }
+  if (process.env.REPORT_RUN_CALLBACK) {
+    process.stdout.write('programmatic Mocha run completed\n')
+  }
   if (process.env.SHOULD_CHECK_RESULTS && failures > 0) {
     process.exit(1)
   }

integration-tests/ci-visibility/test-optimization-wrong-init.spec.js

@@ -76,7 +76,7 @@ testFrameworks.forEach(({ testFramework, command, expectedOutput, extraTestConte
       await receiver.stop()
     })
 
-    it('does not initialize test optimization plugins if Test Optimization mode is not enabled', async () => {
+    it('finishes without initializing test optimization plugins if Test Optimization mode is not enabled', async () => {
       const eventsPromise = receiver
         .gatherPayloadsMaxTimeout(({ url }) => url === '/v0.4/traces', (tracesRequests) => {
           const spans = tracesRequests.flatMap(trace => trace.payload).flatMap(request => request)
@@ -121,11 +121,12 @@ testFrameworks.forEach(({ testFramework, command, expectedOutput, extraTestConte
         processOutput += chunk.toString()
       })
 
-      await Promise.all([
+      const [[exitCode]] = await Promise.all([
         once(childProcess, 'exit'),
         eventsPromise,
       ])
 
+      assert.strictEqual(exitCode, 0, processOutput)
       const reason = 'is not initialized because Test Optimization mode is not enabled.'
       const expectedSubstring = `Plugin "${testFramework}" ${reason}`
       assert.ok(processOutput.includes(expectedSubstring), `Got: ${inspect(processOutput)}`)

integration-tests/cucumber/cucumber.spec.js

@@ -157,6 +157,44 @@ describe(`cucumber@${version} commonJS`, () => {
     ])
   })
 
+  onlyLatestIt('waits for the final payload before the programmatic run resolves', async () => {
+    const completionOrder = []
+    const completedMessage = 'programmatic Cucumber run completed'
+    receiver.setWaitingTime(500)
+
+    const intakePromise = (async () => {
+      await receiver.payloadReceived(({ url, payload }) => (
+        url === '/api/v2/citestcycle' &&
+        payload.events.some(event => event.type === 'test_session_end')
+      ))
+      completionOrder.push('intake')
+    })()
+
+    childProcess = exec(
+      'node ./ci-visibility/run-cucumber-programmatic.js',
+      {
+        cwd,
+        env: getCiVisAgentlessConfig(receiver.port),
+      }
+    )
+    childProcess.stdout?.on('data', (chunk) => {
+      const output = chunk.toString()
+      testOutput += output
+      if (output.includes(completedMessage)) completionOrder.push('run')
+    })
+    childProcess.stderr?.on('data', (chunk) => {
+      testOutput += chunk.toString()
+    })
+
+    const [[exitCode]] = await Promise.all([
+      once(childProcess, 'exit'),
+      intakePromise,
+    ])
+
+    assert.strictEqual(exitCode, 0, testOutput)
+    assert.deepStrictEqual(completionOrder, ['intake', 'run'], testOutput)
+  })
+
   context('with APM protocol (old agents)', () => {
     it('can report tests', async function () {
       receiver.setInfoResponse({ endpoints: [] })

integration-tests/mocha/mocha.spec.js

@@ -312,6 +312,59 @@ describe(`mocha@${MOCHA_VERSION}`, function () {
     ])
   })
 
+  /**
+   * @param {boolean} runInParallel
+   * @returns {Promise<void>}
+   */
+  async function assertProgrammaticRunWaitsForFinalPayload (runInParallel) {
+    const completionOrder = []
+    const completedMessage = 'programmatic Mocha run completed'
+    receiver.setWaitingTime(500)
+
+    const intakePromise = (async () => {
+      await receiver.payloadReceived(({ url, payload }) => (
+        url === '/api/v2/citestcycle' &&
+        payload.events.some(event => event.type === 'test_session_end')
+      ))
+      completionOrder.push('intake')
+    })()
+
+    childProcess = exec(
+      runTestsCommand,
+      {
+        cwd,
+        env: {
+          ...getCiVisAgentlessConfig(receiver.port),
+          REPORT_RUN_CALLBACK: '1',
+          ...(runInParallel && { RUN_IN_PARALLEL: '1' }),
+        },
+      }
+    )
+    childProcess.stdout?.on('data', (chunk) => {
+      const output = chunk.toString()
+      testOutput += output
+      if (output.includes(completedMessage)) completionOrder.push('run')
+    })
+    childProcess.stderr?.on('data', (chunk) => {
+      testOutput += chunk.toString()
+    })
+
+    const [[exitCode]] = await Promise.all([
+      once(childProcess, 'exit'),
+      intakePromise,
+    ])
+
+    assert.strictEqual(exitCode, 0, testOutput)
+    assert.deepStrictEqual(completionOrder, ['intake', 'run'], testOutput)
+  }
+
+  for (const runInParallel of [false, true]) {
+    const mode = runInParallel ? 'parallel' : 'serial'
+    onlyLatestIt(`waits for the final payload before invoking the programmatic run callback (${mode})`, async () => {
+      await assertProgrammaticRunWaitsForFinalPayload(runInParallel)
+    })
+  }
+
   const nonLegacyReportingOptions = ['evp proxy', 'agentless']
 
   nonLegacyReportingOptions.forEach((reportingOption) => {

integration-tests/playwright/playwright-reporting.spec.js

@@ -814,6 +814,24 @@ versions.forEach((version) => {
       })
     })
 
+    if (version === latest) {
+      it('finishes if the plugin is re-enabled after test start', async (receiver, run) => {
+        const proc = run(
+          './node_modules/.bin/playwright test -c playwright.config.js',
+          {
+            cwd,
+            env: {
+              ...getCiVisAgentlessConfig(receiver.port),
+              TEST_DIR: './ci-visibility/playwright-plugin-lifecycle',
+            },
+          }
+        )
+
+        const [exitCode] = await once(proc, 'exit')
+        assert.strictEqual(exitCode, 0)
+      })
+    }
+
     const fullyParallelConfigValue = [true, false]
 
     fullyParallelConfigValue.forEach((parallelism) => {

packages/datadog-instrumentations/src/cucumber.js

@@ -28,6 +28,7 @@ const {
 } = require('../../dd-trace/src/plugins/util/test')
 const { writeCoverageBackfillToCache } = require('../../dd-trace/src/ci-visibility/test-optimization-cache')
 const satisfies = require('../../../vendor/dist/semifies')
+const { getChannelPromise } = require('./helpers/channel')
 const { addHook, channel } = require('./helpers/instrument')
 
 const cucumberWorkerThreadsPatchModule = require.resolve('./cucumber-worker-threads')
@@ -563,12 +564,6 @@ function getErrorFromCucumberResult (cucumberResult) {
   return error
 }
 
-function getChannelPromise (channelToPublishTo, frameworkVersion = null) {
-  return new Promise(resolve => {
-    channelToPublishTo.publish({ onDone: resolve, frameworkVersion })
-  })
-}
-
 function getShouldBeSkippedSuite (pickle, suitesToSkip) {
   const testSuitePath = getTestSuitePath(pickle.uri, process.cwd())
   const isUnskippable = isMarkedAsUnskippable(pickle)
@@ -1056,7 +1051,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     }
     let errorSkippableRequest
 
-    const configurationResponse = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
+    const configurationResponse = await getChannelPromise(libraryConfigurationCh, { frameworkVersion }) || {}
 
     repositoryRoot = configurationResponse.repositoryRoot
     isItrEnabled = configurationResponse.libraryConfig?.isItrEnabled
@@ -1089,7 +1084,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
 
     if (isKnownTestsEnabled) {
       const currentKnownTestsResponse = knownTestsResponse || await getChannelPromise(knownTestsCh)
-      if (currentKnownTestsResponse.err) {
+      if (!currentKnownTestsResponse || currentKnownTestsResponse.err) {
         isEarlyFlakeDetectionEnabled = false
         isKnownTestsEnabled = false
       } else {
@@ -1098,7 +1093,9 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     }
 
     if (isSuitesSkippingEnabled) {
-      const skippableResponse = skippableSuitesResponse || await getChannelPromise(skippableSuitesCh)
+      const skippableResponse = skippableSuitesResponse ||
+        await getChannelPromise(skippableSuitesCh) ||
+        { err: true }
 
       errorSkippableRequest = skippableResponse.err
       skippableSuites = skippableResponse.skippableSuites ?? []
@@ -1147,7 +1144,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
     if (isTestManagementTestsEnabled) {
       const currentTestManagementTestsResponse =
         testManagementTestsResponse || await getChannelPromise(testManagementTestsCh)
-      if (currentTestManagementTestsResponse.err) {
+      if (!currentTestManagementTestsResponse || currentTestManagementTestsResponse.err) {
         isTestManagementTestsEnabled = false
       } else {
         testManagementTests = currentTestManagementTestsResponse.testManagementTests
@@ -1156,7 +1153,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
 
     if (isImpactedTestsEnabled) {
       const impactedTestsResponse = await getChannelPromise(modifiedFilesCh)
-      if (!impactedTestsResponse.err) {
+      if (impactedTestsResponse && !impactedTestsResponse.err) {
         modifiedFiles = impactedTestsResponse.modifiedFiles
       }
     }
@@ -1215,7 +1212,7 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
       global.__coverage__ = fromCoverageMapToCoverage(originalCoverageMap)
     }
 
-    sessionFinishCh.publish({
+    const flushPromise = getChannelPromise(sessionFinishCh, {
       status: success ? 'pass' : 'fail',
       isSuitesSkipped,
       testCodeCoverageLinesTotal,
@@ -1228,9 +1225,11 @@ function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordin
       isTestManagementTestsEnabled,
       isParallel,
     })
+
     logTestOptimizationSummary({ attemptToFixExecutions })
     loggedAttemptToFixTests.clear()
     eventDataCollector = null
+    await flushPromise
     return result
   }
 }