From ed65a7dbca435e4a80b54a5b5701122135da417e Mon Sep 17 00:00:00 2001
From: arjunajesh <34989598+arjunajesh@users.noreply.github.com>
Date: Sat, 15 Jul 2023 22:48:30 -0400
Subject: [PATCH] implement progress bar for index loading

---
 cmd/restic/cmd_backup.go                  |  5 ++++-
 cmd/restic/cmd_backup_integration_test.go |  2 +-
 cmd/restic/cmd_cat.go                     |  6 ++++--
 cmd/restic/cmd_copy.go                    |  7 ++++---
 cmd/restic/cmd_debug.go                   |  3 ++-
 cmd/restic/cmd_diff.go                    |  4 ++--
 cmd/restic/cmd_dump.go                    |  3 ++-
 cmd/restic/cmd_find.go                    |  4 ++--
 cmd/restic/cmd_ls.go                      |  3 ++-
 cmd/restic/cmd_mount.go                   |  3 ++-
 cmd/restic/cmd_prune.go                   |  3 ++-
 cmd/restic/cmd_recover.go                 |  5 +++--
 cmd/restic/cmd_repair_snapshots.go        |  3 ++-
 cmd/restic/cmd_restore.go                 |  3 ++-
 cmd/restic/cmd_rewrite.go                 |  3 ++-
 cmd/restic/cmd_stats.go                   |  4 ++--
 cmd/restic/integration_helpers_test.go    |  2 +-
 cmd/restic/progress.go                    | 23 ++++++++++++++++++-----
 internal/fuse/snapshots_dirstruct.go      |  2 +-
 internal/index/master_index_test.go       |  2 +-
 internal/repository/repack_test.go        |  2 +-
 internal/repository/repository.go         | 19 +++++++++++++++++--
 internal/repository/repository_test.go    |  4 ++--
 internal/restic/repository.go             |  2 +-
 24 files changed, 80 insertions(+), 37 deletions(-)

