Use process polling to ensure DVCLive only PIDs are still running (#3045)

Author: mattseddon

extension/src/experiments/index.ts

@@ -41,7 +41,7 @@ import { createTypedAccumulator } from '../util/object'
 import { pickPaths } from '../path/selection/quickPick'
 import { Toast } from '../vscode/toast'
 import { ConfigKey } from '../vscode/config'
-import { checkSignalFile } from '../fileSystem'
+import { checkSignalFile, pollSignalFileForProcess } from '../fileSystem'
 import { DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants'
 
 export const ExperimentsScale = {
@@ -94,6 +94,9 @@ export class Experiments extends BaseRepository<TableData> {
   private readonly internalCommands: InternalCommands
   private readonly webviewMessages: WebviewMessages
 
+  private dvcLiveOnlyCleanupInitialized = false
+  private dvcLiveOnlySignalFile: string
+
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
@@ -105,6 +108,11 @@ export class Experiments extends BaseRepository<TableData> {
   ) {
     super(dvcRoot, resourceLocator.beaker)
 
+    this.dvcLiveOnlySignalFile = join(
+      this.dvcRoot,
+      DVCLIVE_ONLY_RUNNING_SIGNAL_FILE
+    )
+
     this.internalCommands = internalCommands
 
     this.onDidChangeIsParamsFileFocused = this.paramsFileFocused.event
@@ -164,9 +172,8 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public async setState(data: ExperimentsOutput) {
-    const dvcLiveOnly = await checkSignalFile(
-      join(this.dvcRoot, DVCLIVE_ONLY_RUNNING_SIGNAL_FILE)
-    )
+    const dvcLiveOnly = await this.checkSignalFile()
+
     const commitMessages: { [sha: string]: string } =
       await this.internalCommands.executeCommand(
         AvailableCommands.GIT_GET_LAST_THREE_COMMIT_MESSAGES,
@@ -602,4 +609,19 @@ export class Experiments extends BaseRepository<TableData> {
       this.onDidChangeColumns
     )
   }
+
+  private async checkSignalFile() {
+    const dvcLiveOnly = await checkSignalFile(this.dvcLiveOnlySignalFile)
+
+    if (dvcLiveOnly && !this.dvcLiveOnlyCleanupInitialized) {
+      this.dvcLiveOnlyCleanupInitialized = true
+      pollSignalFileForProcess(this.dvcLiveOnlySignalFile, () => {
+        this.dvcLiveOnlyCleanupInitialized = false
+        if (this.hasRunningExperiment()) {
+          this.cliData.update()
+        }
+      })
+    }
+    return dvcLiveOnly
+  }
 }

extension/src/fileSystem/index.ts

@@ -18,6 +18,7 @@ import { createValidInteger } from '../util/number'
 import { processExists } from '../processExecution'
 import { getFirstWorkspaceFolder } from '../vscode/workspaceFolders'
 import { DOT_DVC } from '../cli/dvc/constants'
+import { delay } from '../util/time'
 
 export const exists = (path: string): boolean => existsSync(path)
 
@@ -176,6 +177,19 @@ export const checkSignalFile = async (path: string): Promise<boolean> => {
   return !!(await getPidFromSignalFile(path))
 }
 
+export const pollSignalFileForProcess = async (
+  path: string,
+  callback: () => void,
+  ms = 5000
+): Promise<void> => {
+  await delay(ms)
+  const signalIsValid = await checkSignalFile(path)
+  if (signalIsValid) {
+    return pollSignalFileForProcess(path, callback, ms)
+  }
+  return callback()
+}
+
 export const getBinDisplayText = (
   path: string | undefined
 ): string | undefined => {

extension/src/test/suite/experiments/index.test.ts

@@ -71,6 +71,7 @@ import { PlotSizeNumber } from '../../../plots/webview/contract'
 import { RegisteredCommands } from '../../../commands/external'
 import { ConfigKey } from '../../../vscode/config'
 import { EXPERIMENT_WORKSPACE_ID } from '../../../cli/dvc/contract'
+import * as Time from '../../../util/time'
 
 suite('Experiments Test Suite', () => {
   const disposable = Disposable.fn()
@@ -1626,4 +1627,55 @@ suite('Experiments Test Suite', () => {
       ).to.deep.equal([])
     })
   })
+
+  describe('setState', () => {
+    it('should clean up after a killed DVCLive process that was running an experiment outside of the DVC context', async () => {
+      const defaultExperimentsData = { workspace: { baseline: { data: {} } } }
+
+      const { experiments, mockCheckSignalFile, mockUpdateExperimentsData } =
+        buildExperiments(disposable, defaultExperimentsData)
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const getCleanupInitialized = (experiments: any) =>
+        experiments.dvcLiveOnlyCleanupInitialized
+
+      await experiments.isReady()
+
+      const mockDelay = stub(Time, 'delay')
+      mockDelay.callsFake(() => mockDelay.wrappedMethod(500))
+
+      let processKilled = false
+
+      const cleanupUpdate = new Promise(resolve =>
+        mockUpdateExperimentsData.callsFake(() => resolve(undefined))
+      )
+
+      mockCheckSignalFile.resetBehavior()
+      mockCheckSignalFile.callsFake(() => {
+        return Promise.resolve(!processKilled)
+      })
+
+      const dataUpdated = new Promise(resolve =>
+        disposable.track(
+          experiments.onDidChangeExperiments(() => resolve(undefined))
+        )
+      )
+
+      experiments.setState(defaultExperimentsData)
+      await dataUpdated
+      expect(experiments.hasRunningExperiment()).to.be.true
+      expect(getCleanupInitialized(experiments)).to.be.true
+
+      processKilled = true
+
+      mockUpdateExperimentsData.resetHistory()
+
+      await cleanupUpdate
+
+      expect(getCleanupInitialized(experiments)).to.be.false
+      expect(mockCheckSignalFile).to.be.called
+      expect(mockDelay).to.be.called
+      expect(mockUpdateExperimentsData).to.be.calledOnce
+    })
+  })
 })

extension/src/test/suite/experiments/util.ts

@@ -60,14 +60,19 @@ export const buildExperiments = (
     resourceLocator
   } = buildDependencies(disposer, experimentShowData)
 
+  const mockUpdateExperimentsData = stub()
+  const mockExperimentsData = buildMockData<ExperimentsData>(
+    mockUpdateExperimentsData
+  )
+
   const experiments = disposer.track(
     new Experiments(
       dvcRoot,
       internalCommands,
       updatesPaused,
       resourceLocator,
       buildMockMemento(),
-      buildMockData<ExperimentsData>(),
+      mockExperimentsData,
       buildMockData<FileSystemData>()
     )
   )
@@ -90,6 +95,7 @@ export const buildExperiments = (
     messageSpy,
     mockCheckSignalFile,
     mockExperimentShow,
+    mockUpdateExperimentsData,
     resourceLocator,
     updatesPaused
   }

extension/src/test/suite/util.ts

@@ -164,10 +164,13 @@ export const buildInternalCommands = (disposer: Disposer) => {
   }
 }
 
-export const buildMockData = <T extends ExperimentsData | FileSystemData>() =>
+export const buildMockData = <T extends ExperimentsData | FileSystemData>(
+  update = stub()
+) =>
   ({
     dispose: stub(),
-    onDidUpdate: stub()
+    onDidUpdate: stub(),
+    update
   } as unknown as T)
 
 export const buildDependencies = (