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.


Discussion:

Patch:
diff --git a/db.go b/db.go
--- 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
--- 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
--- 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
--- 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
--- 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
--- 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)
+	})
 }