Repository: polarsignals/frostdb
Pull request #491 (2023-08-14 19:54:09)

Store snapshot
Small refactor inject a writer dependency into `snapshotAtATX` which allows for a new function `StoreSnapshot` that can write an external snapshot into the snapshot directory. 


Patch:
diff --git a/db.go b/db.go
index e522f2a..eed4253 100644
--- a/db.go
+++ b/db.go
@@ -477,9 +477,7 @@ func (s *ColumnStore) DB(ctx context.Context, name string, opts ...DBOption) (*D
 			}, func() float64 {
 				return float64(db.highWatermark.Load().TxnID)
 			}),
-		}
-		if db.columnStore.snapshotTriggerSize != 0 {
-			db.metrics.snapshotMetrics = newSnapshotMetrics(reg)
+			snapshotMetrics: newSnapshotMetrics(reg),
 		}
 		return nil
 	}(); dbSetupErr != nil {
diff --git a/snapshot.go b/snapshot.go
index 0412c85..668e2fe 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -157,7 +157,7 @@ func (db *DB) asyncSnapshot(ctx context.Context, onSuccess func()) {
 			}
 		}
 
-		if err := db.snapshotAtTX(ctx, tx, metadata); err != nil {
+		if err := db.snapshotAtTX(ctx, tx, metadata, db.snapshotWriter(tx, metadata)); err != nil {
 			level.Error(db.logger).Log(
 				"msg", "failed to snapshot database", "err", err,
 			)
@@ -173,7 +173,7 @@ func (db *DB) asyncSnapshot(ctx context.Context, onSuccess func()) {
 }
 
 // snapshotAtTX takes a snapshot of the state of the database at transaction tx.
-func (db *DB) snapshotAtTX(ctx context.Context, tx uint64, txnMetadata []byte) error {
+func (db *DB) snapshotAtTX(ctx context.Context, tx uint64, txnMetadata []byte, writeSnapshot func(context.Context, io.Writer) error) error {
 	var fileSize int64
 	start := time.Now()
 	if err := func() error {
@@ -190,7 +190,7 @@ func (db *DB) snapshotAtTX(ctx context.Context, tx uint64, txnMetadata []byte) e
 		defer f.Close()
 
 		if err := func() error {
-			if err := WriteSnapshot(ctx, tx, txnMetadata, db, f); err != nil {
+			if err := writeSnapshot(ctx, f); err != nil {
 				return err
 			}
 			if err := f.Sync(); err != nil {
@@ -356,6 +356,12 @@ func (w *offsetWriter) checksum() uint32 {
 	return w.runningChecksum.Sum32()
 }
 
+func (db *DB) snapshotWriter(tx uint64, txnMetadata []byte) func(context.Context, io.Writer) error {
+	return func(ctx context.Context, w io.Writer) error {
+		return WriteSnapshot(ctx, tx, txnMetadata, db, w)
+	}
+}
+
 func WriteSnapshot(ctx context.Context, tx uint64, txnMetadata []byte, db *DB, w io.Writer) error {
 	offW := newOffsetWriter(w)
 	w = offW
@@ -700,3 +706,10 @@ func (db *DB) snapshotsDo(ctx context.Context, dir string, callback func(tx uint
 	}
 	return nil
 }
+
+func StoreSnapshot(ctx context.Context, tx uint64, txnMetadata []byte, db *DB, snapshot io.Reader) error {
+	return db.snapshotAtTX(ctx, tx, txnMetadata, func(ctx context.Context, w io.Writer) error {
+		_, err := io.Copy(w, snapshot)
+		return err
+	})
+}
diff --git a/table.go b/table.go
index 53d7247..7950bf4 100644
--- a/table.go
+++ b/table.go
@@ -481,7 +481,7 @@ func (t *Table) writeBlock(block *TableBlock, skipPersist, snapshotDB bool) {
 			// TODO(asubiotto): Eventually we should register a cancel function
 			// that is called with a grace period on db.Close.
 			ctx := context.Background()
-			if err := t.db.snapshotAtTX(ctx, tx, metadata); err != nil {
+			if err := t.db.snapshotAtTX(ctx, tx, metadata, t.db.snapshotWriter(tx, metadata)); err != nil {
 				level.Error(t.logger).Log(
 					"msg", "failed to write snapshot on block rotation",
 					"err", err,


Test patch:
diff --git a/snapshot_test.go b/snapshot_test.go
index 5c85878..5c2777d 100644
--- a/snapshot_test.go
+++ b/snapshot_test.go
@@ -41,10 +41,11 @@ func TestSnapshot(t *testing.T) {
 		_, _, metadata, commit := db.begin()
 		commit()
 
-		require.NoError(t, db.snapshotAtTX(ctx, db.highWatermark.Load().TxnID, metadata))
+		tx := db.highWatermark.Load()
+		require.NoError(t, db.snapshotAtTX(ctx, tx.TxnID, metadata, db.snapshotWriter(tx.TxnID, metadata)))
 
 		txBefore := db.highWatermark.Load().TxnID
-		tx, err := db.loadLatestSnapshot(ctx)
+		tx, err = db.loadLatestSnapshot(ctx)
 		require.NoError(t, err)
 		require.Equal(t, txBefore, tx.TxnID)
 	})
@@ -87,7 +88,7 @@ func TestSnapshot(t *testing.T) {
 
 		// Insert a sample that should not be snapshot.
 		insertSamples(ctx, t, table, 10)
-		require.NoError(t, db.snapshotAtTX(ctx, highWatermark, nil))
+		require.NoError(t, db.snapshotAtTX(ctx, highWatermark, nil, db.snapshotWriter(highWatermark, nil)))
 
 		// Create another db and verify.
 		snapshotDB, err := c.DB(ctx, "testsnapshot")
@@ -187,7 +188,8 @@ func TestSnapshot(t *testing.T) {
 		defer cancelWrites()
 		snapshotDB, err := c.DB(ctx, "testsnapshot")
 		require.NoError(t, err)
-		require.NoError(t, db.snapshotAtTX(ctx, db.highWatermark.Load().TxnID, nil))
+		tx := db.highWatermark.Load()
+		require.NoError(t, db.snapshotAtTX(ctx, tx.TxnID, nil, db.snapshotWriter(tx.TxnID, nil)))
 		snapshotTx, err := snapshotDB.loadLatestSnapshotFromDir(ctx, db.snapshotsDir())
 		require.NoError(t, err)
 		require.NoError(