From 68b74e359e08acf0ca8ac5724ef67d396dd743c4 Mon Sep 17 00:00:00 2001
From: Alexander Weiss <alex@weissfam.de>
Date: Sat, 5 Dec 2020 13:59:18 +0100
Subject: [PATCH] Count packs directly in RebuildIndexFiles

---
 cmd/restic/cmd_prune.go             | 19 +++++--------------
 cmd/restic/cmd_rebuild_index.go     |  6 +-----
 internal/repository/master_index.go | 10 ++++++++--
 internal/restic/repository.go       |  1 -
 4 files changed, 14 insertions(+), 22 deletions(-)

diff --git a/cmd/restic/cmd_prune.go b/cmd/restic/cmd_prune.go
index fd051e605..c8c665cde 100644
--- a/cmd/restic/cmd_prune.go
+++ b/cmd/restic/cmd_prune.go
@@ -479,11 +479,7 @@ func prune(opts PruneOptions, gopts GlobalOptions, repo restic.Repository, usedB
 		DeleteFiles(gopts, repo, removePacksFirst, restic.PackFile)
 	}
 
-	packsAddedByRepack := 0
 	if len(repackPacks) != 0 {
-		// Remember the number of unique packs before repacking
-		packsBeforeRepacking := len(repo.Index().Packs())
-
 		Verbosef("repacking packs\n")
 		bar := newProgressMax(!gopts.Quiet, uint64(len(repackPacks)), "packs repacked")
 		_, err := repository.Repack(ctx, repo, repackPacks, keepBlobs, bar)
@@ -492,18 +488,12 @@ func prune(opts PruneOptions, gopts GlobalOptions, repo restic.Repository, usedB
 			return err
 		}
 
-		// Since repacking will only add new packs, we can calculate the number
-		// of packs like this:
-		packsAddedByRepack = len(repo.Index().Packs()) - packsBeforeRepacking
-
 		// Also remove repacked packs
 		removePacks.Merge(repackPacks)
 	}
 
 	if len(removePacks) != 0 {
-		totalpacks := int(stats.packs.used+stats.packs.partlyUsed+stats.packs.unused) -
-			len(removePacks) + packsAddedByRepack
-		err = rebuildIndexFiles(gopts, repo, removePacks, nil, uint64(totalpacks))
+		err = rebuildIndexFiles(gopts, repo, removePacks, nil)
 		if err != nil {
 			return err
 		}
@@ -516,12 +506,13 @@ func prune(opts PruneOptions, gopts GlobalOptions, repo restic.Repository, usedB
 	return nil
 }
 
-func rebuildIndexFiles(gopts GlobalOptions, repo restic.Repository, removePacks restic.IDSet, extraObsolete restic.IDs, packcount uint64) error {
+func rebuildIndexFiles(gopts GlobalOptions, repo restic.Repository, removePacks restic.IDSet, extraObsolete restic.IDs) error {
 	Verbosef("rebuilding index\n")
 
+	idx := (repo.Index()).(*repository.MasterIndex)
+	packcount := uint64(len(idx.Packs(removePacks)))
 	bar := newProgressMax(!gopts.Quiet, packcount, "packs processed")
-	obsoleteIndexes, err := (repo.Index()).(*repository.MasterIndex).
-		Save(gopts.ctx, repo, removePacks, extraObsolete, bar)
+	obsoleteIndexes, err := idx.Save(gopts.ctx, repo, removePacks, extraObsolete, bar)
 	bar.Done()
 	if err != nil {
 		return err
diff --git a/cmd/restic/cmd_rebuild_index.go b/cmd/restic/cmd_rebuild_index.go
index 27f36d2b9..81bd121e4 100644
--- a/cmd/restic/cmd_rebuild_index.go
+++ b/cmd/restic/cmd_rebuild_index.go
@@ -60,7 +60,6 @@ func rebuildIndex(opts RebuildIndexOptions, gopts GlobalOptions, repo *repositor
 	var obsoleteIndexes restic.IDs
 	packSizeFromList := make(map[restic.ID]int64)
 	removePacks := restic.NewIDSet()
-	totalPacks := 0
 
 	if opts.ReadAllPacks {
 		// get old index files
@@ -76,7 +75,6 @@ func rebuildIndex(opts RebuildIndexOptions, gopts GlobalOptions, repo *repositor
 		err = repo.List(ctx, restic.PackFile, func(id restic.ID, size int64) error {
 			packSizeFromList[id] = size
 			removePacks.Insert(id)
-			totalPacks++
 			return nil
 		})
 		if err != nil {
@@ -99,7 +97,6 @@ func rebuildIndex(opts RebuildIndexOptions, gopts GlobalOptions, repo *repositor
 				packSizeFromList[id] = packSize
 				removePacks.Insert(id)
 			}
-			totalPacks++
 			delete(packSizeFromIndex, id)
 			return nil
 		})
@@ -123,11 +120,10 @@ func rebuildIndex(opts RebuildIndexOptions, gopts GlobalOptions, repo *repositor
 
 		for _, id := range invalidFiles {
 			Verboseff("skipped incomplete pack file: %v\n", id)
-			totalPacks--
 		}
 	}
 
-	err := rebuildIndexFiles(gopts, repo, removePacks, obsoleteIndexes, uint64(totalPacks))
+	err := rebuildIndexFiles(gopts, repo, removePacks, obsoleteIndexes)
 	if err != nil {
 		return err
 	}
diff --git a/internal/repository/master_index.go b/internal/repository/master_index.go
index 3447be61e..f4de3fdd2 100644
--- a/internal/repository/master_index.go
+++ b/internal/repository/master_index.go
@@ -100,13 +100,19 @@ func (mi *MasterIndex) Has(bh restic.BlobHandle) bool {
 }
 
 // Packs returns all packs that are covered by the index.
-func (mi *MasterIndex) Packs() restic.IDSet {
+// If packBlacklist ist given, those packs are only contained in the
+// resulting IDSet if they are contained in a non-final index.
+func (mi *MasterIndex) Packs(packBlacklist restic.IDSet) restic.IDSet {
 	mi.idxMutex.RLock()
 	defer mi.idxMutex.RUnlock()
 
 	packs := restic.NewIDSet()
 	for _, idx := range mi.idx {
-		packs.Merge(idx.Packs())
+		idxPacks := idx.Packs()
+		if idx.final {
+			idxPacks = idxPacks.Sub(packBlacklist)
+		}
+		packs.Merge(idxPacks)
 	}
 
 	return packs
diff --git a/internal/restic/repository.go b/internal/restic/repository.go
index e03cae4dc..38c611ce6 100644
--- a/internal/restic/repository.go
+++ b/internal/restic/repository.go
@@ -65,7 +65,6 @@ type MasterIndex interface {
 	Has(BlobHandle) bool
 	Lookup(BlobHandle) []PackedBlob
 	Count(BlobType) uint
-	Packs() IDSet
 	PackSize(ctx context.Context, onlyHdr bool) map[ID]int64
 
 	// Each returns a channel that yields all blobs known to the index. When