Repository: polarsignals/frostdb
Pull request #525 (2023-09-07 20:35:31)

table/close: snapshot if not persisting


Patch:
diff --git a/db.go b/db.go
index 3ac45fa..c7b5a09 100644
--- a/db.go
+++ b/db.go
@@ -841,6 +841,19 @@ func (db *DB) Close(options ...CloseOption) error {
 	}
 	level.Info(db.logger).Log("msg", "closed all tables")
 
+	if !shouldPersist && db.columnStore.snapshotTriggerSize != 0 {
+		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 {
+				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 844dbdb..76ea172 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -115,6 +115,10 @@ func getTxFromSnapshotFileName(fileName string) (uint64, error) {
 // started (i.e. no other snapshot was in progress). When the snapshot
 // goroutine successfully completes a snapshot, onSuccess is called.
 func (db *DB) asyncSnapshot(ctx context.Context, onSuccess func()) {
+	db.snapshot(ctx, true, onSuccess)
+}
+
+func (db *DB) snapshot(ctx context.Context, async bool, onSuccess func()) {
 	if !db.columnStore.enableWAL {
 		return
 	}
@@ -131,7 +135,7 @@ func (db *DB) asyncSnapshot(ctx context.Context, onSuccess func()) {
 		"msg", "starting a new snapshot",
 		"tx", tx,
 	)
-	go func() {
+	doSnapshot := func(writeSnapshot func(context.Context, io.Writer) error) {
 		start := time.Now()
 		defer db.snapshotInProgress.Store(false)
 		defer commit()
@@ -157,7 +161,7 @@ func (db *DB) asyncSnapshot(ctx context.Context, onSuccess func()) {
 			}
 		}
 
-		if err := db.snapshotAtTX(ctx, tx, db.snapshotWriter(tx, metadata)); err != nil {
+		if err := db.snapshotAtTX(ctx, tx, writeSnapshot); err != nil {
 			level.Error(db.logger).Log(
 				"msg", "failed to snapshot database", "err", err,
 			)
@@ -169,7 +173,13 @@ func (db *DB) asyncSnapshot(ctx context.Context, onSuccess func()) {
 			"duration", time.Since(start),
 		)
 		onSuccess()
-	}()
+	}
+
+	if async {
+		go doSnapshot(db.snapshotWriter(tx, metadata))
+	} else {
+		doSnapshot(db.offlineSnapshotWriter(tx, metadata))
+	}
 }
 
 // snapshotAtTX takes a snapshot of the state of the database at transaction tx.
@@ -358,11 +368,18 @@ func (w *offsetWriter) checksum() uint32 {
 
 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)
+		return WriteSnapshot(ctx, tx, txnMetadata, db, w, false)
+	}
+}
+
+// offlineSnapshotWriter is used when a database is closing after all the tables have closed.
+func (db *DB) offlineSnapshotWriter(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, true)
 	}
 }
 
-func WriteSnapshot(ctx context.Context, tx uint64, txnMetadata []byte, db *DB, w io.Writer) error {
+func WriteSnapshot(ctx context.Context, tx uint64, txnMetadata []byte, db *DB, w io.Writer, offline bool) error {
 	offW := newOffsetWriter(w)
 	w = offW
 	var tables []*Table
@@ -381,11 +398,18 @@ func WriteSnapshot(ctx context.Context, tx uint64, txnMetadata []byte, db *DB, w
 		if err := func() error {
 			// Obtain a write block to prevent racing with
 			// compaction/persistence.
-			block, done, err := t.ActiveWriteBlock()
-			if err != nil {
-				return err
+			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()
 			}
-			defer done()
 			blockUlid, err := block.ulid.MarshalBinary()
 			if err != nil {
 				return err


Test patch:
diff --git a/db_test.go b/db_test.go
index 797d04e..893781e 100644
--- a/db_test.go
+++ b/db_test.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"strings"
@@ -1569,11 +1570,11 @@ func TestDBRecover(t *testing.T) {
 		snapshotsPath := filepath.Join(dir, "databases", dbAndTableName, "snapshots")
 		files, err := os.ReadDir(snapshotsPath)
 		require.NoError(t, err)
-		require.Equal(t, 2, len(files))
+		require.Equal(t, 3, len(files))
 		require.NoError(t, os.RemoveAll(filepath.Join(snapshotsPath, files[len(files)-1].Name())))
 		files, err = os.ReadDir(snapshotsPath)
 		require.NoError(t, err)
-		require.Equal(t, 1, len(files))
+		require.Equal(t, 2, len(files))
 
 		c, err := New(
 			WithLogger(newTestLogger(t)),
@@ -1915,3 +1916,47 @@ func Test_DB_EngineInMemory(t *testing.T) {
 		})
 	require.NoError(t, err)
 }
+
+func Test_DB_SnapshotOnClose(t *testing.T) {
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+
+	logger := newTestLogger(t)
+	dir := t.TempDir()
+
+	c, err := New(
+		WithLogger(logger),
+		WithStoragePath(dir),
+		WithWAL(),
+		WithActiveMemorySize(1024*1024*1024),
+		WithSnapshotTriggerSize(1024*1024*1024),
+		WithManualBlockRotation(),
+	)
+	require.NoError(t, err)
+	db, err := c.DB(context.Background(), "test")
+	require.NoError(t, err)
+	table, err := db.Table("test", config)
+	require.NoError(t, err)
+
+	samples := dynparquet.NewTestSamples()
+
+	ctx := context.Background()
+	for i := 0; i < 100; i++ {
+		buf, err := samples.ToBuffer(table.Schema())
+		require.NoError(t, err)
+		_, err = table.InsertBuffer(ctx, buf)
+		require.NoError(t, err)
+	}
+	require.NoError(t, c.Close())
+
+	// Check that we have a snapshot
+	found := false
+	require.NoError(t, filepath.WalkDir(dir, func(path string, info fs.DirEntry, err error) error {
+		if filepath.Ext(path) == ".fdbs" {
+			found = true
+		}
+		return nil
+	}))
+	require.True(t, found)
+}