Repository: polarsignals/frostdb (go)

Issue:
snapshot: two snapshots at the same txn causes data loss
If a DB receives no writes between two snapshots, the second snapshot will error out, causing the first snapshot to be removed from disk and resulting in data loss:
```
level=error db=<...> name=storage ts=2024-02-22T08:24:04.677341856Z caller=snapshot.go:140 msg="failed to snapshot database" err="failed to write snapshot for tx 47318: failed to snapshot table <...> index: link /data/database/<...>/index/table/01HQ64SWCRHN5AR6RS01SXYN16/L1/00000000000000000000.idx /data/databases/<...>/snapshots/00000000000000047318/index/table/01HQ64SWCRHN5AR6RS01SXYN16/L1/00000000000000000000.idx: file exists"
```

The above error causes the following log line to be printed on recovery:
```
level=info db=<...> name=storage ts=2024-02-22T08:24:30.532034244Z caller=db.go:700 msg="failed to load latest snapshot" db=<...> err="no valid snapshots found"
```
The culprit seems to be this line:
https://github.com/polarsignals/frostdb/blob/d9f86325ab691f52e33d19b86ff9d78b6ea313c3/snapshot.go#L192

Couple of ways to fix this. Maybe the easiest is to us `reclaimDiskSpace` instead of `RemoveAll` which takes into account snapshot validity on disk. Another is to remove a snapshot with the same txn on disk before performing the new snapshot, or checking whether a valid snapshot already exists on disk (since if it is invalid I think we do want to perform it), but this might be more expensive.


Discussion:

