Add a mapping from revision to branches while using only one exp show call (#3980)

* Add a mapping from revision to branches

* Fix most tests

* Add flag to failing test

* Reduce number of calls to git

* Rework call to get commit messages

* Update structure of code and test fixture

* Do not duplicate commits

* Fix indicator

---------

Co-authored-by: Matt Seddon <[email protected]>

Author: sroy3

extension/src/cli/git/reader.ts

@@ -49,15 +49,19 @@ export class GitReader extends GitCli {
     return !output
   }
 
-  public async getCommitMessages(cwd: string, sha: string): Promise<string> {
+  public async getCommitMessages(
+    cwd: string,
+    revision: string,
+    revisions: string
+  ): Promise<string> {
     const options = getOptions(
       cwd,
       Command.LOG,
-      sha,
+      revision,
       Flag.PRETTY_FORMAT_COMMIT_MESSAGE,
       Flag.SEPARATE_WITH_NULL,
       Flag.NUMBER,
-      '1'
+      revisions
     )
     try {
       return await this.executeProcess(options)

extension/src/data/index.ts

@@ -9,8 +9,14 @@ import { uniqueValues } from '../util/array'
 import { DeferredDisposable } from '../class/deferred'
 import { isSameOrChild } from '../fileSystem'
 
+export type ExperimentsOutput = {
+  expShow: ExpShowOutput
+  rowOrder: { branch: string; sha: string }[]
+  gitLog: string
+}
+
 export abstract class BaseData<
-  T extends { data: PlotsOutputOrError; revs: string[] } | ExpShowOutput
+  T extends { data: PlotsOutputOrError; revs: string[] } | ExperimentsOutput
 > extends DeferredDisposable {
   public readonly onDidUpdate: Event<T>
   public readonly onDidChangeDvcYaml: Event<void>

extension/src/experiments/data/index.ts

@@ -7,14 +7,14 @@ import {
 import { getRelativePattern } from '../../fileSystem/relativePattern'
 import { createFileSystemWatcher } from '../../fileSystem/watcher'
 import { AvailableCommands, InternalCommands } from '../../commands/internal'
-import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../../cli/dvc/contract'
-import { BaseData } from '../../data'
-import { DOT_DVC, ExperimentFlag } from '../../cli/dvc/constants'
-import { gitPath } from '../../cli/git/constants'
+import { ExpShowOutput } from '../../cli/dvc/contract'
+import { BaseData, ExperimentsOutput } from '../../data'
+import { Args, DOT_DVC, ExperimentFlag } from '../../cli/dvc/constants'
+import { COMMITS_SEPARATOR, gitPath } from '../../cli/git/constants'
 import { getGitPath } from '../../fileSystem'
 import { ExperimentsModel } from '../model'
 
-export class ExperimentsData extends BaseData<ExpShowOutput> {
+export class ExperimentsData extends BaseData<ExperimentsOutput> {
   private readonly experiments: ExperimentsModel
 
   constructor(
@@ -47,69 +47,80 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
     )
 
     void this.updateAvailableBranchesToSelect(allBranches)
-    const data: ExpShowOutput = []
 
-    const { branches, currentBranch } = await this.getBranchesToShowWithCurrent(
-      allBranches
-    )
+    const branches = await this.getBranchesToShow(allBranches)
+    let gitLog = ''
+    const rowOrder: { branch: string; sha: string }[] = []
+    const args: Args = []
 
-    await Promise.all(
-      branches.map(async branch => {
-        const branchFlags = [
-          ExperimentFlag.REV,
-          branch,
-          ExperimentFlag.NUM_COMMIT,
-          this.experiments.getNbOfCommitsToShow(branch).toString()
-        ]
-
-        const output = (await this.expShow(
-          branchFlags,
-          branch
-        )) as ExpShowOutput
-
-        if (branch !== currentBranch) {
-          const workspaceIndex = output.findIndex(
-            exp => exp.rev === EXPERIMENT_WORKSPACE_ID
-          )
-          output.splice(workspaceIndex, 1)
-        }
-        data.push(...output)
-      })
+    for (const branch of branches) {
+      gitLog = await this.collectGitLogAndOrder(gitLog, branch, rowOrder, args)
+    }
+
+    const expShow = await this.internalCommands.executeCommand<ExpShowOutput>(
+      AvailableCommands.EXP_SHOW,
+      this.dvcRoot,
+      ...args
     )
 
-    this.collectFiles(data)
+    this.collectFiles({ expShow })
 
-    return this.notifyChanged(data)
+    return this.notifyChanged({
+      expShow,
+      gitLog,
+      rowOrder
+    })
   }
 
-  protected collectFiles(data: ExpShowOutput) {
-    this.collectedFiles = collectFiles(data, this.collectedFiles)
+  protected collectFiles({ expShow }: { expShow: ExpShowOutput }) {
+    this.collectedFiles = collectFiles(expShow, this.collectedFiles)
   }
 
-  private async getBranchesToShowWithCurrent(allBranches: string[]) {
+  private async collectGitLogAndOrder(
+    gitLog: string,
+    branch: string,
+    rowOrder: { branch: string; sha: string }[],
+    args: Args
+  ) {
+    const nbOfCommitsToShow = this.experiments.getNbOfCommitsToShow(branch)
+
+    const branchGitLog = await this.internalCommands.executeCommand(
+      AvailableCommands.GIT_GET_COMMIT_MESSAGES,
+      this.dvcRoot,
+      branch,
+      String(nbOfCommitsToShow)
+    )
+    gitLog = [gitLog, branchGitLog].join(COMMITS_SEPARATOR)
+
+    for (const commit of branchGitLog.split(COMMITS_SEPARATOR)) {
+      const [sha] = commit.split('\n')
+      rowOrder.push({ branch, sha })
+      if (args.includes(sha)) {
+        continue
+      }
+      args.push(ExperimentFlag.REV, sha)
+    }
+    return gitLog
+  }
+
+  private async getBranchesToShow(allBranches: string[]) {
     const currentBranch = await this.internalCommands.executeCommand<string>(
       AvailableCommands.GIT_GET_CURRENT_BRANCH,
       this.dvcRoot
     )
 
     this.experiments.pruneBranchesToShow(allBranches)
 
-    const branches = this.experiments.getBranchesToShow()
+    const currentBranches = [
+      currentBranch,
+      ...this.experiments
+        .getBranchesToShow()
+        .filter(branch => branch !== currentBranch)
+    ]
 
-    if (!branches.includes(currentBranch)) {
-      branches.push(currentBranch)
-      this.experiments.setBranchesToShow(branches)
-    }
-    return { branches, currentBranch }
-  }
+    this.experiments.setBranchesToShow(currentBranches)
 
-  private async expShow(flags: (ExperimentFlag | string)[], branch?: string) {
-    const data = await this.internalCommands.executeCommand<ExpShowOutput>(
-      AvailableCommands.EXP_SHOW,
-      this.dvcRoot,
-      ...flags
-    )
-    return data.map(exp => ({ ...exp, branch }))
+    return currentBranches
   }
 
   private async updateAvailableBranchesToSelect(branches?: string[]) {

extension/src/experiments/index.ts

@@ -33,7 +33,7 @@ import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
-import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../cli/dvc/contract'
+import { ExperimentsOutput } from '../data'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Title } from '../vscode/title'
@@ -43,7 +43,6 @@ import { Toast } from '../vscode/toast'
 import { ConfigKey } from '../vscode/config'
 import { checkSignalFile, pollSignalFileForProcess } from '../fileSystem'
 import { DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants'
-import { COMMITS_SEPARATOR } from '../cli/git/constants'
 
 export const ExperimentsScale = {
   ...omit(ColumnType, 'TIMESTAMP'),
@@ -169,13 +168,12 @@ export class Experiments extends BaseRepository<TableData> {
     return this.data.managedUpdate()
   }
 
-  public async setState(data: ExpShowOutput) {
+  public async setState({ expShow, gitLog, rowOrder }: ExperimentsOutput) {
     const hadCheckpoints = this.hasCheckpoints()
     const dvcLiveOnly = await this.checkSignalFile()
-    const commitsOutput = await this.getCommitOutput(data)
     await Promise.all([
-      this.columns.transformAndSet(data),
-      this.experiments.transformAndSet(data, dvcLiveOnly, commitsOutput)
+      this.columns.transformAndSet(expShow),
+      this.experiments.transformAndSet(expShow, gitLog, dvcLiveOnly, rowOrder)
     ])
 
     if (hadCheckpoints !== this.hasCheckpoints()) {
@@ -538,25 +536,6 @@ export class Experiments extends BaseRepository<TableData> {
     return this.webviewMessages.sendWebviewMessage()
   }
 
-  private async getCommitOutput(data: ExpShowOutput | undefined) {
-    if (!data || data.length === 0) {
-      return
-    }
-    let output = ''
-    for (const commit of data) {
-      if (commit.rev === EXPERIMENT_WORKSPACE_ID) {
-        continue
-      }
-      output += await this.internalCommands.executeCommand(
-        AvailableCommands.GIT_GET_COMMIT_MESSAGES,
-        this.dvcRoot,
-        commit.rev
-      )
-      output += COMMITS_SEPARATOR
-    }
-    return output
-  }
-
   private createWebviewMessageHandler() {
     const webviewMessages = new WebviewMessages(
       this.dvcRoot,
@@ -596,7 +575,7 @@ export class Experiments extends BaseRepository<TableData> {
     }
 
     return await pickExperiment(
-      this.experiments.getUniqueList(),
+      this.experiments.getCombinedList(),
       this.getFirstThreeColumnOrder(),
       Title.SELECT_BASE_EXPERIMENT
     )

extension/src/experiments/model/collect.test.ts

@@ -1,91 +1,64 @@
 import { collectExperiments } from './collect'
-import { COMMITS_SEPARATOR } from '../../cli/git/constants'
 import { generateTestExpShowOutput } from '../../test/util/experiments'
+import { ExpShowOutput } from '../../cli/dvc/contract'
+
+const DEFAULT_DATA: [string, boolean] = ['', false]
 
 describe('collectExperiments', () => {
   it('should return an empty array if no commits are present', () => {
     const { commits } = collectExperiments(
       generateTestExpShowOutput({}),
-      false,
-      ''
+      ...DEFAULT_DATA
     )
     expect(commits).toStrictEqual([])
   })
 
-  const expShowWithTwoCommits = generateTestExpShowOutput(
-    {},
-    {
-      experiments: [{}, {}],
-      name: 'branchA',
-      rev: '61bed4ce8913eca7f73ca754d65bc5daad1520e2'
-    },
-    {
-      name: 'branchB',
-      rev: '351e42ace3cb6a3a853c65bef285e60748cc6341'
-    }
-  )
+  const expShowWithTwoCommits: [ExpShowOutput, string, boolean] = [
+    generateTestExpShowOutput(
+      {},
+      {
+        experiments: [{}, {}],
+        name: 'branchA',
+        rev: '61bed4ce8913eca7f73ca754d65bc5daad1520e2'
+      },
+      {
+        name: 'branchB',
+        rev: '351e42ace3cb6a3a853c65bef285e60748cc6341'
+      }
+    ),
+    ...DEFAULT_DATA
+  ]
 
   it('should define a workspace', () => {
-    const { workspace } = collectExperiments(expShowWithTwoCommits, false, '')
+    const { workspace } = collectExperiments(...expShowWithTwoCommits)
 
     expect(workspace).toBeDefined()
   })
 
   it('should find two branches from a repo with two branches', () => {
-    const { commits } = collectExperiments(expShowWithTwoCommits, false, '')
+    const { commits } = collectExperiments(...expShowWithTwoCommits)
 
     expect(commits.length).toStrictEqual(2)
   })
 
   it('should list commits in the same order as they are collected', () => {
-    const { commits } = collectExperiments(expShowWithTwoCommits, false, '')
+    const { commits } = collectExperiments(...expShowWithTwoCommits)
     const [commitA, commitB] = commits
 
     expect(commitA.id).toStrictEqual('branchA')
     expect(commitB.id).toStrictEqual('branchB')
   })
 
   it('should find two experiments on commitA', () => {
-    const { experimentsByCommit } = collectExperiments(
-      expShowWithTwoCommits,
-      false,
-      ''
-    )
+    const { experimentsByCommit } = collectExperiments(...expShowWithTwoCommits)
     expect(experimentsByCommit.get('branchA')?.length).toStrictEqual(2)
   })
 
   it('should find no experiments on branchB', () => {
-    const { experimentsByCommit } = collectExperiments(
-      expShowWithTwoCommits,
-      false,
-      ''
-    )
+    const { experimentsByCommit } = collectExperiments(...expShowWithTwoCommits)
     expect(experimentsByCommit.get('branchB')).toBeUndefined()
   })
 
-  it('should add data from git to commits if git log output is provided', () => {
-    const { commits } = collectExperiments(
-      expShowWithTwoCommits,
-      false,
-      `61bed4ce8913eca7f73ca754d65bc5daad1520e2\nJohn Smith\n3 days ago\nrefNames:tag: v.1.1\nmessage:add new feature${COMMITS_SEPARATOR}351e42ace3cb6a3a853c65bef285e60748cc6341\nrenovate[bot]\n5 weeks ago\nrefNames:\nmessage:update various dependencies\n* update dvc\n* update dvclive`
-    )
-    const [branch1, branch2] = commits
-    expect(branch1.description).toStrictEqual('add new feature')
-    expect(branch2.description).toStrictEqual('update various dependencies ...')
-    expect(branch1.commit).toStrictEqual({
-      author: 'John Smith',
-      date: '3 days ago',
-      message: 'add new feature',
-      tags: ['v.1.1']
-    })
-    expect(branch2.commit).toStrictEqual({
-      author: 'renovate[bot]',
-      date: '5 weeks ago',
-      message: 'update various dependencies\n* update dvc\n* update dvclive',
-      tags: []
-    })
-  })
-
   it('should not collect the same experiment twice', () => {
     const main = {
       experiments: [
@@ -114,8 +87,7 @@ describe('collectExperiments', () => {
 
     const { experimentsByCommit, commits } = collectExperiments(
       expShowWithDuplicateCommits,
-      false,
-      ''
+      ...DEFAULT_DATA
     )
 
     expect(commits.length).toStrictEqual(3)