From 8a05de537f2503055e9c52eec8af1ec5ea9213c6 Mon Sep 17 00:00:00 2001
From: Pauline Middelink <middelin@google.com>
Date: Wed, 8 Mar 2017 20:30:52 +0100
Subject: [PATCH] Refactor `prune` and `rebuild_index`

Factor out and reuse `rebuildIndex()` in cmd_rebuild_index and cmd_prune.

Use contexts.
---
 src/cmds/restic/cmd_prune.go         | 35 ++++------------------------
 src/cmds/restic/cmd_rebuild_index.go | 25 +++++++++++---------
 2 files changed, 19 insertions(+), 41 deletions(-)

diff --git a/src/cmds/restic/cmd_prune.go b/src/cmds/restic/cmd_prune.go
index fec8126f2..03d14d300 100644
--- a/src/cmds/restic/cmd_prune.go
+++ b/src/cmds/restic/cmd_prune.go
@@ -1,8 +1,8 @@
 package main
 
 import (
+	"context"
 	"fmt"
-	"os"
 	"restic"
 	"restic/debug"
 	"restic/errors"
@@ -81,8 +81,8 @@ func pruneRepository(gopts GlobalOptions, repo restic.Repository) error {
 		return err
 	}
 
-	done := make(chan struct{})
-	defer close(done)
+	ctx, cancel := context.WithCancel(gopts.ctx)
+	defer cancel()
 
 	var stats struct {
 		blobs     int
@@ -92,7 +92,7 @@ func pruneRepository(gopts GlobalOptions, repo restic.Repository) error {
 	}
 
 	Verbosef("counting files in repo\n")
-	for _ = range repo.List(restic.DataFile, done) {
+	for _ = range repo.List(restic.DataFile, ctx.Done()) {
 		stats.packs++
 	}
 
@@ -238,35 +238,10 @@ func pruneRepository(gopts GlobalOptions, repo restic.Repository) error {
 		bar.Done()
 	}
 
-	Verbosef("creating new index\n")
-
-	stats.packs = 0
-	for _ = range repo.List(restic.DataFile, done) {
-		stats.packs++
-	}
-	bar = newProgressMax(!gopts.Quiet, uint64(stats.packs), "packs")
-	idx, err = index.New(repo, bar)
-	if err != nil {
+	if err = rebuildIndex(ctx, repo); err != nil {
 		return err
 	}
 
-	var supersedes restic.IDs
-	for idxID := range repo.List(restic.IndexFile, done) {
-		h := restic.Handle{Type: restic.IndexFile, Name: idxID.String()}
-		err := repo.Backend().Remove(h)
-		if err != nil {
-			fmt.Fprintf(os.Stderr, "unable to remove index %v: %v\n", idxID.Str(), err)
-		}
-
-		supersedes = append(supersedes, idxID)
-	}
-
-	id, err := idx.Save(repo, supersedes)
-	if err != nil {
-		return err
-	}
-	Verbosef("saved new index as %v\n", id.Str())
-
 	Verbosef("done\n")
 	return nil
 }
diff --git a/src/cmds/restic/cmd_rebuild_index.go b/src/cmds/restic/cmd_rebuild_index.go
index 6f6647daa..e392b80ca 100644
--- a/src/cmds/restic/cmd_rebuild_index.go
+++ b/src/cmds/restic/cmd_rebuild_index.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"restic"
 	"restic/index"
 
@@ -35,25 +36,29 @@ func runRebuildIndex(gopts GlobalOptions) error {
 		return err
 	}
 
-	done := make(chan struct{})
-	defer close(done)
+	ctx, cancel := context.WithCancel(gopts.ctx)
+	defer cancel()
+	return rebuildIndex(ctx, repo)
+}
 
+func rebuildIndex(ctx context.Context, repo restic.Repository) error {
 	Verbosef("counting files in repo\n")
 
 	var packs uint64
-	for _ = range repo.List(restic.DataFile, done) {
+	for _ = range repo.List(restic.DataFile, ctx.Done()) {
 		packs++
 	}
 
-	bar := newProgressMax(!gopts.Quiet, packs, "packs")
+	bar := newProgressMax(!globalOptions.Quiet, packs, "packs")
 	idx, err := index.New(repo, bar)
 	if err != nil {
 		return err
 	}
 
-	Verbosef("listing old index files\n")
+	Verbosef("finding old index files\n")
+
 	var supersedes restic.IDs
-	for id := range repo.List(restic.IndexFile, done) {
+	for id := range repo.List(restic.IndexFile, ctx.Done()) {
 		supersedes = append(supersedes, id)
 	}
 
@@ -67,13 +72,11 @@ func runRebuildIndex(gopts GlobalOptions) error {
 	Verbosef("remove %d old index files\n", len(supersedes))
 
 	for _, id := range supersedes {
-		err := repo.Backend().Remove(restic.Handle{
+		if err := repo.Backend().Remove(restic.Handle{
 			Type: restic.IndexFile,
 			Name: id.String(),
-		})
-
-		if err != nil {
-			Warnf("error deleting old index %v: %v\n", id.Str(), err)
+		}); err != nil {
+			Warnf("error removing old index %v: %v\n", id.Str(), err)
 		}
 	}