Repository: polarsignals/frostdb (go)

Issue:
wal: make replaying part of opening the WAL
Currently, callers need to explicitly call `ReplayWAL` to replay a write ahead log. This should be part of opening a WAL. Given the caller needs to provide a replay function, this can be passed in to `wal.Open` by the db. If no replay function is provided, I suggest we just nuke all the files in the WAL.

Pull request description:
Replay wal on open
In efforts to make the WAL/replay more robust I've added two things here:

- On wal read panics instead of continuing the panic we instead truncate the WAL up to the last good tx. Using some corrupted wals from production this works really well.
- On DB open (and wal open) replay the WAL. This is mostly a refactor, but it does remove some weird footguns where you could open a DB with a WAL and not replay it, or write to it while you were replaying it.

Closes #328 

Patch:
diff --git a/db.go b/db.go
index 1463af9..cb22868 100644
--- a/db.go
+++ b/db.go
@@ -103,6 +103,10 @@ func New(
 		return nil, fmt.Errorf("storage path must be configured if WAL is enabled")
 	}
 
+	if err := s.recoverDBsFromStorage(context.Background()); err != nil {
+		return nil, err
+	}
+
 	return s, nil
 }
 
@@ -235,8 +239,8 @@ func (s *ColumnStore) DatabasesDir() string {
 	return filepath.Join(s.storagePath, "databases")
 }
 
