Repository: polarsignals/frostdb (go)

Issue:
db: handle cases where WAL last index != snapshot txn
There might be cases where the WAL hasn't caught up to a given snapshot txn for whatever reason (e.g. WAL corruption causes an end truncation). In cases where the WAL's last index is < the snapshot txn during recovery, we should remove all WAL files and create a new WAL file with the given snapshot's last index. This will ensure that the WAL can be appended to.

Pull request description:
wal: allow Truncating past end of WAL
This PR implements the logic to allow truncating past the end of the WAL,
i.e. resetting the WAL entirely. The only condition is that the next record
logged MUST be the txn immediately following the txn at which the WAL was
truncated.

This is very useful for recovery where we have a snapshot with a txn k and the
last index of a WAL is j < k-1, so we don't really care about the WAL entries
anymore.

Closes #412 

Patch:
diff --git a/db.go b/db.go
index 496a4d1..71c9940 100644
--- a/db.go
+++ b/db.go
@@ -544,18 +544,6 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 		)
 		snapshotTx.TxnID = 0
 	}
-	firstIndex, err := wal.FirstIndex()
-	if err != nil {
-		level.Info(db.logger).Log(
-			"msg", "failed to get WAL first index",
-			"err", err)
-	}
-	lastIndex, err := wal.LastIndex()
-	if err != nil {
-		level.Info(db.logger).Log(
-			"msg", "failed to get WAL last index",
-			"err", err)
-	}
 	snapshotLogArgs := make([]any, 0)
 	if snapshotTx.TxnID != 0 {
 		snapshotLogArgs = append(
@@ -571,30 +559,7 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 				"snapshot_tx", snapshotTx,
 			)
 		}
-		if snapshotTx.TxnID > firstIndex && snapshotTx.TxnID <= lastIndex {
-			if err := wal.Truncate(snapshotTx.TxnID); err != nil {
-				// Since this is a best-effort truncation, move on if there is an
-				// error.
-				level.Info(db.logger).Log(
-					"msg", "WAL truncation after successful snapshot load encountered error",
-					"err", err,
-					"first_index", firstIndex,
-					"last_index", lastIndex,
-					"snapshot_tx", snapshotTx,
-				)
-			} else {
-				snapshotLogArgs = append(
-					snapshotLogArgs,
-					"wal_truncated", "true",
-					"wal_first_index_pre_truncation", firstIndex)
-				firstIndex = snapshotTx.TxnID
-			}
-		} else {
-			snapshotLogArgs = append(
-				snapshotLogArgs,
-				"wal_truncated", "false",
-			)
-		}
+		db.resetToTxn(snapshotTx, wal)
 	}
 
 	// persistedTables is a map from a table name to the last transaction
@@ -771,10 +736,11 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 				}
 			}
 
-			// After every insert we're setting the tx and highWatermark to the replayed tx.
-			// This allows the block's compaction to start working on the inserted data.
-			db.tx.Store(tx)
-			db.highWatermark.Store(Txn{TxnID: tx, TxnMetadata: record.TxnMetadata})
+			// After every insert we're setting the tx and highWatermark to the
+			// replayed tx.
+			// This allows the block's compaction to start working on the
+			// inserted data.
+			db.resetToTxn(Txn{TxnID: tx, TxnMetadata: record.TxnMetadata}, nil)
 		case *walpb.Entry_TableBlockPersisted_:
 			return nil
 		case *walpb.Entry_Snapshot_:
@@ -788,23 +754,12 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 	}
 
 	if lastTx.TxnID >= snapshotTx.TxnID {
-		db.tx.Store(lastTx.TxnID)
-		db.highWatermark.Store(lastTx)
-	}
-	if lastTxn := db.tx.Load(); lastTxn != lastIndex {
-		level.Warn(db.logger).Log(
-			"msg", "WAL last index is != db last txn, won't be able to log records to WAL",
-			"wal_last_index", lastIndex,
-			"last_tx", lastTxn,
-		)
+		db.resetToTxn(lastTx, nil)
 	}