Patch:
diff --git a/snapshot.go b/snapshot.go
--- a/snapshot.go
+++ b/snapshot.go
@@ -163,11 +163,22 @@ func (db *DB) snapshotAtTX(ctx context.Context, tx uint64, writeSnapshot func(co
 	start := time.Now()
 	if err := func() error {
 		snapshotsDir := SnapshotDir(db, tx)
+		fileName := filepath.Join(snapshotsDir, snapshotFileName(tx))
+		_, err := os.Stat(fileName)
+		if err == nil { // Snapshot file already exists
+			if db.validateSnapshotTxn(ctx, tx) == nil {
+				return nil // valid snapshot already exists at tx no need to re-snapshot
+			}
+
+			// Snapshot exists but is invalid. Remove it.
+			if err := os.RemoveAll(SnapshotDir(db, tx)); err != nil {
+				return fmt.Errorf("failed to remove invalid snapshot %v: %w", tx, err)
+			}
+		}
 		if err := os.MkdirAll(snapshotsDir, dirPerms); err != nil {
 			return err
 		}
 
-		fileName := filepath.Join(snapshotsDir, snapshotFileName(tx))
 		f, err := os.OpenFile(fileName, os.O_CREATE|os.O_RDWR|os.O_TRUNC, filePerms)
 		if err != nil {
 			return err
@@ -279,6 +290,33 @@ func LoadSnapshot(ctx context.Context, db *DB, tx uint64, r io.ReaderAt, size in
 	return watermark, nil
 }
 
+func (db *DB) validateSnapshotTxn(ctx context.Context, tx uint64) error {
+	dir := db.snapshotsDir()
+
+	return db.snapshotsDo(ctx, dir, func(parsedTx uint64, entry os.DirEntry) (bool, error) {
+		if parsedTx != tx { // We're only trying to validate a single tx
+			return true, nil
+		}
+
+		return false, func() error {
+			f, err := os.Open(filepath.Join(dir, entry.Name(), snapshotFileName(parsedTx)))
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			info, err := f.Stat()
+			if err != nil {
+				return err
+			}
+			// readFooter validates the checksum.
+			if _, err := readFooter(f, info.Size()); err != nil {
+				return err
+			}
+			return nil
+		}()
+	})
+}
+
 func (db *DB) getLatestValidSnapshotTxn(ctx context.Context) (uint64, error) {
 	dir := db.snapshotsDir()
 	latestValidTxn := uint64(0)


Test patch:
diff --git a/db_test.go b/db_test.go
--- a/db_test.go
+++ b/db_test.go
@@ -2784,3 +2784,177 @@ func Test_DB_SnapshotNewerData(t *testing.T) {
 		})
 	}
 }
+
+// Test_DB_SnapshotDuplicate verifies that if we attempt to take a snapshot again at the same tx that it will abort if the current snapshot is valid.
+func Test_DB_SnapshotDuplicate(t *testing.T) {
+	t.Parallel()
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+	logger := newTestLogger(t)
+
+	cfg := []*index.LevelConfig{
+		{Level: index.L0, MaxSize: 336, Type: index.CompactionTypeParquetDisk},
+		{Level: index.L1, MaxSize: 1024 * 1024 * 128, Type: index.CompactionTypeParquetDisk},
+		{Level: index.L2, MaxSize: 1024 * 1024 * 512},
+	}
+	dir := t.TempDir()
+	c, err := New(
+		WithLogger(logger),
+		WithIndexConfig(cfg),
+		WithStoragePath(dir),
+		WithWAL(),
+		WithManualBlockRotation(),
+		WithSnapshotTriggerSize(1024*1024*1024), // we just need a non-zero value; but don't want the writes to unexpectedly trigger a snapshot
+	)
+	t.Cleanup(func() {
+		require.NoError(t, c.Close())
+	})
+	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)
+
+	ctx := context.Background()
+	samples := dynparquet.GenerateTestSamples(10)
+	r, err := samples.ToRecord()
+	require.NoError(t, err)
+	_, err = table.InsertRecord(ctx, r)
+	require.NoError(t, err)
+
+	require.NoError(t, table.EnsureCompaction())
+
+	// Snapshot the file
+	success := false
+	table.db.snapshot(context.Background(), false, func() {
+		success = true
+	})
+	require.True(t, success)
+
+	require.NoError(t, db.wal.Truncate(2))
+	time.Sleep(1 * time.Second) // wal flushes every 50ms
+
+	// Close the database to trigger a snapshot on close
+	require.NoError(t, c.Close())
+
+	// Reopen the database and expect the data to be there
+	c, err = New(
+		WithLogger(logger),
+		WithIndexConfig(cfg),
+		WithStoragePath(dir),
+		WithWAL(),
+		WithManualBlockRotation(),
+		WithSnapshotTriggerSize(1024*1024*1024), // we just need a non-zero value; but don't want the writes to unexpectedly trigger a snapshot
+	)
+	t.Cleanup(func() {
+		require.NoError(t, c.Close())
+	})
+	require.NoError(t, err)
+	db, err = c.DB(context.Background(), "test")
+	require.NoError(t, err)
+
+	validateRows := func(expected int64) {
+		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+		defer pool.AssertSize(t, 0)
+		rows := int64(0)
+		engine := query.NewEngine(pool, db.TableProvider())
+		err = engine.ScanTable("test").
+			Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+				rows += r.NumRows()
+				return nil
+			})
+		require.NoError(t, err)
+		require.Equal(t, expected, rows)
+	}
+
+	validateRows(10)
+}
+
+func Test_DB_SnapshotDuplicate_Corrupted(t *testing.T) {
+	t.Parallel()
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+	logger := newTestLogger(t)
+
+	cfg := []*index.LevelConfig{
+		{Level: index.L0, MaxSize: 336, Type: index.CompactionTypeParquetDisk},
+		{Level: index.L1, MaxSize: 1024 * 1024 * 128, Type: index.CompactionTypeParquetDisk},
+		{Level: index.L2, MaxSize: 1024 * 1024 * 512},
+	}
+	dir := t.TempDir()
+	c, err := New(
+		WithLogger(logger),
+		WithIndexConfig(cfg),
+		WithStoragePath(dir),
+		WithWAL(),
+		WithManualBlockRotation(),
+		WithSnapshotTriggerSize(1024*1024*1024), // we just need a non-zero value; but don't want the writes to unexpectedly trigger a snapshot
+	)
+	t.Cleanup(func() {
+		require.NoError(t, c.Close())
+	})
+	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)
+
+	ctx := context.Background()
+	samples := dynparquet.GenerateTestSamples(10)
+	r, err := samples.ToRecord()
+	require.NoError(t, err)
+	_, err = table.InsertRecord(ctx, r)
+	require.NoError(t, err)
+
+	require.NoError(t, table.EnsureCompaction())
+
+	// Snapshot the file
+	success := false
+	table.db.snapshot(context.Background(), false, func() {
+		success = true
+	})
+	require.True(t, success)
+
+	require.NoError(t, db.wal.Truncate(2))
+	time.Sleep(1 * time.Second) // wal flushes every 50ms
+
+	// Corrupt the snapshot file.
+	require.NoError(t, os.Truncate(filepath.Join(SnapshotDir(db, 2), snapshotFileName(2)), 100))
+
+	// Close the database to trigger a snapshot on close
+	require.NoError(t, c.Close())
+
+	// Reopen the database and expect the data to be there
+	c, err = New(
+		WithLogger(logger),
+		WithIndexConfig(cfg),
+		WithStoragePath(dir),
+		WithWAL(),
+		WithManualBlockRotation(),
+		WithSnapshotTriggerSize(1024*1024*1024), // we just need a non-zero value; but don't want the writes to unexpectedly trigger a snapshot
+	)
+	t.Cleanup(func() {
+		require.NoError(t, c.Close())
+	})
+	require.NoError(t, err)
+	db, err = c.DB(context.Background(), "test")
+	require.NoError(t, err)
+
+	validateRows := func(expected int64) {
+		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+		defer pool.AssertSize(t, 0)
+		rows := int64(0)
+		engine := query.NewEngine(pool, db.TableProvider())
+		err = engine.ScanTable("test").
+			Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+				rows += r.NumRows()
+				return nil
+			})
+		require.NoError(t, err)
+		require.Equal(t, expected, rows)
+	}
+
+	validateRows(10)
+}