diff --git a/cmd/restic/cmd_backup.go b/cmd/restic/cmd_backup.go
index 6b5706855..09d5191b5 100644
--- a/cmd/restic/cmd_backup.go
+++ b/cmd/restic/cmd_backup.go
@@ -546,7 +546,10 @@ func runBackup(ctx context.Context, opts BackupOptions, gopts GlobalOptions, ter
 	if !gopts.JSON {
 		progressPrinter.V("load index files")
 	}
-	err = repo.LoadIndex(ctx)
+
+	bar := newTerminalProgressMax(!gopts.Quiet, 0, "index files loaded", term)
+
+	err = repo.LoadIndex(ctx, bar)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/restic/cmd_backup_integration_test.go b/cmd/restic/cmd_backup_integration_test.go
index fb7bef633..76c227e3d 100644
--- a/cmd/restic/cmd_backup_integration_test.go
+++ b/cmd/restic/cmd_backup_integration_test.go
@@ -252,7 +252,7 @@ func TestBackupTreeLoadError(t *testing.T) {
 
 	r, err := OpenRepository(context.TODO(), env.gopts)
 	rtest.OK(t, err)
-	rtest.OK(t, r.LoadIndex(context.TODO()))
+	rtest.OK(t, r.LoadIndex(context.TODO(), nil))
 	treePacks := restic.NewIDSet()
 	r.Index().Each(context.TODO(), func(pb restic.PackedBlob) {
 		if pb.Type == restic.TreeBlob {
diff --git a/cmd/restic/cmd_cat.go b/cmd/restic/cmd_cat.go
index 97789d271..fe4314179 100644
--- a/cmd/restic/cmd_cat.go
+++ b/cmd/restic/cmd_cat.go
@@ -169,7 +169,8 @@ func runCat(ctx context.Context, gopts GlobalOptions, args []string) error {
 		return err
 
 	case "blob":
-		err = repo.LoadIndex(ctx)
+		bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+		err = repo.LoadIndex(ctx, bar)
 		if err != nil {
 			return err
 		}
@@ -197,7 +198,8 @@ func runCat(ctx context.Context, gopts GlobalOptions, args []string) error {
 			return errors.Fatalf("could not find snapshot: %v\n", err)
 		}
 
-		err = repo.LoadIndex(ctx)
+		bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+		err = repo.LoadIndex(ctx, bar)
 		if err != nil {
 			return err
 		}
diff --git a/cmd/restic/cmd_copy.go b/cmd/restic/cmd_copy.go
index eaa0ef81a..32d21e50b 100644
--- a/cmd/restic/cmd_copy.go
+++ b/cmd/restic/cmd_copy.go
@@ -99,12 +99,13 @@ func runCopy(ctx context.Context, opts CopyOptions, gopts GlobalOptions, args []
 	}
 
 	debug.Log("Loading source index")
-	if err := srcRepo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err := srcRepo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
-
+	bar = newProgressMax(!gopts.Quiet, 0, "index files loaded")
 	debug.Log("Loading destination index")
-	if err := dstRepo.LoadIndex(ctx); err != nil {
+	if err := dstRepo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
diff --git a/cmd/restic/cmd_debug.go b/cmd/restic/cmd_debug.go
index a54200c45..c7aa82449 100644
--- a/cmd/restic/cmd_debug.go
+++ b/cmd/restic/cmd_debug.go
@@ -469,7 +469,8 @@ func runDebugExamine(ctx context.Context, gopts GlobalOptions, args []string) er
 		}
 	}
 
-	err = repo.LoadIndex(ctx)
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	err = repo.LoadIndex(ctx, bar)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/restic/cmd_diff.go b/cmd/restic/cmd_diff.go
index accd55db2..a82700358 100644
--- a/cmd/restic/cmd_diff.go
+++ b/cmd/restic/cmd_diff.go
@@ -363,8 +363,8 @@ func runDiff(ctx context.Context, opts DiffOptions, gopts GlobalOptions, args []
 	if !gopts.JSON {
 		Verbosef("comparing snapshot %v to %v:\n\n", sn1.ID().Str(), sn2.ID().Str())
 	}
-
-	if err = repo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err = repo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
diff --git a/cmd/restic/cmd_dump.go b/cmd/restic/cmd_dump.go
index 325607017..970d323bb 100644
--- a/cmd/restic/cmd_dump.go
+++ b/cmd/restic/cmd_dump.go
@@ -152,7 +152,8 @@ func runDump(ctx context.Context, opts DumpOptions, gopts GlobalOptions, args []
 		return errors.Fatalf("failed to find snapshot: %v", err)
 	}
 
-	err = repo.LoadIndex(ctx)
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	err = repo.LoadIndex(ctx, bar)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/restic/cmd_find.go b/cmd/restic/cmd_find.go
index 181d8595d..d1c6e6628 100644
--- a/cmd/restic/cmd_find.go
+++ b/cmd/restic/cmd_find.go
@@ -589,8 +589,8 @@ func runFind(ctx context.Context, opts FindOptions, gopts GlobalOptions, args []
 	if err != nil {
 		return err
 	}
-
-	if err = repo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err = repo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
diff --git a/cmd/restic/cmd_ls.go b/cmd/restic/cmd_ls.go
index 256c9e002..9f1b64d60 100644
--- a/cmd/restic/cmd_ls.go
+++ b/cmd/restic/cmd_ls.go
@@ -173,7 +173,8 @@ func runLs(ctx context.Context, opts LsOptions, gopts GlobalOptions, args []stri
 		return err
 	}
 
-	if err = repo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err = repo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
diff --git a/cmd/restic/cmd_mount.go b/cmd/restic/cmd_mount.go
index ec3662d5c..39bcd6729 100644
--- a/cmd/restic/cmd_mount.go
+++ b/cmd/restic/cmd_mount.go
@@ -130,7 +130,8 @@ func runMount(ctx context.Context, opts MountOptions, gopts GlobalOptions, args
 		}
 	}
 
-	err = repo.LoadIndex(ctx)
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	err = repo.LoadIndex(ctx, bar)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/restic/cmd_prune.go b/cmd/restic/cmd_prune.go
index e4c2c7b29..3b9590595 100644
--- a/cmd/restic/cmd_prune.go
+++ b/cmd/restic/cmd_prune.go
@@ -187,7 +187,8 @@ func runPruneWithRepo(ctx context.Context, opts PruneOptions, gopts GlobalOption
 
 	Verbosef("loading indexes...\n")
 	// loading the index before the snapshots is ok, as we use an exclusive lock here
-	err := repo.LoadIndex(ctx)
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	err := repo.LoadIndex(ctx, bar)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/restic/cmd_recover.go b/cmd/restic/cmd_recover.go
index 85dcc23d7..b8035d102 100644
--- a/cmd/restic/cmd_recover.go
+++ b/cmd/restic/cmd_recover.go
@@ -58,7 +58,8 @@ func runRecover(ctx context.Context, gopts GlobalOptions) error {
 	}
 
 	Verbosef("load index files\n")
-	if err = repo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err = repo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
@@ -73,7 +74,7 @@ func runRecover(ctx context.Context, gopts GlobalOptions) error {
 	})
 
 	Verbosef("load %d trees\n", len(trees))
-	bar := newProgressMax(!gopts.Quiet, uint64(len(trees)), "trees loaded")
+	bar = newProgressMax(!gopts.Quiet, uint64(len(trees)), "trees loaded")
 	for id := range trees {
 		tree, err := restic.LoadTree(ctx, repo, id)
 		if err != nil {
diff --git a/cmd/restic/cmd_repair_snapshots.go b/cmd/restic/cmd_repair_snapshots.go
index 03736795c..a7d3af4a8 100644
--- a/cmd/restic/cmd_repair_snapshots.go
+++ b/cmd/restic/cmd_repair_snapshots.go
@@ -89,7 +89,8 @@ func runRepairSnapshots(ctx context.Context, gopts GlobalOptions, opts RepairOpt
 		return err
 	}
 
-	if err := repo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err := repo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
diff --git a/cmd/restic/cmd_restore.go b/cmd/restic/cmd_restore.go
index 9bce90a26..11ce5e4a5 100644
--- a/cmd/restic/cmd_restore.go
+++ b/cmd/restic/cmd_restore.go
@@ -173,7 +173,8 @@ func runRestore(ctx context.Context, opts RestoreOptions, gopts GlobalOptions,
 		return errors.Fatalf("failed to find snapshot: %v", err)
 	}
 
-	err = repo.LoadIndex(ctx)
+	bar := newTerminalProgressMax(!gopts.Quiet, 0, "index files loaded", term)
+	err = repo.LoadIndex(ctx, bar)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/restic/cmd_rewrite.go b/cmd/restic/cmd_rewrite.go
index c08797c48..008c133d1 100644
--- a/cmd/restic/cmd_rewrite.go
+++ b/cmd/restic/cmd_rewrite.go
@@ -212,7 +212,8 @@ func runRewrite(ctx context.Context, opts RewriteOptions, gopts GlobalOptions, a
 		return err
 	}
 
-	if err = repo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err = repo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
diff --git a/cmd/restic/cmd_stats.go b/cmd/restic/cmd_stats.go
index a7ecd438f..900000439 100644
--- a/cmd/restic/cmd_stats.go
+++ b/cmd/restic/cmd_stats.go
@@ -98,8 +98,8 @@ func runStats(ctx context.Context, opts StatsOptions, gopts GlobalOptions, args
 	if err != nil {
 		return err
 	}
-
-	if err = repo.LoadIndex(ctx); err != nil {
+	bar := newProgressMax(!gopts.Quiet, 0, "index files loaded")
+	if err = repo.LoadIndex(ctx, bar); err != nil {
 		return err
 	}
 
diff --git a/cmd/restic/integration_helpers_test.go b/cmd/restic/integration_helpers_test.go
index b7cb5b333..2afdbb938 100644
--- a/cmd/restic/integration_helpers_test.go
+++ b/cmd/restic/integration_helpers_test.go
@@ -258,7 +258,7 @@ func removePacksExcept(gopts GlobalOptions, t testing.TB, keep restic.IDSet, rem
 	rtest.OK(t, err)
 
 	// Get all tree packs
-	rtest.OK(t, r.LoadIndex(context.TODO()))
+	rtest.OK(t, r.LoadIndex(context.TODO(), nil))
 
 	treePacks := restic.NewIDSet()
 	r.Index().Each(context.TODO(), func(pb restic.PackedBlob) {
diff --git a/cmd/restic/progress.go b/cmd/restic/progress.go
index 4b6025a54..56cca9ace 100644
--- a/cmd/restic/progress.go
+++ b/cmd/restic/progress.go
@@ -29,13 +29,12 @@ func calculateProgressInterval(show bool, json bool) time.Duration {
 	return interval
 }
 
-// newProgressMax returns a progress.Counter that prints to stdout.
-func newProgressMax(show bool, max uint64, description string) *progress.Counter {
+// newTerminalProgressMax returns a progress.Counter that prints to stdout or terminal if provided.
+func newGenericProgressMax(show bool, max uint64, description string, print func(status string)) *progress.Counter {
 	if !show {
 		return nil
 	}
 	interval := calculateProgressInterval(show, false)
-	canUpdateStatus := stdoutCanUpdateStatus()
 
 	return progress.NewCounter(interval, max, func(v uint64, max uint64, d time.Duration, final bool) {
 		var status string
@@ -47,14 +46,28 @@ func newProgressMax(show bool, max uint64, description string) *progress.Counter
 				ui.FormatDuration(d), ui.FormatPercent(v, max), v, max, description)
 		}
 
-		printProgress(status, canUpdateStatus)
+		print(status)
 		if final {
 			fmt.Print("\n")
 		}
 	})
 }
 
-func printProgress(status string, canUpdateStatus bool) {
+func newTerminalProgressMax(show bool, max uint64, description string, term *termstatus.Terminal) *progress.Counter {
+	return newGenericProgressMax(show, max, description, func(status string) {
+		term.SetStatus([]string{status})
+	})
+}
+
+// newProgressMax calls newTerminalProgress without a terminal (print to stdout)
+func newProgressMax(show bool, max uint64, description string) *progress.Counter {
+	return newGenericProgressMax(show, max, description, printProgress)
+}
+
+func printProgress(status string) {
+
+	canUpdateStatus := stdoutCanUpdateStatus()
+
 	w := stdoutTerminalWidth()
 	if w > 0 {
 		if w < 3 {
diff --git a/internal/fuse/snapshots_dirstruct.go b/internal/fuse/snapshots_dirstruct.go
index 3080d4de8..a3cd11b14 100644
--- a/internal/fuse/snapshots_dirstruct.go
+++ b/internal/fuse/snapshots_dirstruct.go
@@ -328,7 +328,7 @@ func (d *SnapshotsDirStructure) updateSnapshots(ctx context.Context) error {
 		return nil
 	}
 
-	err = d.root.repo.LoadIndex(ctx)
+	err = d.root.repo.LoadIndex(ctx, nil)
 	if err != nil {
 		return err
 	}
diff --git a/internal/index/master_index_test.go b/internal/index/master_index_test.go
index 5260a11a0..493d11ace 100644
--- a/internal/index/master_index_test.go
+++ b/internal/index/master_index_test.go
@@ -357,7 +357,7 @@ func TestIndexSave(t *testing.T) {
 func testIndexSave(t *testing.T, version uint) {
 	repo := createFilledRepo(t, 3, version)
 
-	err := repo.LoadIndex(context.TODO())
+	err := repo.LoadIndex(context.TODO(), nil)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/repository/repack_test.go b/internal/repository/repack_test.go
index bb31eba77..c8570a9d4 100644
--- a/internal/repository/repack_test.go
+++ b/internal/repository/repack_test.go
@@ -213,7 +213,7 @@ func reloadIndex(t *testing.T, repo restic.Repository) {
 		t.Fatal(err)
 	}
 
-	if err := repo.LoadIndex(context.TODO()); err != nil {
+	if err := repo.LoadIndex(context.TODO(), nil); err != nil {
 		t.Fatalf("error loading new index: %v", err)
 	}
 }
diff --git a/internal/repository/repository.go b/internal/repository/repository.go
index 653c1f774..7c3c86d2b 100644
--- a/internal/repository/repository.go
+++ b/internal/repository/repository.go
@@ -581,15 +581,30 @@ func (r *Repository) SetIndex(i restic.MasterIndex) error {
 }
 
 // LoadIndex loads all index files from the backend in parallel and stores them
-// in the master index. The first error that occurred is returned.
-func (r *Repository) LoadIndex(ctx context.Context) error {
+func (r *Repository) LoadIndex(ctx context.Context, p *progress.Counter) error {
 	debug.Log("Loading index")
 
+	if p != nil {
+		var numIndexFiles uint64
+		err := r.List(ctx, restic.IndexFile, func(id restic.ID, size int64) error {
+			numIndexFiles++
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+		p.SetMax(numIndexFiles)
+		defer p.Done()
+	}
+
 	err := index.ForAllIndexes(ctx, r, func(id restic.ID, idx *index.Index, oldFormat bool, err error) error {
 		if err != nil {
 			return err
 		}
 		r.idx.Insert(idx)
+		if p != nil {
+			p.Add(1)
+		}
 		return nil
 	})
 
diff --git a/internal/repository/repository_test.go b/internal/repository/repository_test.go
index f26bf46f2..d362b0b5e 100644
--- a/internal/repository/repository_test.go
+++ b/internal/repository/repository_test.go
@@ -255,7 +255,7 @@ func TestRepositoryLoadIndex(t *testing.T) {
 	defer cleanup()
 
 	repo := repository.TestOpenLocal(t, repodir)
-	rtest.OK(t, repo.LoadIndex(context.TODO()))
+	rtest.OK(t, repo.LoadIndex(context.TODO(), nil))
 }
 
 // loadIndex loads the index id from backend and returns it.
@@ -324,7 +324,7 @@ func TestRepositoryLoadUnpackedRetryBroken(t *testing.T) {
 	err = repo.SearchKey(context.TODO(), rtest.TestPassword, 10, "")
 	rtest.OK(t, err)
 
-	rtest.OK(t, repo.LoadIndex(context.TODO()))
+	rtest.OK(t, repo.LoadIndex(context.TODO(), nil))
 }
 
 func BenchmarkLoadIndex(b *testing.B) {
diff --git a/internal/restic/repository.go b/internal/restic/repository.go
index 6990200e4..a651f9906 100644
--- a/internal/restic/repository.go
+++ b/internal/restic/repository.go
@@ -24,7 +24,7 @@ type Repository interface {
 	Key() *crypto.Key
 
 	Index() MasterIndex
-	LoadIndex(context.Context) error
+	LoadIndex(context.Context, *progress.Counter) error
 	SetIndex(MasterIndex) error
 	LookupBlobSize(ID, BlobType) (uint, bool)