-// ReplayWALs replays the write-ahead log for each database.
-func (s *ColumnStore) ReplayWALs(ctx context.Context) error {
+// recoverDBsFromStorage replays the snapshots and write-ahead logs for each database.
+func (s *ColumnStore) recoverDBsFromStorage(ctx context.Context) error {
 	if !s.enableWAL {
 		return nil
 	}
@@ -259,11 +263,9 @@ func (s *ColumnStore) ReplayWALs(ctx context.Context) error {
 	for _, f := range files {
 		databaseName := f.Name()
 		g.Go(func() error {
-			db, err := s.DB(ctx, databaseName)
-			if err != nil {
-				return err
-			}
-			return db.recover(ctx)
+			// Open the DB for the side effect of the snapshot and WALs being loaded as part of the open operation.
+			_, err := s.DB(ctx, databaseName)
+			return err
 		})
 	}
 
@@ -350,14 +352,6 @@ func (s *ColumnStore) DB(ctx context.Context, name string) (*DB, error) {
 		db.bucket = storage.NewPrefixedBucket(s.bucket, db.name)
 	}
 
-	if s.enableWAL {
-		var err error
-		db.wal, err = db.openWAL()
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	if dbSetupErr := func() error {
 		db.txPool = NewTxPool(db.highWatermark)
 		db.compactorPool = newCompactorPool(db, s.compactionConfig)
@@ -375,6 +369,20 @@ func (s *ColumnStore) DB(ctx context.Context, name string) (*DB, error) {
 			}
 		}
 
+		if s.enableWAL {
+			snapshotTx, err := db.loadLatestSnapshot(ctx)
+			if err != nil {
+				level.Debug(s.logger).Log(
+					"msg", "failed to load latest snapshot", "db", db.name, "err", err,
+				)
+				snapshotTx = 0
+			}
+			db.wal, err = db.openWAL(ctx, snapshotTx)
+			if err != nil {
+				return err
+			}
+		}
+
 		// Register metrics last to avoid duplicate registration should and of the WAL or storage replay errors occur
 		db.metrics = &dbMetrics{
 			txHighWatermark: promauto.With(reg).NewGaugeFunc(prometheus.GaugeOpts{
@@ -404,12 +412,22 @@ func (s *ColumnStore) DB(ctx context.Context, name string) (*DB, error) {
 	return db, nil
 }
 
-func (db *DB) openWAL() (WAL, error) {
-	return wal.Open(
+func (db *DB) openWAL(ctx context.Context, firstIndex uint64) (WAL, error) {
+	wal, err := wal.Open(
 		db.logger,
 		db.reg,
 		db.walDir(),
 	)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := db.recover(ctx, wal); err != nil {
+		return nil, err
+	}
+
+	wal.RunAsync()
+	return wal, nil
 }
 
 func (db *DB) walDir() string {
@@ -422,7 +440,7 @@ func (db *DB) snapshotsDir() string {
 
 // recover attempts to recover database state from a combination of snapshots
 // and the WAL.
-func (db *DB) recover(ctx context.Context) error {
+func (db *DB) recover(ctx context.Context, wal WAL) error {
 	level.Info(db.logger).Log("msg", "recovering db")
 	snapshotLoadStart := time.Now()
 	snapshotTx, err := db.loadLatestSnapshot(ctx)
@@ -454,13 +472,15 @@ func (db *DB) recover(ctx context.Context) error {
 	}
 	level.Info(db.logger).Log("msg", "replaying WAL", "snapshot_tx", snapshotTx)
 
-	firstIndex, err := db.wal.FirstIndex()
+	firstIndex, err := wal.FirstIndex()
 	if err != nil {
-		return err
+		level.Info(db.logger).Log(
+			"msg", "failed to get WAL first index",
+			"err", err)
 	}
 
 	if snapshotTx > firstIndex {
-		if err := db.wal.Truncate(snapshotTx); err != nil {
+		if err := wal.Truncate(snapshotTx); err != nil {
 			// Since this is a best-effort truncation, move on if there is an
 			// error.
 			level.Info(db.logger).Log(
@@ -477,11 +497,13 @@ func (db *DB) recover(ctx context.Context) error {
 		}
 	}
 
-	start := time.Now()
 	// persistedTables is a map from a table name to the last transaction
 	// persisted.
 	persistedTables := map[string]uint64{}
-	if err := db.wal.Replay(snapshotTx, func(tx uint64, record *walpb.Record) error {
+	lastTx := uint64(0)
+
+	start := time.Now()
+	if err := wal.Replay(firstIndex, func(tx uint64, record *walpb.Record) error {
 		switch e := record.Entry.EntryType.(type) {
 		case *walpb.Entry_TableBlockPersisted_:
 			persistedTables[e.TableBlockPersisted.TableName] = tx
@@ -495,14 +517,15 @@ func (db *DB) recover(ctx context.Context) error {
 				}
 				db.mtx.Unlock()
 			}
+			return nil
+		default:
+			return nil
 		}
-		return nil
 	}); err != nil {
-		return fmt.Errorf("first WAL replay: %w", err)
+		return err
 	}
 
-	lastTx := uint64(0)
-	if err := db.wal.Replay(firstIndex, func(tx uint64, record *walpb.Record) error {
+	if err := wal.Replay(firstIndex, func(tx uint64, record *walpb.Record) error {
 		lastTx = tx
 		switch e := record.Entry.EntryType.(type) {
 		case *walpb.Entry_NewTableBlock_:
@@ -552,7 +575,7 @@ func (db *DB) recover(ctx context.Context) error {
 						db.reg,
 						db.logger,
 						db.tracer,
-						db.wal,
+						wal,
 					)
 				}
 				if err != nil {
@@ -652,14 +675,12 @@ func (db *DB) recover(ctx context.Context) error {
 		}
 		return nil
 	}); err != nil {
-		return fmt.Errorf("second WAL replay: %w", err)
+		return err
 	}
 
-	level.Info(db.logger).Log("msg", "replaying WAL completed", "duration", time.Since(start))
-
 	db.tx.Store(lastTx)
 	db.highWatermark.Store(lastTx)
-
+	level.Info(db.logger).Log("msg", "replaying WAL completed", "duration", time.Since(start))
 	return nil
 }
 
diff --git a/table.go b/table.go
index 15e4d2c..55156dc 100644
--- a/table.go
+++ b/table.go
@@ -41,6 +41,7 @@ import (
 	"github.com/polarsignals/frostdb/parts"
 	"github.com/polarsignals/frostdb/pqarrow"
 	"github.com/polarsignals/frostdb/query/logicalplan"
+	"github.com/polarsignals/frostdb/wal"
 	walpkg "github.com/polarsignals/frostdb/wal"
 )
 
@@ -146,7 +147,7 @@ type WAL interface {
 	// 0, the first index read from the WAL is used (i.e. given a truncation,
 	// using 0 is still valid). If the given firstIndex is less than the WAL's
 	// first index on disk, the replay happens from the first index on disk.
-	Replay(firstIndex uint64, handler func(tx uint64, record *walpb.Record) error) error
+	Replay(tx uint64, handler wal.ReplayHandlerFunc) error
 	Truncate(tx uint64) error
 	FirstIndex() (uint64, error)
 }
diff --git a/wal/wal.go b/wal/wal.go
index 249ffd0..e10f721 100644
--- a/wal/wal.go
+++ b/wal/wal.go
@@ -22,6 +22,8 @@ import (
 	walpb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/wal/v1alpha1"
 )
 
+type ReplayHandlerFunc func(tx uint64, record *walpb.Record) error
+
 type NopWAL struct{}
 
 func (w *NopWAL) Close() error {
@@ -32,11 +34,11 @@ func (w *NopWAL) Log(tx uint64, record *walpb.Record) error {
 	return nil
 }
 
-func (w *NopWAL) LogRecord(tx uint64, table string, record arrow.Record) error {
+func (w *NopWAL) Replay(tx uint64, handler ReplayHandlerFunc) error {
 	return nil
 }
 
-func (w *NopWAL) Replay(firstIndex uint64, handler func(tx uint64, record *walpb.Record) error) error {
+func (w *NopWAL) LogRecord(tx uint64, table string, record arrow.Record) error {
 	return nil
 }
 
@@ -168,13 +170,6 @@ func Open(
 		shutdownCh: make(chan struct{}),
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	w.cancel = cancel
-	go func() {
-		w.run(ctx)
-		close(w.shutdownCh)
-	}()
-
 	return w, nil
 }
 
@@ -297,6 +292,9 @@ func (w *FileWAL) Truncate(tx uint64) error {
 }
 
 func (w *FileWAL) Close() error {
+	if w.cancel == nil { // wal was never started
+		return nil
+	}
 	w.cancel()
 	<-w.shutdownCh
 	return w.log.Close()
@@ -387,13 +385,17 @@ func (w *FileWAL) LastIndex() (uint64, error) {
 	return w.log.LastIndex()
 }
 
-func (w *FileWAL) Replay(firstIndex uint64, handler func(tx uint64, record *walpb.Record) error) error {
+func (w *FileWAL) Replay(tx uint64, handler ReplayHandlerFunc) (err error) {
+	if handler == nil { // no handler provided
+		return nil
+	}
+
 	logFirstIndex, err := w.log.FirstIndex()
 	if err != nil {
 		return fmt.Errorf("read first index: %w", err)
 	}
-	if firstIndex == 0 || firstIndex < logFirstIndex {
-		firstIndex = logFirstIndex
+	if tx == 0 || tx < logFirstIndex {
+		tx = logFirstIndex
 	}
 
 	lastIndex, err := w.log.LastIndex()
@@ -402,21 +404,26 @@ func (w *FileWAL) Replay(firstIndex uint64, handler func(tx uint64, record *walp
 	}
 
 	// FirstIndex and LastIndex returns zero when there is no WAL files.
-	if firstIndex == 0 || lastIndex == 0 {
+	if tx == 0 || lastIndex == 0 {
 		return nil
 	}
 
-	level.Debug(w.logger).Log("msg", "replaying WAL", "first_index", firstIndex, "last_index", lastIndex)
+	level.Debug(w.logger).Log("msg", "replaying WAL", "first_index", tx, "last_index", lastIndex)
 
 	defer func() {
-		// recover the panic to print more context. Exit afterwards regardless.
-		if err := recover(); err != nil {
-			level.Error(w.logger).Log("msg", "replaying WAL failed", "path", w.path, "first_index", firstIndex, "last_index", lastIndex, "err", err)
-			panic(err)
+		// recover a panic of reading a transaction. Truncate the wal to the last valid transaction.
+		if r := recover(); r != nil {
+			level.Error(w.logger).Log("msg", "replaying WAL failed", "path", w.path, "first_index", tx, "last_index", lastIndex, "offending_index", tx, "err", err)
+			if err = w.log.TruncateBack(tx - 1); err != nil {
+				return
+			}
+			w.txmtx.Lock()
+			w.nextTx = tx
+			w.txmtx.Unlock()
 		}
 	}()
 
-	for tx := firstIndex; tx <= lastIndex; tx++ {
+	for ; tx <= lastIndex; tx++ {
 		level.Debug(w.logger).Log("msg", "replaying WAL record", "tx", tx)
 		data, err := w.log.Read(tx)
 		if err != nil {
@@ -436,5 +443,15 @@ func (w *FileWAL) Replay(firstIndex uint64, handler func(tx uint64, record *walp
 	w.txmtx.Lock()
 	w.nextTx = lastIndex + 1
 	w.txmtx.Unlock()
+
 	return nil
 }
+
+func (w *FileWAL) RunAsync() {
+	ctx, cancel := context.WithCancel(context.Background())
+	w.cancel = cancel
+	go func() {
+		w.run(ctx)
+		close(w.shutdownCh)
+	}()
+}


Test patch:
diff --git a/bench_test.go b/bench_test.go
index 2f2a9c3..34dcf48 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -38,6 +38,8 @@ const (
 func newDBForBenchmarks(ctx context.Context, b testing.TB) (*ColumnStore, *DB, error) {
 	b.Helper()
 
+	b.Logf("replaying WAL")
+	start := time.Now()
 	col, err := New(
 		WithWAL(),
 		WithStoragePath(storagePath),
@@ -45,12 +47,6 @@ func newDBForBenchmarks(ctx context.Context, b testing.TB) (*ColumnStore, *DB, e
 	if err != nil {
 		return nil, nil, err
 	}
-
-	b.Logf("replaying WAL")
-	start := time.Now()
-	if err := col.ReplayWALs(ctx); err != nil {
-		return nil, nil, err
-	}
 	b.Logf("replayed WAL in %s", time.Since(start))
 
 	colDB, err := col.DB(ctx, dbName)
@@ -363,9 +359,6 @@ func BenchmarkReplay(b *testing.B) {
 			)
 			require.NoError(b, err)
 			defer col.Close()
-			if err := col.ReplayWALs(context.Background()); err != nil {
-				b.Fatal(err)
-			}
 		}()
 	}
 }
diff --git a/db_test.go b/db_test.go
index 2409f14..e06791e 100644
--- a/db_test.go
+++ b/db_test.go
@@ -72,7 +72,6 @@ func TestDBWithWALAndBucket(t *testing.T) {
 	)
 	require.NoError(t, err)
 	defer c.Close()
-	require.NoError(t, c.ReplayWALs(context.Background()))
 }
 
 func TestDBWithWAL(t *testing.T) {
@@ -252,8 +251,6 @@ func TestDBWithWAL(t *testing.T) {
 		require.NoError(t, err)
 		defer c.Close()
 
-		require.NoError(t, c.ReplayWALs(context.Background()))
-
 		db, err = c.DB(context.Background(), "test")
 		require.NoError(t, err)
 		table, err = db.Table("test", config)
@@ -1258,13 +1255,10 @@ func Test_DB_TableNotExist(t *testing.T) {
 }
 
 func TestReplayBackwardsCompatibility(t *testing.T) {
-	ctx := context.Background()
 	const storagePath = "testdata/oldwal"
 	c, err := New(WithWAL(), WithStoragePath(storagePath))
 	require.NoError(t, err)
 	defer c.Close()
-
-	require.NoError(t, c.ReplayWALs(ctx))
 }
 
 func Test_DB_TableWrite_ArrowRecord(t *testing.T) {
@@ -1431,7 +1425,6 @@ func Test_DB_ReadOnlyQuery(t *testing.T) {
 	)
 	require.NoError(t, err)
 	defer c.Close()
-	require.NoError(t, c.ReplayWALs(context.Background()))
 
 	// Query with an aggregat query
 	pool := memory.NewGoAllocator()
@@ -1528,34 +1521,17 @@ func TestDBRecover(t *testing.T) {
 	// already been persisted.
 	require.NoError(t, os.Remove(filepath.Join(db.snapshotsDir(), snapshotFileName(8))))
 
-	numDistinctTimestamps := func(t *testing.T, db *DB) int {
-		t.Helper()
-		engine := query.NewEngine(memory.DefaultAllocator, db.TableProvider())
-		nrows := 0
-		require.NoError(t, engine.ScanTable(dbAndTableName).
-			Distinct(logicalplan.Col("timestamp")).
-			Execute(
-				ctx,
-				func(_ context.Context, r arrow.Record) error {
-					nrows += int(r.NumRows())
-					return nil
-				}))
-		return nrows
-	}
-
-	// Before we check there is no data when recovering our expected state,
-	// manually load the previous snapshot (without WAL) to verify that
-	// there is indeed some data.
-	verifyDB, err := c.DB(ctx, "verifyDB")
-	require.NoError(t, err)
-	snapshotTx, err := verifyDB.loadLatestSnapshotFromDir(ctx, db.snapshotsDir())
-	require.NoError(t, err)
-	require.Equal(t, uint64(5), snapshotTx)
-	require.Equal(t, numInserts, numDistinctTimestamps(t, verifyDB))
-
-	// Recover the full expected state.
-	require.NoError(t, db.recover(ctx))
+	engine := query.NewEngine(memory.DefaultAllocator, db.TableProvider())
+	nrows := 0
+	require.NoError(t, engine.ScanTable(dbAndTableName).
+		Distinct(logicalplan.Col("timestamp")).
+		Execute(
+			ctx,
+			func(_ context.Context, r arrow.Record) error {
+				nrows += int(r.NumRows())
+				return nil
+			}))
 	// No more timestamps if querying in-memory only, since the data has
 	// been rotated.
-	require.Equal(t, 0, numDistinctTimestamps(t, db))
+	require.Equal(t, 0, nrows)
 }
diff --git a/snapshot_test.go b/snapshot_test.go
index 973fd69..6ea16be 100644
--- a/snapshot_test.go
+++ b/snapshot_test.go
@@ -356,8 +356,6 @@ func TestSnapshotWithWAL(t *testing.T) {
 	// insert.
 	require.NoError(t, verifyDB.wal.Truncate(snapshotTx+1))
 
-	require.NoError(t, verifyC.ReplayWALs(ctx))
-
 	engine := query.NewEngine(memory.DefaultAllocator, verifyDB.TableProvider())
 	require.NoError(
 		t,
diff --git a/wal/wal_test.go b/wal/wal_test.go
index dbb03f4..583dc64 100644
--- a/wal/wal_test.go
+++ b/wal/wal_test.go
@@ -20,6 +20,7 @@ func TestWAL(t *testing.T) {
 		dir,
 	)
 	require.NoError(t, err)
+	w.RunAsync()
 
 	require.NoError(t, w.Log(1, &walpb.Record{
 		Entry: &walpb.Entry{
@@ -40,6 +41,7 @@ func TestWAL(t *testing.T) {
 		dir,
 	)
 	require.NoError(t, err)
+	w.RunAsync()
 
 	err = w.Replay(0, func(tx uint64, r *walpb.Record) error {
 		require.Equal(t, uint64(1), tx)
@@ -70,12 +72,8 @@ func TestWAL(t *testing.T) {
 		dir,
 	)
 	require.NoError(t, err)
+	w.RunAsync()
 	defer w.Close()
-
-	err = w.Replay(0, func(tx uint64, r *walpb.Record) error {
-		return nil
-	})
-	require.NoError(t, err)
 }
 
 func TestCorruptWAL(t *testing.T) {
@@ -87,6 +85,7 @@ func TestCorruptWAL(t *testing.T) {
 		path,
 	)
 	require.NoError(t, err)
+	w.RunAsync()
 
 	require.NoError(t, w.Log(1, &walpb.Record{
 		Entry: &walpb.Entry{
@@ -110,6 +109,7 @@ func TestCorruptWAL(t *testing.T) {
 		path,
 	)
 	require.NoError(t, err)
+	w.RunAsync()
 	defer w.Close()
 
 	lastIdx, err := w.LastIndex()
@@ -130,6 +130,7 @@ func TestUnexpectedTxn(t *testing.T) {
 			walDir,
 		)
 		require.NoError(t, err)
+		w.RunAsync()
 		defer w.Close()
 
 		emptyRecord := &walpb.Record{}
@@ -148,6 +149,7 @@ func TestUnexpectedTxn(t *testing.T) {
 		walDir,
 	)
 	require.NoError(t, err)
+	w.RunAsync()
 	defer w.Close()
 	lastIndex, err := w.LastIndex()
 	require.NoError(t, err)