restore: Add progress bar

Co-authored-by: Mark Herrmann <mark.herrmann@mailbox.org>

internal/restorer/filerestorer.go

@@ -12,6 +12,7 @@ import (
 	"github.com/restic/restic/internal/errors"
 	"github.com/restic/restic/internal/repository"
 	"github.com/restic/restic/internal/restic"
+	"github.com/restic/restic/internal/ui/restore"
 )
 
 // TODO if a blob is corrupt, there may be good blob copies in other packs
@@ -54,6 +55,7 @@ type fileRestorer struct {
 	filesWriter *filesWriter
 	zeroChunk   restic.ID
 	sparse      bool
+	progress    *restore.Progress
 
 	dst   string
 	files []*fileInfo
@@ -65,7 +67,8 @@ func newFileRestorer(dst string,
 	key *crypto.Key,
 	idx func(restic.BlobHandle) []restic.PackedBlob,
 	connections uint,
-	sparse bool) *fileRestorer {
+	sparse bool,
+	progress *restore.Progress) *fileRestorer {
 
 	// as packs are streamed the concurrency is limited by IO
 	workerCount := int(connections)
@@ -77,6 +80,7 @@ func newFileRestorer(dst string,
 		filesWriter: newFilesWriter(workerCount),
 		zeroChunk:   repository.ZeroChunk(),
 		sparse:      sparse,
+		progress:    progress,
 		workerCount: workerCount,
 		dst:         dst,
 		Error:       restorerAbortOnAllErrors,
@@ -268,7 +272,13 @@ func (r *fileRestorer) downloadPack(ctx context.Context, pack *packInfo) error {
 						file.inProgress = true
 						createSize = file.size
 					}
-					return r.filesWriter.writeToFile(r.targetPath(file.location), blobData, offset, createSize, file.sparse)
+					writeErr := r.filesWriter.writeToFile(r.targetPath(file.location), blobData, offset, createSize, file.sparse)
+
+					if r.progress != nil {
+						r.progress.AddProgress(file.location, uint64(len(blobData)), uint64(file.size))
+					}
+
+					return writeErr
 				}
 				err := sanitizeError(file, writeToFile())
 				if err != nil {

internal/restorer/filerestorer_test.go

@@ -150,7 +150,7 @@ func newTestRepo(content []TestFile) *TestRepo {
 func restoreAndVerify(t *testing.T, tempdir string, content []TestFile, files map[string]bool, sparse bool) {
 	repo := newTestRepo(content)
 
-	r := newFileRestorer(tempdir, repo.loader, repo.key, repo.Lookup, 2, sparse)
+	r := newFileRestorer(tempdir, repo.loader, repo.key, repo.Lookup, 2, sparse, nil)
 
 	if files == nil {
 		r.files = repo.files
@@ -265,7 +265,7 @@ func TestErrorRestoreFiles(t *testing.T) {
 		return loadError
 	}
 
-	r := newFileRestorer(tempdir, repo.loader, repo.key, repo.Lookup, 2, false)
+	r := newFileRestorer(tempdir, repo.loader, repo.key, repo.Lookup, 2, false, nil)
 	r.files = repo.files
 
 	err := r.restoreFiles(context.TODO())
@@ -304,7 +304,7 @@ func testPartialDownloadError(t *testing.T, part int) {
 		return loader(ctx, h, length, offset, fn)
 	}
 
-	r := newFileRestorer(tempdir, repo.loader, repo.key, repo.Lookup, 2, false)
+	r := newFileRestorer(tempdir, repo.loader, repo.key, repo.Lookup, 2, false, nil)
 	r.files = repo.files
 	r.Error = func(s string, e error) error {
 		// ignore errors as in the `restore` command

internal/restorer/restorer.go

@@ -10,6 +10,7 @@ import (
 	"github.com/restic/restic/internal/errors"
 	"github.com/restic/restic/internal/fs"
 	"github.com/restic/restic/internal/restic"
+	restoreui "github.com/restic/restic/internal/ui/restore"
 
 	"golang.org/x/sync/errgroup"
 )
@@ -20,6 +21,8 @@ type Restorer struct {
 	sn     *restic.Snapshot
 	sparse bool
 
+	progress *restoreui.Progress
+
 	Error        func(location string, err error) error
 	SelectFilter func(item string, dstpath string, node *restic.Node) (selectedForRestore bool, childMayBeSelected bool)
 }
@@ -27,12 +30,14 @@ type Restorer struct {
 var restorerAbortOnAllErrors = func(location string, err error) error { return err }
 
 // NewRestorer creates a restorer preloaded with the content from the snapshot id.
-func NewRestorer(ctx context.Context, repo restic.Repository, sn *restic.Snapshot, sparse bool) *Restorer {
+func NewRestorer(ctx context.Context, repo restic.Repository, sn *restic.Snapshot, sparse bool,
+	progress *restoreui.Progress) *Restorer {
 	r := &Restorer{
 		repo:         repo,
 		sparse:       sparse,
 		Error:        restorerAbortOnAllErrors,
 		SelectFilter: func(string, string, *restic.Node) (bool, bool) { return true, true },
+		progress:     progress,
 		sn:           sn,
 	}
 
@@ -186,6 +191,11 @@ func (res *Restorer) restoreHardlinkAt(node *restic.Node, target, path, location
 	if err != nil {
 		return errors.WithStack(err)
 	}
+
+	if res.progress != nil {
+		res.progress.AddProgress(location, 0, 0)
+	}
+
 	// TODO investigate if hardlinks have separate metadata on any supported system
 	return res.restoreNodeMetadataTo(node, path, location)
 }
@@ -200,6 +210,10 @@ func (res *Restorer) restoreEmptyFileAt(node *restic.Node, target, location stri
 		return err
 	}
 
+	if res.progress != nil {
+		res.progress.AddProgress(location, 0, 0)
+	}
+
 	return res.restoreNodeMetadataTo(node, target, location)
 }
 
@@ -215,7 +229,8 @@ func (res *Restorer) RestoreTo(ctx context.Context, dst string) error {
 	}
 
 	idx := NewHardlinkIndex()
-	filerestorer := newFileRestorer(dst, res.repo.Backend().Load, res.repo.Key(), res.repo.Index().Lookup, res.repo.Connections(), res.sparse)
+	filerestorer := newFileRestorer(dst, res.repo.Backend().Load, res.repo.Key(), res.repo.Index().Lookup,
+		res.repo.Connections(), res.sparse, res.progress)
 	filerestorer.Error = res.Error
 
 	debug.Log("first pass for %q", dst)
@@ -242,6 +257,10 @@ func (res *Restorer) RestoreTo(ctx context.Context, dst string) error {
 				return nil
 			}
 
+			if res.progress != nil {
+				res.progress.AddFile(node.Size)
+			}
+
 			if node.Size == 0 {
 				return nil // deal with empty files later
 			}

internal/restorer/restorer_test.go

@@ -325,7 +325,7 @@ func TestRestorer(t *testing.T) {
 			sn, id := saveSnapshot(t, repo, test.Snapshot)
 			t.Logf("snapshot saved as %v", id.Str())
 
-			res := NewRestorer(context.TODO(), repo, sn, false)
+			res := NewRestorer(context.TODO(), repo, sn, false, nil)
 
 			tempdir := rtest.TempDir(t)
 			// make sure we're creating a new subdir of the tempdir
@@ -442,7 +442,7 @@ func TestRestorerRelative(t *testing.T) {
 			sn, id := saveSnapshot(t, repo, test.Snapshot)
 			t.Logf("snapshot saved as %v", id.Str())
 
-			res := NewRestorer(context.TODO(), repo, sn, false)
+			res := NewRestorer(context.TODO(), repo, sn, false, nil)
 
 			tempdir := rtest.TempDir(t)
 			cleanup := rtest.Chdir(t, tempdir)
@@ -671,7 +671,7 @@ func TestRestorerTraverseTree(t *testing.T) {
 			repo := repository.TestRepository(t)
 			sn, _ := saveSnapshot(t, repo, test.Snapshot)
 
-			res := NewRestorer(context.TODO(), repo, sn, false)
+			res := NewRestorer(context.TODO(), repo, sn, false, nil)
 
 			res.SelectFilter = test.Select
 
@@ -747,7 +747,7 @@ func TestRestorerConsistentTimestampsAndPermissions(t *testing.T) {
 		},
 	})
 
-	res := NewRestorer(context.TODO(), repo, sn, false)
+	res := NewRestorer(context.TODO(), repo, sn, false, nil)
 
 	res.SelectFilter = func(item string, dstpath string, node *restic.Node) (selectedForRestore bool, childMayBeSelected bool) {
 		switch filepath.ToSlash(item) {
@@ -802,7 +802,7 @@ func TestVerifyCancel(t *testing.T) {
 	repo := repository.TestRepository(t)
 	sn, _ := saveSnapshot(t, repo, snapshot)
 
-	res := NewRestorer(context.TODO(), repo, sn, false)
+	res := NewRestorer(context.TODO(), repo, sn, false, nil)
 
 	tempdir := rtest.TempDir(t)
 	ctx, cancel := context.WithCancel(context.Background())
@@ -844,7 +844,7 @@ func TestRestorerSparseFiles(t *testing.T) {
 		archiver.SnapshotOptions{})
 	rtest.OK(t, err)
 
-	res := NewRestorer(context.TODO(), repo, sn, true)
+	res := NewRestorer(context.TODO(), repo, sn, true, nil)
 
 	tempdir := rtest.TempDir(t)
 	ctx, cancel := context.WithCancel(context.Background())

internal/restorer/restorer_unix_test.go

@@ -29,7 +29,7 @@ func TestRestorerRestoreEmptyHardlinkedFileds(t *testing.T) {
 		},
 	})
 
-	res := NewRestorer(context.TODO(), repo, sn, false)
+	res := NewRestorer(context.TODO(), repo, sn, false, nil)
 
 	res.SelectFilter = func(item string, dstpath string, node *restic.Node) (selectedForRestore bool, childMayBeSelected bool) {
 		return true, true

internal/ui/restore/progressformatter.go

@@ -0,0 +1,131 @@
+package restore
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/restic/restic/internal/ui"
+	"github.com/restic/restic/internal/ui/progress"
+)
+
+type Progress struct {
+	updater progress.Updater
+	m       sync.Mutex
+
+	progressInfoMap map[string]progressInfoEntry
+	filesFinished   uint64
+	filesTotal      uint64
+	allBytesWritten uint64
+	allBytesTotal   uint64
+	started         time.Time
+
+	printer ProgressPrinter
+}
+
+type progressInfoEntry struct {
+	bytesWritten uint64
+	bytesTotal   uint64
+}
+
+type ProgressPrinter interface {
+	Update(filesFinished, filesTotal, allBytesWritten, allBytesTotal uint64, duration time.Duration)
+	Finish(filesFinished, filesTotal, allBytesWritten, allBytesTotal uint64, duration time.Duration)
+}
+
+func NewProgress(printer ProgressPrinter, interval time.Duration) *Progress {
+	p := &Progress{
+		progressInfoMap: make(map[string]progressInfoEntry),
+		started:         time.Now(),
+		printer:         printer,
+	}
+	p.updater = *progress.NewUpdater(interval, p.update)
+	return p
+}
+
+func (p *Progress) update(runtime time.Duration, final bool) {
+	p.m.Lock()
+	defer p.m.Unlock()
+
+	if !final {
+		p.printer.Update(p.filesFinished, p.filesTotal, p.allBytesWritten, p.allBytesTotal, runtime)
+	} else {
+		p.printer.Finish(p.filesFinished, p.filesTotal, p.allBytesWritten, p.allBytesTotal, runtime)
+	}
+}
+
+// AddFile starts tracking a new file with the given size
+func (p *Progress) AddFile(size uint64) {
+	p.m.Lock()
+	defer p.m.Unlock()
+
+	p.filesTotal++
+	p.allBytesTotal += size
+}
+
+// AddProgress accumulates the number of bytes written for a file
+func (p *Progress) AddProgress(name string, bytesWrittenPortion uint64, bytesTotal uint64) {
+	p.m.Lock()
+	defer p.m.Unlock()
+
+	entry, exists := p.progressInfoMap[name]
+	if !exists {
+		entry.bytesTotal = bytesTotal
+	}
+	entry.bytesWritten += bytesWrittenPortion
+	p.progressInfoMap[name] = entry
+
+	p.allBytesWritten += bytesWrittenPortion
+	if entry.bytesWritten == entry.bytesTotal {
+		delete(p.progressInfoMap, name)
+		p.filesFinished++
+	}
+}
+
+func (p *Progress) Finish() {
+	p.updater.Done()
+}
+
+type term interface {
+	Print(line string)
+	SetStatus(lines []string)
+}
+
+type textPrinter struct {
+	terminal term
+}
+
+func NewProgressPrinter(terminal term) ProgressPrinter {
+	return &textPrinter{
+		terminal: terminal,
+	}
+}
+
+func (t *textPrinter) Update(filesFinished, filesTotal, allBytesWritten, allBytesTotal uint64, duration time.Duration) {
+	timeLeft := ui.FormatDuration(duration)
+	formattedAllBytesWritten := ui.FormatBytes(allBytesWritten)
+	formattedAllBytesTotal := ui.FormatBytes(allBytesTotal)
+	allPercent := ui.FormatPercent(allBytesWritten, allBytesTotal)
+	progress := fmt.Sprintf("[%s] %s  %v files %s, total %v files %v",
+		timeLeft, allPercent, filesFinished, formattedAllBytesWritten, filesTotal, formattedAllBytesTotal)
+
+	t.terminal.SetStatus([]string{progress})
+}
+
+func (t *textPrinter) Finish(filesFinished, filesTotal, allBytesWritten, allBytesTotal uint64, duration time.Duration) {
+	t.terminal.SetStatus([]string{})
+
+	timeLeft := ui.FormatDuration(duration)
+	formattedAllBytesTotal := ui.FormatBytes(allBytesTotal)
+
+	var summary string
+	if filesFinished == filesTotal && allBytesWritten == allBytesTotal {
+		summary = fmt.Sprintf("Summary: Restored %d Files (%s) in %s", filesTotal, formattedAllBytesTotal, timeLeft)
+	} else {
+		formattedAllBytesWritten := ui.FormatBytes(allBytesWritten)
+		summary = fmt.Sprintf("Summary: Restored %d / %d Files (%s / %s) in %s",
+			filesFinished, filesTotal, formattedAllBytesWritten, formattedAllBytesTotal, timeLeft)
+	}
+
+	t.terminal.Print(summary)
+}

internal/ui/restore/progressformatter_test.go

@@ -0,0 +1,170 @@
+package restore
+
+import (
+	"testing"
+	"time"
+
+	"github.com/restic/restic/internal/test"
+)
+
+type printerTraceEntry struct {
+	filesFinished, filesTotal, allBytesWritten, allBytesTotal uint64
+
+	duration   time.Duration
+	isFinished bool
+}
+
+type printerTrace []printerTraceEntry
+
+type mockPrinter struct {
+	trace printerTrace
+}
+
+const mockFinishDuration = 42 * time.Second
+
+func (p *mockPrinter) Update(filesFinished, filesTotal, allBytesWritten, allBytesTotal uint64, duration time.Duration) {
+	p.trace = append(p.trace, printerTraceEntry{filesFinished, filesTotal, allBytesWritten, allBytesTotal, duration, false})
+}
+func (p *mockPrinter) Finish(filesFinished, filesTotal, allBytesWritten, allBytesTotal uint64, duration time.Duration) {
+	p.trace = append(p.trace, printerTraceEntry{filesFinished, filesTotal, allBytesWritten, allBytesTotal, mockFinishDuration, true})
+}
+
+func testProgress(fn func(progress *Progress) bool) printerTrace {
+	printer := &mockPrinter{}
+	progress := NewProgress(printer, 0)
+	final := fn(progress)
+	progress.update(0, final)
+	trace := append(printerTrace{}, printer.trace...)
+	// cleanup to avoid goroutine leak, but copy trace first
+	progress.Finish()
+	return trace
+}
+
+func TestNew(t *testing.T) {
+	result := testProgress(func(progress *Progress) bool {
+		return false
+	})
+	test.Equals(t, printerTrace{
+		printerTraceEntry{0, 0, 0, 0, 0, false},
+	}, result)
+}
+
+func TestAddFile(t *testing.T) {
+	fileSize := uint64(100)
+
+	result := testProgress(func(progress *Progress) bool {
+		progress.AddFile(fileSize)
+		return false
+	})
+	test.Equals(t, printerTrace{
+		printerTraceEntry{0, 1, 0, fileSize, 0, false},
+	}, result)
+}
+
+func TestFirstProgressOnAFile(t *testing.T) {
+	expectedBytesWritten := uint64(5)
+	expectedBytesTotal := uint64(100)
+
+	result := testProgress(func(progress *Progress) bool {
+		progress.AddFile(expectedBytesTotal)
+		progress.AddProgress("test", expectedBytesWritten, expectedBytesTotal)
+		return false
+	})
+	test.Equals(t, printerTrace{
+		printerTraceEntry{0, 1, expectedBytesWritten, expectedBytesTotal, 0, false},
+	}, result)
+}
+
+func TestLastProgressOnAFile(t *testing.T) {
+	fileSize := uint64(100)
+
+	result := testProgress(func(progress *Progress) bool {
+		progress.AddFile(fileSize)
+		progress.AddProgress("test", 30, fileSize)
+		progress.AddProgress("test", 35, fileSize)
+		progress.AddProgress("test", 35, fileSize)
+		return false
+	})
+	test.Equals(t, printerTrace{
+		printerTraceEntry{1, 1, fileSize, fileSize, 0, false},
+	}, result)
+}
+
+func TestLastProgressOnLastFile(t *testing.T) {
+	fileSize := uint64(100)
+
+	result := testProgress(func(progress *Progress) bool {
+		progress.AddFile(fileSize)
+		progress.AddFile(50)
+		progress.AddProgress("test1", 50, 50)
+		progress.AddProgress("test2", 50, fileSize)
+		progress.AddProgress("test2", 50, fileSize)
+		return false
+	})
+	test.Equals(t, printerTrace{
+		printerTraceEntry{2, 2, 50 + fileSize, 50 + fileSize, 0, false},
+	}, result)
+}
+
+func TestSummaryOnSuccess(t *testing.T) {
+	fileSize := uint64(100)
+
+	result := testProgress(func(progress *Progress) bool {
+		progress.AddFile(fileSize)
+		progress.AddFile(50)
+		progress.AddProgress("test1", 50, 50)
+		progress.AddProgress("test2", fileSize, fileSize)
+		return true
+	})
+	test.Equals(t, printerTrace{
+		printerTraceEntry{2, 2, 50 + fileSize, 50 + fileSize, mockFinishDuration, true},
+	}, result)
+}
+
+func TestSummaryOnErrors(t *testing.T) {
+	fileSize := uint64(100)
+
+	result := testProgress(func(progress *Progress) bool {
+		progress.AddFile(fileSize)
+		progress.AddFile(50)
+		progress.AddProgress("test1", 50, 50)
+		progress.AddProgress("test2", fileSize/2, fileSize)
+		return true
+	})
+	test.Equals(t, printerTrace{
+		printerTraceEntry{1, 2, 50 + fileSize/2, 50 + fileSize, mockFinishDuration, true},
+	}, result)
+}
+
+type mockTerm struct {
+	output []string
+}
+
+func (m *mockTerm) Print(line string) {
+	m.output = append(m.output, line)
+}
+
+func (m *mockTerm) SetStatus(lines []string) {
+	m.output = append([]string{}, lines...)
+}
+
+func TestPrintUpdate(t *testing.T) {
+	term := &mockTerm{}
+	printer := NewProgressPrinter(term)
+	printer.Update(3, 11, 29, 47, 5*time.Second)
+	test.Equals(t, []string{"[0:05] 61.70%  3 files 29 B, total 11 files 47 B"}, term.output)
+}
+
+func TestPrintSummaryOnSuccess(t *testing.T) {
+	term := &mockTerm{}
+	printer := NewProgressPrinter(term)
+	printer.Finish(11, 11, 47, 47, 5*time.Second)
+	test.Equals(t, []string{"Summary: Restored 11 Files (47 B) in 0:05"}, term.output)
+}
+
+func TestPrintSummaryOnErrors(t *testing.T) {
+	term := &mockTerm{}
+	printer := NewProgressPrinter(term)
+	printer.Finish(3, 11, 29, 47, 5*time.Second)
+	test.Equals(t, []string{"Summary: Restored 3 / 11 Files (29 B / 47 B) in 0:05"}, term.output)
+}