Repository: polarsignals/frostdb
Pull request #711 (2024-01-26 20:30:52)

Fix index deadlocks
Snapshot's and Compaction could deadlock. This fixes that.

Compactions wait for the watermark to reach the newest transaction in the L0 level.
Snapshots were creating a write tx, but where then calling index.Snapshot which acquires the index.Compaction lock.

This could result in a deadlock if a write came in that was newer than the snapshot tx and the compaction was waiting for the watermark to get to that, but the snapshot was waiting to commit it's tx waiting for the compaction lock.

Instead of acquiring a write tx we need only to get a read tx so we don't block the watermark. We also do not need to log to the WAL for snapshots because we do nothing with that information on replay.


Patch:
diff --git a/index/lsm.go b/index/lsm.go
index 30af97f..71430c6 100644
--- a/index/lsm.go
+++ b/index/lsm.go
@@ -158,6 +158,10 @@ func NewLSM(dir string, schema *dynparquet.Schema, levels []*LevelConfig, wait f
 
 	if lsm.metrics == nil {
 		lsm.metrics = NewLSMMetrics(prometheus.NewRegistry())
+	} else {
+		for _, lvl := range levels {
+			lsm.metrics.LevelSize.WithLabelValues(lvl.Level.String()).Set(0)
+		}
 	}
 
 	// Replay the recovered parts
@@ -371,6 +375,7 @@ func (l *LSM) Scan(ctx context.Context, _ string, _ *dynparquet.Schema, filter l
 	if err != nil {
 		return fmt.Errorf("boolean expr: %w", err)
 	}
+
 	var iterError error
 	l.partList.Iterate(func(node *Node) bool {
 		if node.part == nil { // encountered a sentinel node; continue on
@@ -587,10 +592,10 @@ func (l *LSM) merge(level SentinelType) error {
 
 	// release the old parts
 	l.Lock()
-	defer l.Unlock()
 	for _, part := range mergeList {
 		part.Release()
 	}
+	l.Unlock()
 
 	// Reset the level that was just compacted
 	if level != L0 {
diff --git a/snapshot.go b/snapshot.go
index 6c57297..1bc1b03 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -25,7 +25,6 @@ import (
 	"github.com/polarsignals/frostdb/dynparquet"
 	snapshotpb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/snapshot/v1alpha1"
 	tablepb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/table/v1alpha1"
-	walpb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/wal/v1alpha1"
 	"github.com/polarsignals/frostdb/index"
 	"github.com/polarsignals/frostdb/parts"
 )
@@ -131,7 +130,7 @@ func (db *DB) snapshot(ctx context.Context, async bool, onSuccess func()) {
 		return
 	}
 
-	tx, _, commit := db.begin()
+	tx := db.beginRead()
 	level.Debug(db.logger).Log(
 		"msg", "starting a new snapshot",
 		"tx", tx,
@@ -139,28 +138,6 @@ func (db *DB) snapshot(ctx context.Context, async bool, onSuccess func()) {
 	doSnapshot := func(writeSnapshot func(context.Context, io.Writer) error) {
 		start := time.Now()
 		defer db.snapshotInProgress.Store(false)
-		defer commit()
-		if db.columnStore.enableWAL {
-			// Appending a snapshot record to the WAL is necessary,
-			// since the WAL expects a 1:1 relationship between txn ids
-			// and record indexes. This is done before the actual snapshot so
-			// that a failure to snapshot still appends a record to the WAL,
-			// avoiding a WAL deadlock.
-			if err := db.wal.Log(
-				tx,
-				&walpb.Record{
-					Entry: &walpb.Entry{
-						EntryType: &walpb.Entry_Snapshot_{Snapshot: &walpb.Entry_Snapshot{Tx: tx}},
-					},
-				},
-			); err != nil {
-				level.Error(db.logger).Log(
-					"msg", "failed to append snapshot record to WAL", "err", err,
-				)
-				return
-			}
-		}
-
 		if err := db.snapshotAtTX(ctx, tx, writeSnapshot); err != nil {
 			level.Error(db.logger).Log(
 				"msg", "failed to snapshot database", "err", err,
@@ -187,7 +164,7 @@ func (db *DB) snapshotAtTX(ctx context.Context, tx uint64, writeSnapshot func(co
 	var fileSize int64
 	start := time.Now()
 	if err := func() error {
-		snapshotsDir := snapshotDir(db, tx)
+		snapshotsDir := SnapshotDir(db, tx)
 		if err := os.MkdirAll(snapshotsDir, dirPerms); err != nil {
 			return err
 		}
@@ -801,10 +778,10 @@ func restoreIndexFilesFromSnapshot(db *DB, table, snapshotDir, blockID string) e
 	})
 }
 
-func snapshotDir(db *DB, tx uint64) string {
+func SnapshotDir(db *DB, tx uint64) string {
 	return filepath.Join(db.snapshotsDir(), fmt.Sprintf("%020d", tx))
 }
 
 func snapshotIndexDir(db *DB, tx uint64, table, block string) string {
-	return filepath.Join(snapshotDir(db, tx), "index", table, block)
+	return filepath.Join(SnapshotDir(db, tx), "index", table, block)
 }
diff --git a/table.go b/table.go
index 18996fc..c5a960b 100644
--- a/table.go
+++ b/table.go
@@ -516,26 +516,32 @@ func (t *Table) writeBlock(block *TableBlock, skipPersist, snapshotDB bool) {
 		return
 	}
 
-	tx, _, commit := t.db.begin()
-	defer commit()
+	if err := func() error {
+		tx, _, commit := t.db.begin()
+		defer commit()
 
-	buf, err := block.ulid.MarshalBinary()
-	if err != nil {
-		level.Error(t.logger).Log("msg", "failed to record block persistence in WAL: marshal ulid", "err", err)
-		return
-	}
+		buf, err := block.ulid.MarshalBinary()
+		if err != nil {
+			level.Error(t.logger).Log("msg", "failed to record block persistence in WAL: marshal ulid", "err", err)
+			return err
+		}
 
-	if err := t.wal.Log(tx, &walpb.Record{
-		Entry: &walpb.Entry{
-			EntryType: &walpb.Entry_TableBlockPersisted_{
-				TableBlockPersisted: &walpb.Entry_TableBlockPersisted{
-					TableName: t.name,
-					BlockId:   buf,
+		if err := t.wal.Log(tx, &walpb.Record{
+			Entry: &walpb.Entry{
+				EntryType: &walpb.Entry_TableBlockPersisted_{
+					TableBlockPersisted: &walpb.Entry_TableBlockPersisted{
+						TableName: t.name,
+						BlockId:   buf,
+					},
 				},
 			},
-		},
-	}); err != nil {
-		level.Error(t.logger).Log("msg", "failed to record block persistence in WAL", "err", err)
+		}); err != nil {
+			level.Error(t.logger).Log("msg", "failed to record block persistence in WAL", "err", err)
+			return err
+		}
+
+		return nil
+	}(); err != nil {
 		return
 	}
 
@@ -577,6 +583,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()
+			tx := t.db.beginRead()
 			if err := t.db.snapshotAtTX(ctx, tx, t.db.snapshotWriter(tx)); err != nil {
 				level.Error(t.logger).Log(
 					"msg", "failed to write snapshot on block rotation",
@@ -766,6 +773,16 @@ func (t *Table) Iterator(
 		return errors.New("no callbacks provided")
 	}
 	rowGroups := make(chan any, len(callbacks)*4) // buffer up to 4 row groups per callback
+	defer func() {                                // Drain the channel of any leftover parts due to cancellation or error
+		for rg := range rowGroups {
+			switch t := rg.(type) {
+			case index.ReleaseableRowGroup:
+				t.Release()
+			case arrow.Record:
+				t.Release()
+			}
+		}
+	}()
 
 	// Previously we sorted all row groups into a single row group here,
 	// but it turns out that none of the downstream uses actually rely on
@@ -892,6 +909,16 @@ func (t *Table) SchemaIterator(
 	}
 
 	rowGroups := make(chan any, len(callbacks)*4) // buffer up to 4 row groups per callback
+	defer func() {                                // Drain the channel of any leftover parts due to cancellation or error
+		for rg := range rowGroups {
+			switch t := rg.(type) {
+			case index.ReleaseableRowGroup:
+				t.Release()
+			case arrow.Record:
+				t.Release()
+			}
+		}
+	}()
 
 	schema := arrow.NewSchema(
 		[]arrow.Field{
@@ -1189,12 +1216,8 @@ func (t *Table) collectRowGroups(
 	}()
 	for _, block := range memoryBlocks {
 		if err := block.index.Scan(ctx, "", t.schema, filterExpr, tx, func(ctx context.Context, v any) error {
-			select {
-			case rowGroups <- v:
-				return nil
-			case <-ctx.Done():
-				return ctx.Err()
-			}
+			rowGroups <- v
+			return nil
 		}); err != nil {
 			return err
 		}
@@ -1208,12 +1231,8 @@ func (t *Table) collectRowGroups(
 	for _, source := range t.db.sources {
 		span.AddEvent(fmt.Sprintf("source/%s", source.String()))
 		if err := source.Scan(ctx, filepath.Join(t.db.name, t.name), t.schema, filterExpr, lastBlockTimestamp, func(ctx context.Context, v any) error {
-			select {
-			case rowGroups <- v:
-				return nil
-			case <-ctx.Done():
-				return ctx.Err()
-			}
+			rowGroups <- v
+			return nil
 		}); err != nil {
 			return err
 		}


Test patch:
diff --git a/db_test.go b/db_test.go
index 38e0623..3f8c624 100644
--- a/db_test.go
+++ b/db_test.go
@@ -1408,7 +1408,6 @@ func TestDBRecover(t *testing.T) {
 		require.NoError(t, err)
 
 		// Insert 3 txns.
-		var lastWriteTx uint64
 		for i := 0; i < numInserts; i++ {
 			samples := dynparquet.NewTestSamples()
 			for j := range samples {
@@ -1416,15 +1415,20 @@ func TestDBRecover(t *testing.T) {
 			}
 			r, err := samples.ToRecord()
 			require.NoError(t, err)
-			writeTx, err := table.InsertRecord(ctx, r)
+			_, err = table.InsertRecord(ctx, r)
 			require.NoError(t, err)
+			// Wait until a snapshot is written for each write (it is the txn
+			// immediately preceding the write). This has to be done in a loop,
+			// otherwise writes may not cause a snapshot given that there
+			// might be a snapshot in progress.
 			if i > 0 {
-				// Wait until a snapshot is written for each write (it is the txn
-				// immediately preceding the write). This has to be done in a loop,
-				// otherwise writes may not cause a snapshot given that there
-				// might be a snapshot in progress.
-				db.Wait(writeTx - 1)
-				lastWriteTx = writeTx
+				require.Eventually(t, func() bool {
+					files, err := os.ReadDir(db.snapshotsDir())
+					if err != nil {
+						return false
+					}
+					return len(files) == i && !db.snapshotInProgress.Load()
+				}, 30*time.Second, 100*time.Millisecond)
 			}
 		}
 		// At this point, there should be 2 snapshots. One was triggered before
@@ -1432,8 +1436,14 @@ func TestDBRecover(t *testing.T) {
 		if blockRotation {
 			// A block rotation should trigger the third snapshot.
 			require.NoError(t, table.RotateBlock(ctx, table.ActiveBlock(), false))
-			// Wait for both the new block txn, and the old block rotation txn.
-			db.Wait(lastWriteTx + 2)
+			// Wait for the snapshot to complete
+			require.Eventually(t, func() bool {
+				files, err := os.ReadDir(db.snapshotsDir())
+				if err != nil {
+					return false
+				}
+				return len(files) == 3 && !db.snapshotInProgress.Load()
+			}, 30*time.Second, 100*time.Millisecond)
 		}
 
 		// Verify that there are now 3 snapshots and their txns.
@@ -1445,9 +1455,9 @@ func TestDBRecover(t *testing.T) {
 			require.NoError(t, err)
 			snapshotTxns = append(snapshotTxns, tx)
 		}
-		expectedSnapshots := []uint64{3, 5}
+		expectedSnapshots := []uint64{2, 3}
 		if blockRotation {
-			expectedSnapshots = append(expectedSnapshots, 8)
+			expectedSnapshots = append(expectedSnapshots, 6)
 		}
 		require.Equal(t, expectedSnapshots, snapshotTxns)
 		return dir