prune: Parallelize repack command

changelog/unreleased/pull-2941

@@ -0,0 +1,8 @@
+Enhancement: Speed up repacking step of prune command
+
+The repack step of the prune command, which moves still used file parts into
+new pack files such that the old ones can be garbage collected later on, now
+processes multiple pack files in parallel. This is especially beneficial for
+high latency backends or when using a fast network connection.
+
+https://github.com/restic/restic/pull/2941

internal/repository/repack.go

@@ -2,14 +2,19 @@ package repository
 
 import (
 	"context"
+	"os"
+	"sync"
 
 	"github.com/restic/restic/internal/debug"
 	"github.com/restic/restic/internal/errors"
 	"github.com/restic/restic/internal/fs"
 	"github.com/restic/restic/internal/pack"
 	"github.com/restic/restic/internal/restic"
+	"golang.org/x/sync/errgroup"
 )
 
+const numRepackWorkers = 8
+
 // Repack takes a list of packs together with a list of blobs contained in
 // these packs. Each pack is loaded and the blobs listed in keepBlobs is saved
 // into a new pack. Returned is the list of obsolete packs which can then
@@ -22,36 +27,91 @@ func Repack(ctx context.Context, repo restic.Repository, packs restic.IDSet, kee
 
 	debug.Log("repacking %d packs while keeping %d blobs", len(packs), len(keepBlobs))
 
+	wg, ctx := errgroup.WithContext(ctx)
+
+	downloadQueue := make(chan restic.ID)
+	wg.Go(func() error {
+		defer close(downloadQueue)
 		for packID := range packs {
+			select {
+			case downloadQueue <- packID:
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+		}
+		return nil
+	})
+
+	type repackJob struct {
+		tempfile   *os.File
+		hash       restic.ID
+		packLength int64
+	}
+	processQueue := make(chan repackJob)
+	// used to close processQueue once all downloaders have finished
+	var downloadWG sync.WaitGroup
+
+	downloader := func() error {
+		defer downloadWG.Done()
+		for packID := range downloadQueue {
 			// load the complete pack into a temp file
 			h := restic.Handle{Type: restic.PackFile, Name: packID.String()}
 
 			tempfile, hash, packLength, err := DownloadAndHash(ctx, repo.Backend(), h)
 			if err != nil {
-			return nil, errors.Wrap(err, "Repack")
+				return errors.Wrap(err, "Repack")
 			}
 
 			debug.Log("pack %v loaded (%d bytes), hash %v", packID, packLength, hash)
 
 			if !packID.Equal(hash) {
-			return nil, errors.Errorf("hash does not match id: want %v, got %v", packID, hash)
+				return errors.Errorf("hash does not match id: want %v, got %v", packID, hash)
 			}
 
+			select {
+			case processQueue <- repackJob{tempfile, hash, packLength}:
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+		}
+		return nil
+	}
+
+	downloadWG.Add(numRepackWorkers)
+	for i := 0; i < numRepackWorkers; i++ {
+		wg.Go(downloader)
+	}
+	wg.Go(func() error {
+		downloadWG.Wait()
+		close(processQueue)
+		return nil
+	})
+
+	var keepMutex sync.Mutex
+	worker := func() error {
+		for job := range processQueue {
+			tempfile, packID, packLength := job.tempfile, job.hash, job.packLength
+
 			_, err = tempfile.Seek(0, 0)
 			if err != nil {
-			return nil, errors.Wrap(err, "Seek")
+				return errors.Wrap(err, "Seek")
 			}
 
 			blobs, err := pack.List(repo.Key(), tempfile, packLength)
 			if err != nil {
-			return nil, err
+				return err
 			}
 
 			debug.Log("processing pack %v, blobs: %v", packID, len(blobs))
 			var buf []byte
 			for _, entry := range blobs {
 				h := restic.BlobHandle{ID: entry.ID, Type: entry.Type}
-			if !keepBlobs.Has(h) {
+
+				keepMutex.Lock()
+				shouldKeep := keepBlobs.Has(h)
+				keepMutex.Unlock()
+
+				if !shouldKeep {
 					continue
 				}
 
@@ -64,50 +124,70 @@ func Repack(ctx context.Context, repo restic.Repository, packs restic.IDSet, kee
 
 				n, err := tempfile.ReadAt(buf, int64(entry.Offset))
 				if err != nil {
-				return nil, errors.Wrap(err, "ReadAt")
+					return errors.Wrap(err, "ReadAt")
 				}
 
 				if n != len(buf) {
-				return nil, errors.Errorf("read blob %v from %v: not enough bytes read, want %v, got %v",
+					return errors.Errorf("read blob %v from %v: not enough bytes read, want %v, got %v",
 						h, tempfile.Name(), len(buf), n)
 				}
 
 				nonce, ciphertext := buf[:repo.Key().NonceSize()], buf[repo.Key().NonceSize():]
 				plaintext, err := repo.Key().Open(ciphertext[:0], nonce, ciphertext, nil)
 				if err != nil {
-				return nil, err
+					return err
 				}
 
 				id := restic.Hash(plaintext)
 				if !id.Equal(entry.ID) {
 					debug.Log("read blob %v/%v from %v: wrong data returned, hash is %v",
 						h.Type, h.ID, tempfile.Name(), id)
-				return nil, errors.Errorf("read blob %v from %v: wrong data returned, hash is %v",
+					return errors.Errorf("read blob %v from %v: wrong data returned, hash is %v",
 						h, tempfile.Name(), id)
 				}
 
+				keepMutex.Lock()
+				// recheck whether some other worker was faster
+				shouldKeep = keepBlobs.Has(h)
+				if shouldKeep {
+					keepBlobs.Delete(h)
+				}
+				keepMutex.Unlock()
+
+				if !shouldKeep {
+					continue
+				}
+
 				// We do want to save already saved blobs!
 				_, _, err = repo.SaveBlob(ctx, entry.Type, plaintext, entry.ID, true)
 				if err != nil {
-				return nil, err
+					return err
 				}
 
 				debug.Log("  saved blob %v", entry.ID)
-
-			keepBlobs.Delete(h)
 			}
 
 			if err = tempfile.Close(); err != nil {
-			return nil, errors.Wrap(err, "Close")
+				return errors.Wrap(err, "Close")
 			}
 
 			if err = fs.RemoveIfExists(tempfile.Name()); err != nil {
-			return nil, errors.Wrap(err, "Remove")
+				return errors.Wrap(err, "Remove")
 			}
 			if p != nil {
 				p.Report(restic.Stat{Blobs: 1})
 			}
 		}
+		return nil
+	}
+
+	for i := 0; i < numRepackWorkers; i++ {
+		wg.Go(worker)
+	}
+
+	if err := wg.Wait(); err != nil {
+		return nil, err
+	}
 
 	if err := repo.Flush(ctx); err != nil {
 		return nil, err