Repository: polarsignals/frostdb (go)

Issue:
wal/snapshot: reclaim space more agressively
Currently, truncations happen in two cases:
1) At the minimum transaction persisted across all tables when a table block is persisted. As covered elsewhere, this could result in never truncating the WAL because small tables are never persisted, which means that the minimum transaction persisted across all tables is 0.
2) At startup when successfully loading a snapshot at txn k, the WAL can be truncated at index k as well as all snapshots < k.

In addition to truncating at startup, we could truncate the WAL as soon as we've taken a successful snapshot on block persistence and remove snapshots at previous txns, since those should never be used. To be cautious, I would start with re-opening the snapshot file after closing it just to double check the checksum since I'm wary of storage guarantees. I would maybe add a temporary metric to track if snapshots we expect to be durable are not to double check our assumptions, and remove this code after some time of running in prod.


Discussion:

Patch:
diff --git a/db.go b/db.go
--- a/db.go
+++ b/db.go
@@ -56,6 +56,11 @@ type ColumnStore struct {
 
 	sources []DataSource
 	sinks   []DataSink
+
+	// testingOptions are options only used for testing purposes.
+	testingOptions struct {
+		disableReclaimDiskSpaceOnSnapshot bool
+	}
 }
 
 type metrics struct {
@@ -859,6 +864,25 @@ 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 {
+	if db.columnStore.testingOptions.disableReclaimDiskSpaceOnSnapshot {
+		return nil
+	}
+	validSnapshotTxn, err := db.getLatestValidSnapshotTxn(ctx)
+	if err != nil {
+		return err
+	}
+	if validSnapshotTxn == 0 {
+		return nil
+	}
+	if err := db.truncateSnapshotsLessThanTX(ctx, validSnapshotTxn); err != nil {
+		return err
+	}
+	return db.wal.Truncate(validSnapshotTxn)
+}
+
 func (db *DB) getMinTXPersisted() uint64 {
 	db.mtx.RLock()
 	defer db.mtx.RUnlock()
diff --git a/snapshot.go b/snapshot.go
--- a/snapshot.go
+++ b/snapshot.go
@@ -178,30 +178,14 @@ func (db *DB) loadLatestSnapshot(ctx context.Context) (uint64, error) {
 }
 
 func (db *DB) loadLatestSnapshotFromDir(ctx context.Context, dir string) (uint64, error) {
-	files, err := os.ReadDir(dir)
-	if err != nil {
-		return 0, err
-	}
-
-	if len(files) == 0 {
-		return 0, nil
-	}
-
-	var lastErr error
-	// Iterate in reverse order so that the newest snapshot is first.
-	for i := len(files) - 1; i >= 0; i-- {
-		entry := files[i]
-		name := entry.Name()
-		if entry.IsDir() || len(name) < 20 {
-			continue
-		}
-		parsedTx, err := getTxFromSnapshotFileName(name)
-		if err != nil {
-			continue
-		}
-
+	var (
+		lastErr   error
+		loadedTxn uint64
+	)
+	// No error should be returned from snapshotsDo.
+	_ = db.snapshotsDo(ctx, dir, func(parsedTx uint64, entry os.DirEntry) (bool, error) {
 		if err := func() error {
-			f, err := os.Open(filepath.Join(dir, name))
+			f, err := os.Open(filepath.Join(dir, entry.Name()))
 			if err != nil {
 				return err
 			}
@@ -217,17 +201,23 @@ func (db *DB) loadLatestSnapshotFromDir(ctx context.Context, dir string) (uint64
 			db.highWatermark.Store(parsedTx)
 			return nil
 		}(); err != nil {
-			err = fmt.Errorf("unable to read snapshot file %s: %w", name, err)
+			err = fmt.Errorf("unable to read snapshot file %s: %w", entry.Name(), err)
 			level.Debug(db.logger).Log(
 				"msg", "error reading snapshot",
 				"error", err,
 			)
 			lastErr = err
-			continue
+			return true, nil
 		}
 		// Success.
-		return parsedTx, nil
+		loadedTxn = parsedTx
+		return false, nil
+	})
+	if loadedTxn != 0 {
+		// Successfully loaded a snapshot.
+		return loadedTxn, nil
 	}
+
 	errString := "no valid snapshots found"
 	if lastErr != nil {
 		return 0, fmt.Errorf("%s: lastErr: %w", errString, lastErr)
@@ -235,6 +225,41 @@ func (db *DB) loadLatestSnapshotFromDir(ctx context.Context, dir string) (uint64
 	return 0, fmt.Errorf("%s", errString)
 }
 
+func (db *DB) getLatestValidSnapshotTxn(ctx context.Context) (uint64, error) {
+	dir := db.snapshotsDir()
+	latestValidTxn := uint64(0)
+	// No error should be returned from snapshotsDo.
+	_ = db.snapshotsDo(ctx, dir, func(parsedTx uint64, entry os.DirEntry) (bool, error) {
+		if err := func() error {
+			f, err := os.Open(filepath.Join(dir, entry.Name()))
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			info, err := entry.Info()
+			if err != nil {
+				return err
+			}
+			// readFooter validates the checksum.
+			if _, err := readFooter(f, info.Size()); err != nil {
+				return err
+			}
+			return nil
+		}(); err != nil {
+			level.Debug(db.logger).Log(
+				"msg", "error reading snapshot",
+				"error", err,
+			)
+			// Continue to the next snapshot.
+			return true, nil
+		}
+		// Valid snapshot found.
+		latestValidTxn = parsedTx
+		return false, nil
+	})
+	return latestValidTxn, nil
+}
+
 type offsetWriter struct {
 	w               io.Writer
 	runningChecksum hash.Hash32
@@ -556,11 +581,29 @@ func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) error
 // than the given tx.
 func (db *DB) truncateSnapshotsLessThanTX(ctx context.Context, tx uint64) error {
 	dir := db.snapshotsDir()
+	return db.snapshotsDo(ctx, dir, func(fileTx uint64, entry os.DirEntry) (bool, error) {
+		if fileTx >= tx {
+			// Continue.
+			return true, nil
+		}
+		if err := os.Remove(filepath.Join(dir, entry.Name())); err != nil {
+			return false, err
+		}
+		return true, nil
+	})
+}
+
+// snapshotsDo executes the given callback with the filename of each snapshot
+// in dir in reverse lexicographical order (most recent snapshot first). If
+// false or an error is returned by the callback, the iteration is aborted and
+// the error returned.
+func (db *DB) snapshotsDo(ctx context.Context, dir string, callback func(tx uint64, entry os.DirEntry) (bool, error)) error {
 	files, err := os.ReadDir(dir)
 	if err != nil {
 		return err
 	}
-	for _, entry := range files {
+	for i := len(files) - 1; i >= 0; i-- {
+		entry := files[i]
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
@@ -572,12 +615,10 @@ func (db *DB) truncateSnapshotsLessThanTX(ctx context.Context, tx uint64) error
 		if err != nil {
 			continue
 		}
-		if parsedTx >= tx {
-			// All interesting files have been processed.
-			return nil
-		}
-		if err := os.Remove(filepath.Join(dir, name)); err != nil {
+		if ok, err := callback(parsedTx, entry); err != nil {
 			return err
+		} else if !ok {
+			return nil
 		}
 	}
 	return nil
diff --git a/table.go b/table.go
--- a/table.go
+++ b/table.go
@@ -473,12 +473,20 @@ func (t *Table) writeBlock(block *TableBlock, snapshotDB bool) {
 			// snapshot is not. Note that block.Persist does the same.
 			// TODO(asubiotto): Eventually we should register a cancel function
 			// that is called with a grace period on db.Close.
-			if err := t.db.snapshotAtTX(context.Background(), tx); err != nil {
+			ctx := context.Background()
+			if err := t.db.snapshotAtTX(ctx, tx); err != nil {
 				level.Error(t.logger).Log(
 					"msg", "failed to write snapshot on block rotation",
 					"err", err,
 				)
 			}
+			if err := t.db.reclaimDiskSpace(ctx); err != nil {
+				level.Error(t.logger).Log(
+					"msg", "failed to reclaim disk space after snapshot on block rotation",
+					"err", err,
+				)
+				return
+			}
 		}()
 	}
 }
@@ -759,6 +767,13 @@ func (t *Table) appender(ctx context.Context) (*TableBlock, func(), error) {
 					"last_snapshot_size", humanize.IBytes(uint64(block.lastSnapshotSize.Load())),
 				)
 				block.lastSnapshotSize.Store(blockSize)
+				if err := t.db.reclaimDiskSpace(ctx); err != nil {
+					level.Error(t.logger).Log(
+						"msg", "failed to reclaim disk space after snapshot",
+						"err", err,
+					)
+					return
+				}
 			})
 		}
 		if blockSize < t.db.columnStore.activeMemorySize {


Test patch:
diff --git a/db_test.go b/db_test.go
--- a/db_test.go
+++ b/db_test.go
@@ -1442,6 +1442,10 @@ func TestDBRecover(t *testing.T) {
 				WithStoragePath(dir),
 				WithWAL(),
 				WithSnapshotTriggerSize(1),
+				// Disable reclaiming disk space on snapshot (i.e. deleting
+				// old snapshots and WAL). This allows us to modify on-disk
+				// state for some tests.
+				WithTestingNoDiskSpaceReclaimOnSnapshot(),
 			},
 				options...,
 			)...,
diff --git a/snapshot_test.go b/snapshot_test.go
--- a/snapshot_test.go
+++ b/snapshot_test.go
@@ -35,6 +35,11 @@ func TestSnapshot(t *testing.T) {
 		db, err := c.DB(ctx, "test")
 		require.NoError(t, err)
 
+		// Complete a txn so that the snapshot is created at txn 1, snapshots at
+		// txn 0 are considered empty so ignored.
+		_, _, commit := db.begin()
+		commit()
+
 		require.NoError(t, db.snapshotAtTX(ctx, db.highWatermark.Load()))
 
 		txBefore := db.highWatermark.Load()
diff --git a/utils_test.go b/utils_test.go
new file mode 100644
--- /dev/null
+++ b/utils_test.go
@@ -0,0 +1,8 @@
+package frostdb
+
+func WithTestingNoDiskSpaceReclaimOnSnapshot() Option {
+	return func(s *ColumnStore) error {
+		s.testingOptions.disableReclaimDiskSpaceOnSnapshot = true
+		return nil
+	}
+}