Do not block data from going to webview after fetching once (#1727)

Author: mattseddon

extension/src/data/index.ts

@@ -9,7 +9,7 @@ import { definedAndNonEmpty, sameContents, uniqueValues } from '../util/array'
 import { DeferredDisposable } from '../class/deferred'
 
 export abstract class BaseData<
-  T extends PlotsOutput | ExperimentsOutput
+  T extends { data: PlotsOutput; revs: string[] } | ExperimentsOutput
 > extends DeferredDisposable {
   public readonly onDidUpdate: Event<T>
 

extension/src/plots/data/index.ts

@@ -9,7 +9,7 @@ import {
 } from '../../util/array'
 import { PlotsModel } from '../model'
 
-export class PlotsData extends BaseData<PlotsOutput> {
+export class PlotsData extends BaseData<{ data: PlotsOutput; revs: string[] }> {
   private model?: PlotsModel
 
   constructor(
@@ -26,7 +26,7 @@ export class PlotsData extends BaseData<PlotsOutput> {
   }
 
   public async update(): Promise<void> {
-    const revisions = flattenUnique([
+    const revs = flattenUnique([
       this.model?.getMissingRevisions() || [],
       this.model?.getMutableRevisions() || []
     ])
@@ -35,31 +35,31 @@ export class PlotsData extends BaseData<PlotsOutput> {
       (await this.internalCommands.executeCommand<boolean>(
         AvailableCommands.IS_EXPERIMENT_RUNNING
       )) &&
-      !definedAndNonEmpty(revisions)
+      !definedAndNonEmpty(revs)
     ) {
       return
     }
 
-    const args = sameContents(revisions, ['workspace']) ? [] : revisions
+    const args = sameContents(revs, ['workspace']) ? [] : revs
 
     const data = await this.internalCommands.executeCommand<PlotsOutput>(
       AvailableCommands.PLOTS_DIFF,
       this.dvcRoot,
       ...args
     )
 
-    const files = this.collectFiles(data)
+    const files = this.collectFiles({ data })
 
     this.compareFiles(files)
 
-    return this.notifyChanged(data)
+    return this.notifyChanged({ data, revs })
   }
 
   public managedUpdate() {
     return this.processManager.run('update')
   }
 
-  public collectFiles(data: PlotsOutput) {
+  public collectFiles({ data }: { data: PlotsOutput }) {
     return Object.keys(data)
   }
 

extension/src/plots/index.ts

@@ -63,9 +63,9 @@ export class Plots extends BaseRepository<TPlotsData> {
     )
 
     this.dispose.track(
-      this.data.onDidUpdate(async data => {
+      this.data.onDidUpdate(async ({ data, revs }) => {
         await Promise.all([
-          this.plots?.transformAndSetPlots(data),
+          this.plots?.transformAndSetPlots(data, revs),
           this.paths?.transformAndSet(data)
         ])
         this.notifyChanged()
@@ -156,7 +156,7 @@ export class Plots extends BaseRepository<TPlotsData> {
 
     if (
       this.paths?.hasPaths() &&
-      definedAndNonEmpty(this.plots?.getMissingRevisions())
+      definedAndNonEmpty(this.plots?.getUnfetchedRevisions())
     ) {
       this.sendCheckpointPlotsData()
       return this.data.managedUpdate()

extension/src/plots/model/index.ts

@@ -41,6 +41,8 @@ export class PlotsModel extends ModelWithPersistence {
   private branchRevisions: Record<string, string> = {}
   private workspaceRunningCheckpoint: string | undefined
 
+  private fetchedRevs = new Set<string>()
+
   private comparisonData: ComparisonData = {}
   private comparisonOrder: string[]
 
@@ -97,7 +99,9 @@ export class PlotsModel extends ModelWithPersistence {
     return this.removeStaleData()
   }
 
-  public async transformAndSetPlots(data: PlotsOutput) {
+  public async transformAndSetPlots(data: PlotsOutput, revs: string[]) {
+    this.fetchedRevs = new Set([...this.fetchedRevs, ...revs])
+
     const [{ comparisonData, revisionData }, templates] = await Promise.all([
       collectData(data),
       collectTemplates(data)
@@ -153,6 +157,12 @@ export class PlotsModel extends ModelWithPersistence {
     }
   }
 
+  public getUnfetchedRevisions() {
+    return this.getSelectedRevisions().filter(
+      revision => !this.fetchedRevs.has(revision)
+    )
+  }
+
   public getMissingRevisions() {
     const cachedRevisions = new Set([
       ...Object.keys(this.comparisonData),
@@ -306,6 +316,7 @@ export class PlotsModel extends ModelWithPersistence {
       if (sha && this.branchRevisions[id] !== sha) {
         delete this.revisionData[id]
         delete this.comparisonData[id]
+        this.fetchedRevs.delete(id)
         this.branchRevisions[id] = sha
       }
     }