-
 	level.Info(db.logger).Log(
 		append(
 			[]any{
 				"msg", "db recovered",
-				"wal_first_index", firstIndex,
-				"wal_last_index", lastIndex,
 				"wal_replay_duration", time.Since(start),
 			},
 			snapshotLogArgs...,
@@ -1127,6 +1082,25 @@ func (db *DB) HighWatermark() Txn {
 	return db.highWatermark.Load()
 }
 
+// resetToTxn resets the DB's internal state to resume from the given
+// transaction. If the given wal is non-nil, it is also reset so that the next
+// expected transaction will log correctly to the WAL. Note that db.wal is not
+// used since callers might be calling resetToTxn before db.wal has been
+// initialized or might not want the WAL to be reset.
+func (db *DB) resetToTxn(txn Txn, wal WAL) {
+	db.tx.Store(txn.TxnID)
+	db.highWatermark.Store(txn)
+	if wal != nil {
+		if err := wal.Truncate(txn.TxnID); err != nil {
+			level.Warn(db.logger).Log(
+				"msg", "failed to truncate WAL when resetting DB to txn",
+				"txnID", txn.TxnID,
+				"err", err,
+			)
+		}
+	}
+}
+
 // validateName ensures that the passed in name doesn't violate any constrainsts.
 func validateName(name string) bool {
 	return !strings.Contains(name, "/")
diff --git a/go.mod b/go.mod
index e5a2b35..351c6d1 100644
--- a/go.mod
+++ b/go.mod
@@ -16,7 +16,7 @@ require (
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/parquet-go/parquet-go v0.0.0-20230717201923-bb12c19a1110
 	github.com/pingcap/tidb/parser v0.0.0-20220921115303-5aab87679fde
-	github.com/polarsignals/wal v0.0.0-20230509073041-6731e243de20
+	github.com/polarsignals/wal v0.0.0-20230809145250-99ef415c5a6c
 	github.com/prometheus/client_golang v1.16.0
 	github.com/stretchr/testify v1.8.3
 	github.com/thanos-io/objstore v0.0.0-20230713070940-eb01c83b89a4
diff --git a/go.sum b/go.sum
index e41bb01..5c1d8ed 100644
--- a/go.sum
+++ b/go.sum
@@ -92,8 +92,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/polarsignals/wal v0.0.0-20230509073041-6731e243de20 h1:YverywcwPHYj5iQq1/Zy/EihWVyw2F0S7j52FvGWK4w=
-github.com/polarsignals/wal v0.0.0-20230509073041-6731e243de20/go.mod h1:EVDHAAe+7GQ33A1/x+/gE+sBPN4toQ0XG5RoLD49xr8=
+github.com/polarsignals/wal v0.0.0-20230809145250-99ef415c5a6c h1:fVYVdHe8YeIXkmYr+JRNtrobo53lYthop0jXqS6Xamk=
+github.com/polarsignals/wal v0.0.0-20230809145250-99ef415c5a6c/go.mod h1:EVDHAAe+7GQ33A1/x+/gE+sBPN4toQ0XG5RoLD49xr8=
 github.com/prometheus/client_golang v1.16.0 h1:yk/hx9hDbrGHovbci4BY+pRMfSuuat626eFsHb7tmT8=
 github.com/prometheus/client_golang v1.16.0/go.mod h1:Zsulrv/L9oM40tJ7T815tM89lFEugiJ9HzIqaAx4LKc=
 github.com/prometheus/client_model v0.4.0 h1:5lQXD3cAg1OXBf4Wq03gTrXHeaV0TQvGfUooCfx1yqY=
diff --git a/snapshot.go b/snapshot.go
index f7a315b..0412c85 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -247,7 +247,7 @@ func (db *DB) loadLatestSnapshotFromDir(ctx context.Context, dir string) (Txn, e
 			if err != nil {
 				return err
 			}
-			watermark, err := LoadSnapshot(ctx, db, parsedTx, f, info.Size())
+			watermark, err := LoadSnapshot(ctx, db, parsedTx, f, info.Size(), false)
 			if err != nil {
 				return err
 			}
@@ -277,14 +277,17 @@ func (db *DB) loadLatestSnapshotFromDir(ctx context.Context, dir string) (Txn, e
 	return Txn{}, fmt.Errorf("%s", errString)
 }
 
-func LoadSnapshot(ctx context.Context, db *DB, tx uint64, r io.ReaderAt, size int64) (Txn, error) {
+func LoadSnapshot(ctx context.Context, db *DB, tx uint64, r io.ReaderAt, size int64, truncateWAL bool) (Txn, error) {
 	txnMetadata, err := loadSnapshot(ctx, db, r, size)
 	if err != nil {
 		return Txn{}, err
 	}
-	db.tx.Store(tx)
 	watermark := Txn{TxnID: tx, TxnMetadata: txnMetadata}
-	db.highWatermark.Store(watermark)
+	var wal WAL
+	if truncateWAL {
+		wal = db.wal
+	}
+	db.resetToTxn(watermark, wal)
 	return watermark, nil
 }
 
diff --git a/wal/wal.go b/wal/wal.go
index 7623b3b..b7ac1d9 100644
--- a/wal/wal.go
+++ b/wal/wal.go
@@ -79,6 +79,9 @@ type FileWAL struct {
 		// then perform a truncation. If another truncate call occurs in the
 		// meantime, the highest txn will be used.
 		truncateTx uint64
+		// nextTx is the next expected txn. The FileWAL will only log a record
+		// with this txn.
+		nextTx uint64
 	}
 
 	cancel     func()
@@ -127,6 +130,11 @@ func Open(
 		return nil, err
 	}
 
+	lastIndex, err := logStore.LastIndex()
+	if err != nil {
+		return nil, err
+	}
+
 	w := &FileWAL{
 		logger:       logger,
 		path:         path,
@@ -165,6 +173,8 @@ func Open(
 		shutdownCh: make(chan struct{}),
 	}
 
+	w.protected.nextTx = lastIndex + 1
+
 	return w, nil
 }
 
@@ -207,26 +217,16 @@ func (w *FileWAL) run(ctx context.Context) {
 			return
 		case <-ticker.C:
 			batch := batch[:0]
-			nextTx, err := w.log.LastIndex()
-			if err != nil {
-				// TODO(asubiotto): Should we do something else here? This
-				// most likely indicates that the WAL is corrupt.
-				level.Error(w.logger).Log(
-					"msg", "WAL failed to get last index",
-					"err", err,
-				)
-			}
-			nextTx++
 			w.protected.Lock()
 			for w.protected.queue.Len() > 0 {
-				if minTx := w.protected.queue[0].tx; minTx != nextTx {
-					if minTx < nextTx {
+				if minTx := w.protected.queue[0].tx; minTx != w.protected.nextTx {
+					if minTx < w.protected.nextTx {
 						// The next entry must be dropped otherwise progress
 						// will never be made. Log a warning given this could
 						// lead to missing data.
 						level.Warn(w.logger).Log(
 							"msg", "WAL cannot log a txn id that has already been seen; dropping entry",
-							"expected", nextTx,
+							"expected", w.protected.nextTx,
 							"found", minTx,
 						)
 						_ = heap.Pop(&w.protected.queue)
@@ -235,13 +235,13 @@ func (w *FileWAL) run(ctx context.Context) {
 				}
 				r := heap.Pop(&w.protected.queue).(*logRequest)
 				batch = append(batch, r)
-				nextTx++
+				w.protected.nextTx++
 			}
 			// truncateTx will be non-zero if we either are about to log a
 			// record with a txn past the txn to truncate, or we have logged one
 			// in the past.
 			truncateTx := uint64(0)
-			if w.protected.truncateTx < nextTx {
+			if w.protected.truncateTx != 0 {
 				truncateTx = w.protected.truncateTx
 				w.protected.truncateTx = 0
 			}
@@ -280,6 +280,17 @@ func (w *FileWAL) run(ctx context.Context) {
 				if err := w.log.TruncateFront(truncateTx); err != nil {
 					level.Error(w.logger).Log("msg", "failed to truncate WAL", "tx", truncateTx, "err", err)
 				} else {
+					w.protected.Lock()
+					if truncateTx > w.protected.nextTx {
+						// truncateTx is the new firstIndex of the WAL. If it is
+						// greater than the next expected transaction, this was
+						// a full WAL truncation/reset so both the first and
+						// last index are now 0. The underlying WAL will allow a
+						// record with any index to be written, however we only
+						// want to allow the next index to be logged.
+						w.protected.nextTx = truncateTx + 1
+					}
+					w.protected.Unlock()
 					level.Debug(w.logger).Log("msg", "truncated WAL", "tx", truncateTx)
 				}
 			}


Test patch:
diff --git a/wal/wal_test.go b/wal/wal_test.go
index cc08563..17e090c 100644
--- a/wal/wal_test.go
+++ b/wal/wal_test.go
@@ -159,6 +159,18 @@ func TestUnexpectedTxn(t *testing.T) {
 }
 
 func TestWALTruncate(t *testing.T) {
+	logRecord := func(data string) *walpb.Record {
+		return &walpb.Record{
+			Entry: &walpb.Entry{
+				EntryType: &walpb.Entry_Write_{
+					Write: &walpb.Entry_Write{
+						Data:      []byte(data),
+						TableName: "test-table",
+					},
+				},
+			},
+		}
+	}
 	for i, tc := range []string{"BeforeLog", "AfterLog"} {
 		t.Run(tc, func(t *testing.T) {
 			dir := t.TempDir()
@@ -171,17 +183,8 @@ func TestWALTruncate(t *testing.T) {
 			defer w.Close()
 			w.RunAsync()
 
-			for i := uint64(0); i < 10; i++ {
-				require.NoError(t, w.Log(i, &walpb.Record{
-					Entry: &walpb.Entry{
-						EntryType: &walpb.Entry_Write_{
-							Write: &walpb.Entry_Write{
-								Data:      []byte(fmt.Sprintf("test-data-%d", i)),
-								TableName: "test-table",
-							},
-						},
-					},
-				}))
+			for j := uint64(1); j < 10; j++ {
+				require.NoError(t, w.Log(j, logRecord(fmt.Sprintf("test-data-%d", j))))
 			}
 			if i == 1 {
 				// Wait until the last entry is written before issuing the
@@ -212,4 +215,57 @@ func TestWALTruncate(t *testing.T) {
 			require.Equal(t, 1, numRecords)
 		})
 	}
+
+	t.Run("Reset", func(t *testing.T) {
+		dir := t.TempDir()
+		w, err := Open(
+			log.NewNopLogger(),
+			prometheus.NewRegistry(),
+			dir,
+		)
+		require.NoError(t, err)
+		defer w.Close()
+		w.RunAsync()
+
+		for i := uint64(1); i < 10; i++ {
+			require.NoError(t, w.Log(i, logRecord("test-data-%d")))
+		}
+
+		// Truncate way past the current last index, which should be 9.
+		const truncateIdx = 20
+		require.NoError(t, w.Truncate(truncateIdx))
+
+		require.Eventually(t, func() bool {
+			first, _ := w.FirstIndex()
+			last, _ := w.LastIndex()
+			return first == 0 && last == 0
+		}, time.Second, 10*time.Millisecond)
+
+		// Even though the WAL has been reset, we should not allow logging a
+		// record with a txn lower than the last truncation. This will only
+		// be observed on replay.
+		require.NoError(t, w.Log(1, logRecord("should-not-be-logged")))
+
+		// The only valid record to log is the one after truncateIdx.
+		require.NoError(t, w.Log(truncateIdx+1, logRecord("should-be-logged")))
+
+		// Wait for record to be logged.
+		require.Eventually(t, func() bool {
+			first, _ := w.FirstIndex()
+			last, _ := w.LastIndex()
+			return first == truncateIdx+1 && last == truncateIdx+1
+		}, time.Second, 10*time.Millisecond)
+
+		numRecords := 0
+		require.NoError(
+			t,
+			w.Replay(0, func(tx uint64, r *walpb.Record) error {
+				numRecords++
+				require.Equal(t, uint64(truncateIdx+1), tx)
+				require.Equal(t, []byte("should-be-logged"), r.Entry.GetWrite().Data)
+				return nil
+			}),
+		)
+		require.Equal(t, 1, numRecords)
+	})
 }