Repository: benbjohnson/litestream (go)

Issue:
Manual checkpoint before snapshot loses WAL position
If a database is not written to by the main application, Litestream can lose its own WAL position by checkpointing before taking a snapshot without writing into the sequence table after. This requires an extremely specific sequence of events to happen.

1. Start Litestream against a new database
2. Litestream snapshots the initial generation
3. Restart Litestream
4. Litestream snapshots <- this is when the only page in the WAL is checkpointed away
5. Write to database
6. Litestream complains that the WAL was overwritten by another process, starts a new generation

The issue happens here:
https://github.com/benbjohnson/litestream/blob/85ddf322258b96dbde40772661a39e8a7d5a32d9/replica.go#L484-L487

Normally when checkpointing, Litestream will issue a write to the sequence table to force at least one page to be there after restarting the long read transaction which will ensure that regardless of what checkpoint is done that page will stay in until the next "full" checkpoint is made by Litestream and the sequence table is written again.

In this case the read transaction keeps a lock on the page in the WAL at step 2 but at step 4 the read transaction had been restarted after the page was written during previous incarnation of Litestream which makes it free to be checkpointed from the WAL to the main database. Because snapshot will do that rogue checkpoint it will move it silently out without Litestream realizing it until the next write which checks the last page doesn't match the shadow WAL.

The intention of the code is to call `r.db.Checkpoint(ctx, CheckpointModePassive)` rather than do it manually as it will do the required write to sequence table and also make sure the database file is as up-to-date as possible before taking the snapshot.

Pull request description:
Use safe checkpointing before snapshots
The intention of the code is to do a checkpoint before starting to write out the snapshot. Executing a manual checkpoint will at worst make Litestream lose the track of the WAL as it will in specific circumtances flush out the single page in the WAL file out.

Using the public checkpoint function of the database to issue a managed passive checkpoint will ensure the sequence table is written after and the position of the WAL isn't lost.

Fixes #521

