Repository: polarsignals/frostdb
Pull request #569 (2023-10-23 15:32:39)

db: perform a snapshot if WAL contains persisted entries


Patch:
diff --git a/db.go b/db.go
index 453a37d..f43f103 100644
--- a/db.go
+++ b/db.go
@@ -681,6 +681,11 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 		return err
 	}
 
+	// performSnapshot is set to true if a snapshot should be performed after
+	// replay. This is set in cases where there could be "dead bytes" in the
+	// WAL (i.e. entries that occupy space on disk but are useless).
+	performSnapshot := false
+
 	// Writes are performed concurrently to speed up replay.
 	var writeWg errgroup.Group
 	writeWg.SetLimit(runtime.GOMAXPROCS(0))
@@ -755,6 +760,11 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 
 			// If we get to this point it means a block was finished but did
 			// not get persisted.
+			level.Info(db.logger).Log(
+				"msg", "writing unfinished block in recovery",
+				"table", tableName,
+				"tx", tx,
+			)
 			table.pendingBlocks[table.active] = struct{}{}
 			go table.writeBlock(table.active, db.columnStore.manualBlockRotation, false)
 
@@ -824,6 +834,10 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 			})
 
 		case *walpb.Entry_TableBlockPersisted_:
+			// If a block was persisted but the entry still exists in the WAL,
+			// a snapshot was not performed after persisting the block. Perform
+			// one now to clean up the WAL.
+			performSnapshot = true
 			return nil
 		case *walpb.Entry_Snapshot_:
 			return nil
@@ -843,7 +857,28 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 	if lastTx.TxnID > resetTxn.TxnID {
 		resetTxn = lastTx
 	}
+
+	db.mtx.Lock()
+	for _, table := range db.tables {
+		block := table.ActiveBlock()
+		block.uncompressedInsertsSize.Store(block.Index().LevelSize(index.L0))
+	}
+	db.mtx.Unlock()
+
 	db.resetToTxn(resetTxn, nil)
