From 1c01af4e69da164d7e43a3b38958387f2f8f6eed Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Wed, 21 Apr 2021 12:09:05 -0600
Subject: [PATCH] Fix snapshot selection during restore-by-index

This commit fixes a bug where restoring to a specific index will
incorrectly choose the latest snapshot instead of choosing the
latest snapshot that occurred before the given index.
---
 Makefile   |  2 +-
 db.go      | 31 ++++++++++++++++++++++++-------
 replica.go | 36 ++++++++++++++++++++++++++++++++----
 3 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/Makefile b/Makefile
index c5058cd..e3d75e4 100644
--- a/Makefile
+++ b/Makefile
@@ -6,7 +6,7 @@ docker:
 dist-linux:
 	mkdir -p dist
 	cp etc/litestream.yml dist/litestream.yml
-	docker run --rm -v "${PWD}":/usr/src/litestream -w /usr/src/litestream -e GOOS=linux -e GOARCH=amd64 golang:1.16 go build -v -o dist/litestream ./cmd/litestream
+	docker run --rm -v "${PWD}":/usr/src/litestream -w /usr/src/litestream -e GOOS=linux -e GOARCH=amd64 golang:1.16 go build -v -ldflags "-s -w" -o dist/litestream ./cmd/litestream
 	tar -cz -f dist/litestream-linux-amd64.tar.gz -C dist litestream
 
 dist-linux-arm:
diff --git a/db.go b/db.go
index 3635dc7..2e2da6b 100644
--- a/db.go
+++ b/db.go
@@ -1434,7 +1434,7 @@ func (db *DB) monitor() {
 // replica or generation or it will automatically choose the best one. Finally,
 // a timestamp can be specified to restore the database to a specific
 // point-in-time.
-func RestoreReplica(ctx context.Context, r Replica, opt RestoreOptions) error {
+func RestoreReplica(ctx context.Context, r Replica, opt RestoreOptions) (err error) {
 	// Validate options.
 	if opt.OutputPath == "" {
 		return fmt.Errorf("output path required")
@@ -1462,10 +1462,16 @@ func RestoreReplica(ctx context.Context, r Replica, opt RestoreOptions) error {
 		return err
 	}
 
-	// Find lastest snapshot that occurs before timestamp.
-	minWALIndex, err := SnapshotIndexAt(ctx, r, opt.Generation, opt.Timestamp)
-	if err != nil {
-		return fmt.Errorf("cannot find snapshot index for restore: %w", err)
+	// Find lastest snapshot that occurs before timestamp or index.
+	var minWALIndex int
+	if opt.Index < math.MaxInt32 {
+		if minWALIndex, err = SnapshotIndexByIndex(ctx, r, opt.Generation, opt.Index); err != nil {
+			return fmt.Errorf("cannot find snapshot index: %w", err)
+		}
+	} else {
+		if minWALIndex, err = SnapshotIndexAt(ctx, r, opt.Generation, opt.Timestamp); err != nil {
+			return fmt.Errorf("cannot find snapshot index by timestamp: %w", err)
+		}
 	}
 
 	// Find the maximum WAL index that occurs before timestamp.
@@ -1487,7 +1493,7 @@ func RestoreReplica(ctx context.Context, r Replica, opt RestoreOptions) error {
 
 	// Restore each WAL file until we reach our maximum index.
 	for index := minWALIndex; index <= maxWALIndex; index++ {
-		logger.Printf("%s: restoring wal %s/%08x", logPrefix, opt.Generation, index)
+		restoreStartTime := time.Now()
 		if err = restoreWAL(ctx, r, opt.Generation, index, tmpPath); os.IsNotExist(err) && index == minWALIndex && index == maxWALIndex {
 			logger.Printf("%s: no wal available, snapshot only", logPrefix)
 			break // snapshot file only, ignore error
@@ -1495,10 +1501,21 @@ func RestoreReplica(ctx context.Context, r Replica, opt RestoreOptions) error {
 			return fmt.Errorf("cannot restore wal: %w", err)
 		}
 
-		logger.Printf("%s: applying wal %s/%08x", logPrefix, opt.Generation, index)
+		fi, err := os.Stat(tmpPath + "-wal")
+		if err != nil {
+			return fmt.Errorf("cannot stat restored wal: %w", err)
+		}
+
+		applyStartTime := time.Now()
 		if err = applyWAL(ctx, tmpPath); err != nil {
 			return fmt.Errorf("cannot apply wal: %w", err)
 		}
+		logger.Printf("%s: restored wal %s/%08x (sz=%d restore=%s apply=%s)",
+			logPrefix, opt.Generation, index,
+			fi.Size(),
+			applyStartTime.Sub(restoreStartTime).String(),
+			time.Since(applyStartTime).String(),
+		)
 	}
 
 	// Copy file to final location.
diff --git a/replica.go b/replica.go
index 2b00ade..9d2cc47 100644
--- a/replica.go
+++ b/replica.go
@@ -1064,7 +1064,7 @@ func SnapshotIndexAt(ctx context.Context, r Replica, generation string, timestam
 		return 0, ErrNoSnapshots
 	}
 
-	index := -1
+	snapshotIndex := -1
 	var max time.Time
 	for _, snapshot := range snapshots {
 		if !timestamp.IsZero() && snapshot.CreatedAt.After(timestamp) {
@@ -1073,14 +1073,42 @@ func SnapshotIndexAt(ctx context.Context, r Replica, generation string, timestam
 
 		// Use snapshot if it newer.
 		if max.IsZero() || snapshot.CreatedAt.After(max) {
-			index, max = snapshot.Index, snapshot.CreatedAt
+			snapshotIndex, max = snapshot.Index, snapshot.CreatedAt
 		}
 	}
 
-	if index == -1 {
+	if snapshotIndex == -1 {
 		return 0, ErrNoSnapshots
 	}
-	return index, nil
+	return snapshotIndex, nil
+}
+
+// SnapshotIndexbyIndex returns the highest index for a snapshot within a generation
+// that occurs before a given index. If index is MaxInt32, returns the latest snapshot.
+func SnapshotIndexByIndex(ctx context.Context, r Replica, generation string, index int) (int, error) {
+	snapshots, err := r.Snapshots(ctx)
+	if err != nil {
+		return 0, err
+	} else if len(snapshots) == 0 {
+		return 0, ErrNoSnapshots
+	}
+
+	snapshotIndex := -1
+	for _, snapshot := range snapshots {
+		if index < math.MaxInt32 && snapshot.Index > index {
+			continue // after index, skip
+		}
+
+		// Use snapshot if it newer.
+		if snapshotIndex == -1 || snapshotIndex >= snapshotIndex {
+			snapshotIndex = snapshot.Index
+		}
+	}
+
+	if snapshotIndex == -1 {
+		return 0, ErrNoSnapshots
+	}
+	return snapshotIndex, nil
 }
 
 // WALIndexAt returns the highest index for a WAL file that occurs before maxIndex & timestamp.