Patch:
diff --git a/replica.go b/replica.go
index 3d2b099..a5daf2d 100644
--- a/replica.go
+++ b/replica.go
@@ -482,16 +482,16 @@ func (r *Replica) Snapshot(ctx context.Context) (info SnapshotInfo, err error) {
 	r.muf.Lock()
 	defer r.muf.Unlock()
 
-	// Prevent checkpoints during snapshot.
-	r.db.BeginSnapshot()
-	defer r.db.EndSnapshot()
-
 	// Issue a passive checkpoint to flush any pages to disk before snapshotting.
-	if _, err := r.db.db.ExecContext(ctx, `PRAGMA wal_checkpoint(PASSIVE);`); err != nil {
+	if err := r.db.Checkpoint(ctx, CheckpointModePassive); err != nil {
 		return info, fmt.Errorf("pre-snapshot checkpoint: %w", err)
 	}
 
-	// Acquire a read lock on the database during snapshot to prevent checkpoints.
+	// Prevent internal checkpoints during snapshot.
+	r.db.BeginSnapshot()
+	defer r.db.EndSnapshot()
+
+	// Acquire a read lock on the database during snapshot to prevent external checkpoints.
 	tx, err := r.db.db.Begin()
 	if err != nil {
 		return info, err


Test patch:
diff --git a/replica_test.go b/replica_test.go
index 7f42c08..78ba620 100644
--- a/replica_test.go
+++ b/replica_test.go
@@ -13,6 +13,13 @@ import (
 	"github.com/pierrec/lz4/v4"
 )
 
+func nextIndex(pos litestream.Pos) litestream.Pos {
+	return litestream.Pos{
+		Generation: pos.Generation,
+		Index:      pos.Index + 1,
+	}
+}
+
 func TestReplica_Name(t *testing.T) {
 	t.Run("WithName", func(t *testing.T) {
 		if got, want := litestream.NewReplica(nil, "NAME").Name(), "NAME"; got != want {
@@ -32,11 +39,6 @@ func TestReplica_Sync(t *testing.T) {
 	db, sqldb := MustOpenDBs(t)
 	defer MustCloseDBs(t, db, sqldb)
 
-	// Execute a query to force a write to the WAL.
-	if _, err := sqldb.Exec(`CREATE TABLE foo (bar TEXT);`); err != nil {
-		t.Fatal(err)
-	}
-
 	// Issue initial database sync to setup generation.
 	if err := db.Sync(context.Background()); err != nil {
 		t.Fatal(err)
@@ -66,10 +68,47 @@ func TestReplica_Sync(t *testing.T) {
 		t.Fatalf("generations[0]=%v, want %v", got, want)
 	}
 
+	// Verify we synced checkpoint page to WAL.
+	if r, err := c.WALSegmentReader(context.Background(), nextIndex(dpos)); err != nil {
+		t.Fatal(err)
+	} else if b, err := io.ReadAll(lz4.NewReader(r)); err != nil {
+		t.Fatal(err)
+	} else if err := r.Close(); err != nil {
+		t.Fatal(err)
+	} else if len(b) == db.PageSize() {
+		t.Fatalf("wal mismatch: len(%d), len(%d)", len(b), db.PageSize())
+	}
+
+	// Reset WAL so the next write will only write out the segment we are checking.
+	if err := db.Checkpoint(context.Background(), litestream.CheckpointModeTruncate); err != nil {
+		t.Fatal(err)
+	}
+
+	// Execute a query to write something into the truncated WAL.
+	if _, err := sqldb.Exec(`CREATE TABLE foo (bar TEXT);`); err != nil {
+		t.Fatal(err)
+	}
+
+	// Sync database to catch up the shadow WAL.
+	if err := db.Sync(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	// Save position after sync, it should be after our write.
+	dpos, err = db.Pos()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Sync WAL segment out to replica.
+	if err := r.Sync(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
 	// Verify WAL matches replica WAL.
 	if b0, err := os.ReadFile(db.Path() + "-wal"); err != nil {
 		t.Fatal(err)
-	} else if r, err := c.WALSegmentReader(context.Background(), litestream.Pos{Generation: generations[0], Index: 0, Offset: 0}); err != nil {
+	} else if r, err := c.WALSegmentReader(context.Background(), dpos.Truncate()); err != nil {
 		t.Fatal(err)
 	} else if b1, err := io.ReadAll(lz4.NewReader(r)); err != nil {
 		t.Fatal(err)
@@ -103,7 +142,7 @@ func TestReplica_Snapshot(t *testing.T) {
 		t.Fatal(err)
 	} else if info, err := r.Snapshot(context.Background()); err != nil {
 		t.Fatal(err)
-	} else if got, want := info.Pos(), pos0.Truncate(); got != want {
+	} else if got, want := info.Pos(), nextIndex(pos0); got != want {
 		t.Fatalf("pos=%s, want %s", got, want)
 	}
 
@@ -127,18 +166,20 @@ func TestReplica_Snapshot(t *testing.T) {
 		t.Fatal(err)
 	} else if info, err := r.Snapshot(context.Background()); err != nil {
 		t.Fatal(err)
-	} else if got, want := info.Pos(), pos1.Truncate(); got != want {
+	} else if got, want := info.Pos(), nextIndex(pos1); got != want {
 		t.Fatalf("pos=%v, want %v", got, want)
 	}
 
-	// Verify two snapshots exist.
+	// Verify three snapshots exist.
 	if infos, err := r.Snapshots(context.Background()); err != nil {
 		t.Fatal(err)
-	} else if got, want := len(infos), 2; got != want {
+	} else if got, want := len(infos), 3; got != want {
 		t.Fatalf("len=%v, want %v", got, want)
 	} else if got, want := infos[0].Pos(), pos0.Truncate(); got != want {
 		t.Fatalf("info[0]=%s, want %s", got, want)
-	} else if got, want := infos[1].Pos(), pos1.Truncate(); got != want {
+	} else if got, want := infos[1].Pos(), nextIndex(pos0); got != want {
 		t.Fatalf("info[1]=%s, want %s", got, want)
+	} else if got, want := infos[2].Pos(), nextIndex(pos1); got != want {
+		t.Fatalf("info[2]=%s, want %s", got, want)
 	}
 }