+	if performSnapshot && db.columnStore.snapshotTriggerSize != 0 {
+		level.Info(db.logger).Log(
+			"msg", "performing snapshot after recovery",
+		)
+		db.snapshot(ctx, false, func() {
+			if err := db.reclaimDiskSpace(ctx, wal); err != nil {
+				level.Error(db.logger).Log(
+					"msg", "failed to reclaim disk space after snapshot during recovery",
+					"err", err,
+				)
+			}
+		})
+	}
 	level.Info(db.logger).Log(
 		append(
 			[]any{
@@ -894,7 +929,7 @@ func (db *DB) Close(options ...CloseOption) error {
 		start := time.Now()
 		db.snapshot(context.Background(), false, func() {
 			level.Info(db.logger).Log("msg", "snapshot on close completed", "duration", time.Since(start))
-			if err := db.reclaimDiskSpace(context.Background()); err != nil {
+			if err := db.reclaimDiskSpace(context.Background(), nil); err != nil {
 				level.Error(db.logger).Log(
 					"msg", "failed to reclaim disk space after snapshot",
 					"err", err,
@@ -938,8 +973,10 @@ func (db *DB) maintainWAL() {
 }
 
 // reclaimDiskSpace attempts to read the latest valid snapshot txn and removes
-// any snapshots/wal entries that are older than the snapshot tx.
-func (db *DB) reclaimDiskSpace(ctx context.Context) error {
+// any snapshots/wal entries that are older than the snapshot tx. Since this can
+// be called before db.wal is set, the caller may optionally pass in a WAL to
+// truncate.
+func (db *DB) reclaimDiskSpace(ctx context.Context, wal WAL) error {
 	if db.columnStore.testingOptions.disableReclaimDiskSpaceOnSnapshot {
 		return nil
 	}
@@ -953,7 +990,10 @@ func (db *DB) reclaimDiskSpace(ctx context.Context) error {
 	if err := db.cleanupSnapshotDir(ctx, validSnapshotTxn); err != nil {
 		return err
 	}
-	return db.wal.Truncate(validSnapshotTxn)
+	if wal == nil {
+		wal = db.wal
+	}
+	return wal.Truncate(validSnapshotTxn)
 }
 
 func (db *DB) getMinTXPersisted() uint64 {
diff --git a/index/lsm.go b/index/lsm.go
index 161e6d3..926bdde 100644
--- a/index/lsm.go
+++ b/index/lsm.go
@@ -134,6 +134,11 @@ func (l *LSM) Size() int64 {
 	return size
 }
 
+// LevelSize returns the size of a specific level in bytes.
+func (l *LSM) LevelSize(t SentinelType) int64 {
+	return l.sizes[t].Load()
+}
+
 func validateLevels(levels []*LevelConfig) error {
 	for i, l := range levels {
 		if int(l.Level) != i {
diff --git a/snapshot.go b/snapshot.go
index 853c961..da92111 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -220,14 +220,18 @@ func (db *DB) snapshotAtTX(ctx context.Context, tx uint64, writeSnapshot func(co
 		}
 		return nil
 	}(); err != nil {
-		db.metrics.snapshotMetrics.snapshotsTotal.WithLabelValues("false").Inc()
+		if db.metrics != nil {
+			db.metrics.snapshotMetrics.snapshotsTotal.WithLabelValues("false").Inc()
+		}
 		return err
 	}
-	db.metrics.snapshotMetrics.snapshotsTotal.WithLabelValues("true").Inc()
-	if fileSize > 0 {
-		db.metrics.snapshotMetrics.snapshotFileSizeBytes.Set(float64(fileSize))
+	if db.metrics != nil {
+		db.metrics.snapshotMetrics.snapshotsTotal.WithLabelValues("true").Inc()
+		if fileSize > 0 {
+			db.metrics.snapshotMetrics.snapshotFileSizeBytes.Set(float64(fileSize))
+		}
+		db.metrics.snapshotMetrics.snapshotDurationHistogram.Observe(time.Since(start).Seconds())
 	}
-	db.metrics.snapshotMetrics.snapshotDurationHistogram.Observe(time.Since(start).Seconds())
 	// TODO(asubiotto): If snapshot file sizes become too large, investigate
 	// adding compression.
 	return nil
diff --git a/table.go b/table.go
index 50056b8..96fbcdb 100644
--- a/table.go
+++ b/table.go
@@ -466,7 +466,7 @@ func (t *Table) writeBlock(block *TableBlock, skipPersist, snapshotDB bool) {
 					"err", err,
 				)
 			}
-			if err := t.db.reclaimDiskSpace(ctx); err != nil {
+			if err := t.db.reclaimDiskSpace(ctx, nil); err != nil {
 				level.Error(t.logger).Log(
 					"msg", "failed to reclaim disk space after snapshot on block rotation",
 					"err", err,
@@ -608,7 +608,7 @@ func (t *Table) appender(ctx context.Context) (*TableBlock, func(), error) {
 					"last_snapshot_size", humanize.IBytes(uint64(block.lastSnapshotSize.Load())),
 				)
 				block.lastSnapshotSize.Store(uncompressedInsertsSize)
-				if err := t.db.reclaimDiskSpace(context.Background()); err != nil {
+				if err := t.db.reclaimDiskSpace(context.Background(), nil); err != nil {
 					level.Error(t.logger).Log(
 						"msg", "failed to reclaim disk space after snapshot",
 						"err", err,


Test patch:
diff --git a/db_test.go b/db_test.go
index 918580e..f748d1f 100644
--- a/db_test.go
+++ b/db_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/polarsignals/frostdb/dynparquet"
 	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
+	walpb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/wal/v1alpha1"
 	"github.com/polarsignals/frostdb/query"
 	"github.com/polarsignals/frostdb/query/logicalplan"
 	"github.com/polarsignals/frostdb/query/physicalplan"
@@ -1630,6 +1631,81 @@ func TestDBRecover(t *testing.T) {
 		require.NoError(t, err)
 		newWriteAndExpectWALRecord(t, db, table)
 	})
+
+	// SnapshotOnRecovery verifies that a snapshot is taken on recovery if the
+	// WAL indicates that a block was rotated but no snapshot was taken.
+	t.Run("SnapshotOnRecovery", func(t *testing.T) {
+		dir := setup(t, false)
+		c, err := New(
+			WithLogger(newTestLogger(t)),
+			WithStoragePath(dir),
+			WithWAL(),
+			// This option will disable snapshots on block rotation.
+			WithSnapshotTriggerSize(0),
+		)
+		require.NoError(t, err)
+
+		snapshotsPath := filepath.Join(dir, "databases", dbAndTableName, "snapshots")
+		snapshots, err := os.ReadDir(snapshotsPath)
+		require.NoError(t, err)
+
+		seenSnapshots := make(map[string]struct{})
+		for _, s := range snapshots {
+			seenSnapshots[s.Name()] = struct{}{}
+		}
+
+		db, err := c.DB(ctx, dbAndTableName)
+		require.NoError(t, err)
+		table, err := db.Table(dbAndTableName, nil)
+		require.NoError(t, err)
+
+		require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+
+		rec, err := dynparquet.NewTestSamples().ToRecord()
+		require.NoError(t, err)
+
+		insertTx, err := table.InsertRecord(ctx, rec)
+		require.NoError(t, err)
+
+		// RotateBlock again, this should log a couple of persisted block WAL
+		// entries.
+		require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
+		require.NoError(t, c.Close())
+
+		c, err = New(
+			WithLogger(newTestLogger(t)),
+			WithStoragePath(dir),
+			WithWAL(),
+			// Enable snapshots.
+			WithSnapshotTriggerSize(1),
+		)
+		require.NoError(t, err)
+		defer c.Close()
+
+		snapshots, err = os.ReadDir(snapshotsPath)
+		require.NoError(t, err)
+
+		for _, s := range snapshots {
+			tx, err := getTxFromSnapshotFileName(s.Name())
+			require.NoError(t, err)
+			require.GreaterOrEqual(t, tx, insertTx, "expected only snapshots after insert txn")
+		}
+		db, err = c.DB(ctx, dbAndTableName)
+		require.NoError(t, err)
+
+		require.Eventually(t, func() bool {
+			numBlockPersists := 0
+			require.NoError(t, db.wal.Replay(0, func(tx uint64, entry *walpb.Record) error {
+				if _, ok := entry.Entry.EntryType.(*walpb.Entry_TableBlockPersisted_); ok {
+					numBlockPersists++
+				}
+				return nil
+			}))
+			return numBlockPersists <= 1
+		}, 1*time.Second, 10*time.Millisecond,
+			"expected at most one block persist entry; the others should have been snapshot and truncated",
+		)
+	})
 }
 
 func Test_DB_WalReplayTableConfig(t *testing.T) {