Repository: polarsignals/frostdb
Pull request #745 (2024-02-15 22:32:50)

Snapshot: Always acquire an ActiveWriteBlock
When we were closing the database we may take a snapshot on close. If we do, that snapshot was taking on "offline" snapshot which doesn't acquire the active writeblock.

Because it doesn't acquire an active write block it menat it could race with an active block rotation, which could close the index before the snapshot occurs resulting in a bad snapshot.

This changes the snapshot behavior to always be "online" and acquire the active write block.

asubiotto: Actually, isn't it enough to just move snapshotting before closing the tables?
thorfour: No, because moving it to before would still cause it to race with the block rotation 

Patch:
diff --git a/db.go b/db.go
index 4bf531c..b2c63f1 100644
--- a/db.go
+++ b/db.go
@@ -984,8 +984,22 @@ func (db *DB) Close(options ...CloseOption) error {
 	for _, opt := range options {
 		opt(opts)
 	}
-	level.Info(db.logger).Log("msg", "closing DB")
+
 	shouldPersist := len(db.sinks) > 0 && !db.columnStore.manualBlockRotation
+	if !shouldPersist && db.columnStore.snapshotTriggerSize != 0 && !opts.clearStorage {
+		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(), nil); err != nil {
+				level.Error(db.logger).Log(
+					"msg", "failed to reclaim disk space after snapshot",
+					"err", err,
+				)
+			}
+		})
+	}
+
+	level.Info(db.logger).Log("msg", "closing DB")
 	for _, table := range db.tables {
 		table.close()
 		if shouldPersist {
@@ -1000,19 +1014,6 @@ func (db *DB) Close(options ...CloseOption) error {
 	}
 	level.Info(db.logger).Log("msg", "closed all tables")
 
-	if !shouldPersist && db.columnStore.snapshotTriggerSize != 0 && !opts.clearStorage {
-		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(), nil); err != nil {
-				level.Error(db.logger).Log(
-					"msg", "failed to reclaim disk space after snapshot",
-					"err", err,
-				)
-			}
-		})
-	}
-
 	if err := db.closeInternal(); err != nil {
 		return err
 	}
diff --git a/snapshot.go b/snapshot.go
index c4a2d3d..20374a3 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -346,18 +346,18 @@ func (w *offsetWriter) checksum() uint32 {
 
 func (db *DB) snapshotWriter(tx uint64) func(context.Context, io.Writer) error {
 	return func(ctx context.Context, w io.Writer) error {
-		return WriteSnapshot(ctx, tx, db, w, false)
+		return WriteSnapshot(ctx, tx, db, w)
 	}
 }
 
 // offlineSnapshotWriter is used when a database is closing after all the tables have closed.
 func (db *DB) offlineSnapshotWriter(tx uint64) func(context.Context, io.Writer) error {
 	return func(ctx context.Context, w io.Writer) error {
-		return WriteSnapshot(ctx, tx, db, w, true)
+		return WriteSnapshot(ctx, tx, db, w)
 	}
 }
 
-func WriteSnapshot(ctx context.Context, tx uint64, db *DB, w io.Writer, offline bool) error {
+func WriteSnapshot(ctx context.Context, tx uint64, db *DB, w io.Writer) error {
 	offW := newOffsetWriter(w)
 	w = offW
 	var tables []*Table
@@ -376,18 +376,11 @@ func WriteSnapshot(ctx context.Context, tx uint64, db *DB, w io.Writer, offline
 		if err := func() error {
 			// Obtain a write block to prevent racing with
 			// compaction/persistence.
-			var block *TableBlock
-			if offline {
-				block = t.ActiveBlock()
-			} else {
-				var done func()
-				var err error
-				block, done, err = t.ActiveWriteBlock()
-				if err != nil {
-					return err
-				}
-				defer done()
+			block, done, err := t.ActiveWriteBlock()
+			if err != nil {
+				return err
 			}
+			defer done()
 			blockUlid, err := block.ulid.MarshalBinary()
 			if err != nil {
 				return err


Test patch:
diff --git a/bench_test.go b/bench_test.go
index 81b3934..98c8fea 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -409,7 +409,7 @@ func BenchmarkSnapshot(b *testing.B) {
 	bytesWritten := 0
 	for i := 0; i < b.N; i++ {
 		wc := &writeCounter{Writer: io.Discard}
-		require.NoError(b, WriteSnapshot(ctx, db.HighWatermark(), db, wc, false))
+		require.NoError(b, WriteSnapshot(ctx, db.HighWatermark(), db, wc))
 		bytesWritten += wc.count
 	}
 	b.ReportMetric(float64(bytesWritten)/float64(b.N), "size/op")