Implement multi-level restore (#652)

Author: benbjohnson

cmd/litestream/restore.go

@@ -23,7 +23,6 @@ func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) {
 	fs := flag.NewFlagSet("litestream-restore", flag.ContinueOnError)
 	configPath, noExpandEnv := registerConfigFlag(fs)
 	fs.StringVar(&opt.OutputPath, "o", "", "output path")
-	fs.StringVar(&opt.ReplicaName, "replica", "", "replica name")
 	fs.Var((*txidVar)(&opt.TXID), "txid", "transaction ID")
 	fs.IntVar(&opt.Parallelism, "parallelism", opt.Parallelism, "parallelism")
 	ifDBNotExists := fs.Bool("if-db-not-exists", false, "")
@@ -160,10 +159,6 @@ Arguments:
 	-no-expand-env
 	    Disables environment variable expansion in configuration file.
 
-	-replica NAME
-	    Restore from a specific replica.
-	    Defaults to replica with latest data.
-
 	-index NUM
 	    Restore up to a specific hex-encoded WAL index (inclusive).
 	    Defaults to use the highest available index.

db.go

@@ -1336,10 +1336,6 @@ type RestoreOptions struct {
 	// If blank, the original DB path is used.
 	OutputPath string
 
-	// Specific replica to restore from.
-	// If blank, all replicas are considered.
-	ReplicaName string
-
 	// Specific transaction to restore to.
 	// If zero, TXID is ignored.
 	TXID ltx.TXID

go.mod

@@ -15,7 +15,7 @@ require (
 	github.com/mattn/go-sqlite3 v1.14.19
 	github.com/pkg/sftp v1.13.6
 	github.com/prometheus/client_golang v1.17.0
-	github.com/superfly/ltx v0.3.15
+	github.com/superfly/ltx v0.3.16
 	golang.org/x/crypto v0.17.0
 	golang.org/x/sync v0.5.0
 	golang.org/x/sys v0.15.0

go.sum

@@ -166,8 +166,8 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/superfly/ltx v0.3.15 h1:gc+iTncNH5FMPMzxAtRS7DvL0ayxCOCrKjS9S8Q7PQA=
-github.com/superfly/ltx v0.3.15/go.mod h1:Nf50QAIXU/ET4ua3AuQ2fh31MbgNQZA7r/DYx6Os77s=
+github.com/superfly/ltx v0.3.16 h1:osibxhNf/4+DOA4fKuPxXVMPKvqoHgtCdy0tYb+AdVQ=
+github.com/superfly/ltx v0.3.16/go.mod h1:Nf50QAIXU/ET4ua3AuQ2fh31MbgNQZA7r/DYx6Os77s=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
 github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=

replica.go

@@ -78,6 +78,12 @@ func NewReplica(db *DB) *Replica {
 	return r
 }
 
+func NewReplicaWithClient(db *DB, client ReplicaClient) *Replica {
+	r := NewReplica(db)
+	r.Client = client
+	return r
+}
+
 // Logger returns the DB sub-logger for this replica.
 func (r *Replica) Logger() *slog.Logger {
 	logger := slog.Default()
@@ -663,6 +669,57 @@ func (r *Replica) Restore(ctx context.Context, opt RestoreOptions) (err error) {
 	return nil
 }
 
+// CalcRestorePlan returns a list of storage paths to restore a snapshot at the given TXID.
+func (r *Replica) CalcRestorePlan(ctx context.Context, txID ltx.TXID) ([]*ltx.FileInfo, error) {
+	var infos ltx.FileInfoSlice
+	logger := r.Logger().With("target", txID)
+
+	// Start with latest snapshot before target TXID.
+	if a, err := FindLTXFiles(ctx, r.Client, SnapshotLevel, func(info *ltx.FileInfo) (bool, error) {
+		return info.MaxTXID <= txID, nil
+	}); err != nil {
+		return nil, err
+	} else if len(a) > 0 {
+		logger.Debug("found snapshot before target TXID", "snapshot", a[len(a)-1].MaxTXID)
+		infos = append(infos, a[len(a)-1])
+	}
+
+	// Starting from the highest compaction level, collect all paths after the
+	// latest TXID for each level. Compactions are based on the previous level's
+	// TXID granularity so the TXIDs should align between compaction levels.
+	const maxLevel = SnapshotLevel - 1
+	for level := maxLevel; level >= 0; level-- {
+		a, err := FindLTXFiles(ctx, r.Client, level, func(info *ltx.FileInfo) (bool, error) {
+			return info.MinTXID > infos.MaxTXID() && info.MaxTXID <= txID, nil
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		// Append each storage path to the list
+		for i := range a {
+			// Ensure TXIDs are contiguous between each paths.
+			if infos.MaxTXID()+1 != a[i].MinTXID {
+				return nil, fmt.Errorf("non-contiguous transaction files: prev=%s filename=%s",
+					infos.MaxTXID().String(), ltx.FormatFilename(a[i].MinTXID, a[i].MaxTXID))
+			}
+
+			logger.Debug("matching LTX file for restore",
+				"filename", ltx.FormatFilename(a[len(a)-1].MinTXID, a[len(a)-1].MaxTXID))
+			infos = append(infos, a[i])
+		}
+	}
+
+	// Return an error if we are unable to find any set of LTX files before
+	// target TXID. This shouldn't happen under normal circumstances. Only if
+	// lower level LTX files are removed before a snapshot has occurred.
+	if len(infos) == 0 {
+		return nil, ErrTxNotAvailable
+	}
+
+	return infos, nil
+}
+
 // Replica metrics.
 var (
 	replicaValidationTotalCounterVec = promauto.NewCounterVec(prometheus.CounterOpts{

replica_client.go

@@ -2,11 +2,14 @@ package litestream
 
 import (
 	"context"
+	"errors"
 	"io"
 
 	"github.com/superfly/ltx"
 )
 
+var ErrStopIter = errors.New("stop iterator")
+
 // ReplicaClient represents client to connect to a Replica.
 type ReplicaClient interface {
 	// Type returns the type of client.
@@ -30,3 +33,33 @@ type ReplicaClient interface {
 	// DeleteAll deletes all files.
 	DeleteAll(ctx context.Context) error
 }
+
+// FindLTXFiles returns a list of files that match filter.
+func FindLTXFiles(ctx context.Context, client ReplicaClient, level int, filter func(*ltx.FileInfo) (bool, error)) ([]*ltx.FileInfo, error) {
+	itr, err := client.LTXFiles(ctx, level, 0)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = itr.Close() }()
+
+	var a []*ltx.FileInfo
+	for itr.Next() {
+		item := itr.Item()
+
+		match, err := filter(item)
+		if match {
+			a = append(a, item)
+		}
+
+		if err == ErrStopIter {
+			break
+		} else if err != nil {
+			return a, err
+		}
+	}
+
+	if err := itr.Close(); err != nil {
+		return nil, err
+	}
+	return a, nil
+}

replica_test.go

@@ -3,9 +3,11 @@ package litestream_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/benbjohnson/litestream"
 	"github.com/benbjohnson/litestream/file"
+	"github.com/benbjohnson/litestream/mock"
 	"github.com/superfly/ltx"
 )
 
@@ -29,8 +31,7 @@ func TestReplica_Sync(t *testing.T) {
 	t.Logf("position after sync: %s", dpos.String())
 
 	c := file.NewReplicaClient(t.TempDir())
-	r := litestream.NewReplica(db)
-	c.Replica, r.Client = r, c
+	r := litestream.NewReplicaWithClient(db, c)
 
 	if err := r.Sync(context.Background()); err != nil {
 		t.Fatal(err)
@@ -82,3 +83,143 @@ func TestReplica_Sync(t *testing.T) {
 
 	// TODO(ltx): Restore snapshot and verify
 }
+
+func TestReplica_CalcRestorePlan(t *testing.T) {
+	db, sqldb := MustOpenDBs(t)
+	defer MustCloseDBs(t, db, sqldb)
+
+	t.Run("SnapshotOnly", func(t *testing.T) {
+		var c mock.ReplicaClient
+		r := litestream.NewReplicaWithClient(db, &c)
+		c.LTXFilesFunc = func(ctx context.Context, level int, seek ltx.TXID) (ltx.FileIterator, error) {
+			if level == litestream.SnapshotLevel {
+				return ltx.NewFileInfoSliceIterator([]*ltx.FileInfo{{
+					Level:     litestream.SnapshotLevel,
+					MinTXID:   1,
+					MaxTXID:   10,
+					Size:      1024,
+					CreatedAt: time.Now(),
+				}}), nil
+			}
+			return ltx.NewFileInfoSliceIterator(nil), nil
+		}
+
+		plan, err := r.CalcRestorePlan(context.Background(), 10)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if got, want := len(plan), 1; got != want {
+			t.Fatalf("n=%d, want %d", got, want)
+		}
+		if plan[0].MaxTXID != 10 {
+			t.Fatalf("expected MaxTXID 10, got %d", plan[0].MaxTXID)
+		}
+	})
+
+	t.Run("SnapshotAndIncremental", func(t *testing.T) {
+		var c mock.ReplicaClient
+		r := litestream.NewReplicaWithClient(db, &c)
+		c.LTXFilesFunc = func(ctx context.Context, level int, seek ltx.TXID) (ltx.FileIterator, error) {
+			switch level {
+			case litestream.SnapshotLevel:
+				return ltx.NewFileInfoSliceIterator([]*ltx.FileInfo{
+					{Level: litestream.SnapshotLevel, MinTXID: 1, MaxTXID: 5},
+					{Level: litestream.SnapshotLevel, MinTXID: 1, MaxTXID: 15},
+				}), nil
+			case 1:
+				return ltx.NewFileInfoSliceIterator([]*ltx.FileInfo{
+					{Level: 1, MinTXID: 6, MaxTXID: 7},
+					{Level: 1, MinTXID: 8, MaxTXID: 9},
+					{Level: 1, MinTXID: 10, MaxTXID: 12},
+				}), nil
+			case 0:
+				return ltx.NewFileInfoSliceIterator([]*ltx.FileInfo{
+					{Level: 0, MinTXID: 7, MaxTXID: 7},
+					{Level: 0, MinTXID: 8, MaxTXID: 8},
+					{Level: 0, MinTXID: 9, MaxTXID: 9},
+					{Level: 0, MinTXID: 10, MaxTXID: 10},
+					{Level: 0, MinTXID: 11, MaxTXID: 11},
+				}), nil
+			default:
+				return ltx.NewFileInfoSliceIterator(nil), nil
+			}
+		}
+
+		plan, err := r.CalcRestorePlan(context.Background(), 10)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if got, want := len(plan), 4; got != want {
+			t.Fatalf("n=%v, want %v", got, want)
+		}
+		if got, want := *plan[0], (ltx.FileInfo{Level: litestream.SnapshotLevel, MinTXID: 1, MaxTXID: 5}); got != want {
+			t.Fatalf("plan[0]=%#v, want %#v", got, want)
+		}
+		if got, want := *plan[1], (ltx.FileInfo{Level: 1, MinTXID: 6, MaxTXID: 7}); got != want {
+			t.Fatalf("plan[1]=%#v, want %#v", got, want)
+		}
+		if got, want := *plan[2], (ltx.FileInfo{Level: 1, MinTXID: 8, MaxTXID: 9}); got != want {
+			t.Fatalf("plan[2]=%#v, want %#v", got, want)
+		}
+		if got, want := *plan[3], (ltx.FileInfo{Level: 0, MinTXID: 10, MaxTXID: 10}); got != want {
+			t.Fatalf("plan[2]=%#v, want %#v", got, want)
+		}
+	})
+
+	t.Run("ErrNonContiguousFiles", func(t *testing.T) {
+		var c mock.ReplicaClient
+		r := litestream.NewReplicaWithClient(db, &c)
+		c.LTXFilesFunc = func(ctx context.Context, level int, seek ltx.TXID) (ltx.FileIterator, error) {
+			switch level {
+			case litestream.SnapshotLevel:
+				return ltx.NewFileInfoSliceIterator([]*ltx.FileInfo{
+					{Level: litestream.SnapshotLevel, MinTXID: 1, MaxTXID: 5},
+				}), nil
+			case 1:
+				return ltx.NewFileInfoSliceIterator([]*ltx.FileInfo{
+					{Level: 1, MinTXID: 8, MaxTXID: 9},
+				}), nil
+			default:
+				return ltx.NewFileInfoSliceIterator(nil), nil
+			}
+		}
+
+		_, err := r.CalcRestorePlan(context.Background(), 10)
+		if err == nil || err.Error() != `non-contiguous transaction files: prev=0000000000000005 filename=0000000000000008-0000000000000009.ltx` {
+			t.Fatalf("unexpected error: %q", err)
+		}
+	})
+
+	t.Run("ErrTxNotAvailable", func(t *testing.T) {
+		var c mock.ReplicaClient
+		r := litestream.NewReplicaWithClient(db, &c)
+		c.LTXFilesFunc = func(ctx context.Context, level int, seek ltx.TXID) (ltx.FileIterator, error) {
+			switch level {
+			case litestream.SnapshotLevel:
+				return ltx.NewFileInfoSliceIterator([]*ltx.FileInfo{
+					{Level: litestream.SnapshotLevel, MinTXID: 1, MaxTXID: 10},
+				}), nil
+			default:
+				return ltx.NewFileInfoSliceIterator(nil), nil
+			}
+		}
+
+		_, err := r.CalcRestorePlan(context.Background(), 5)
+		if err != litestream.ErrTxNotAvailable {
+			t.Fatalf("expected ErrTxNotAvailable, got %v", err)
+		}
+	})
+
+	t.Run("ErrNoFiles", func(t *testing.T) {
+		var c mock.ReplicaClient
+		c.LTXFilesFunc = func(ctx context.Context, level int, seek ltx.TXID) (ltx.FileIterator, error) {
+			return ltx.NewFileInfoSliceIterator(nil), nil
+		}
+		r := litestream.NewReplicaWithClient(db, &c)
+
+		_, err := r.CalcRestorePlan(context.Background(), 5)
+		if err != litestream.ErrTxNotAvailable {
+			t.Fatalf("expected ErrTxNotAvailable, got %v", err)
+		}
+	})
+}

store.go

@@ -20,6 +20,9 @@ var (
 	// since the last compaction time. This is used to prevent frequent
 	// re-compaction when restarting the process.
 	ErrCompactionTooEarly = errors.New("compaction too early")
+
+	// ErrTxNotAvailable is returned when a transaction does not exist.
+	ErrTxNotAvailable = errors.New("transaction not available")
 )
 
 // Store defaults
@@ -81,12 +84,14 @@ func (s *Store) Open(ctx context.Context) error {
 			if lvl.Level == 0 {
 				continue
 			}
+			s.wg.Add(1)
 			go func() {
 				defer s.wg.Done()
 				s.monitorCompactionLevel(s.ctx, lvl)
 			}()
 		}
 
+		s.wg.Add(1)
 		go func() {
 			defer s.wg.Done()
 			s.monitorCompactionLevel(s.ctx, s.SnapshotLevel())