Repository: polarsignals/frostdb (go)

Issue:
wal: truncations can sometimes unexpectedly fail
In #428, CI uncovered an interesting edge case in our WAL usage. We have two goroutines: one in https://github.com/polarsignals/frostdb/blob/06fe1fc7c4904417cfbe9bf89fb4acb0e9b190e3/wal/wal.go#L178
and another arbitrary goroutine that can perform truncations concurrently:
https://github.com/polarsignals/frostdb/blob/06fe1fc7c4904417cfbe9bf89fb4acb0e9b190e3/wal/wal.go#L282

The edge case is that a goroutine logs a record with index `k` (e.g. block persistence) and then performs a truncation at that same index. If the `run` goroutine has not read and logged the record with this index, the truncation returns an `out of range` error, since the expected index has not been persisted yet.

To fix these cases, we should move truncation into the `run` goroutine so that it can be executed in the expected order.


Discussion:

Patch:
diff --git a/wal/wal.go b/wal/wal.go
--- a/wal/wal.go
+++ b/wal/wal.go
@@ -71,10 +71,17 @@ type FileWAL struct {
 
 	metrics        *fileWALMetrics
 	logRequestCh   chan *logRequest
-	queue          *logRequestQueue
 	logRequestPool *sync.Pool
 	arrowBufPool   *sync.Pool
-	mtx            *sync.Mutex
+	protected      struct {
+		sync.Mutex
+		queue logRequestQueue
+		// truncateTx is set when the caller wishes to perform a truncation. The
+		// WAL will keep on logging records up to and including this txn and
+		// then perform a truncation. If another truncate call occurs in the
+		// meantime, the highest txn will be used.
+		truncateTx uint64
+	}
 
 	cancel     func()
 	shutdownCh chan struct{}
@@ -139,8 +146,6 @@ func Open(
 				return &bytes.Buffer{}
 			},
 		},
-		mtx:   &sync.Mutex{},
-		queue: &logRequestQueue{},
 		metrics: &fileWALMetrics{
 			// TODO(asubiotto): Move these metrics to the underlying WAL
 			// package.
@@ -179,26 +184,29 @@ func (w *FileWAL) run(ctx context.Context) {
 	ticker := time.NewTicker(50 * time.Millisecond)
 	defer ticker.Stop()
 	walBatch := make([]types.LogEntry, 0)
-	batch := make([]*logRequest, 0, 128) // random number is random
+	batch := make([]*logRequest, 0, 128)
 	// lastQueueSize is only used on shutdown to reduce debug logging verbosity.
 	lastQueueSize := 0
 	for {
 		select {
 		case <-ctx.Done():
-			w.mtx.Lock()
-			n := w.queue.Len()
-			w.mtx.Unlock()
+			w.protected.Lock()
+			n := w.protected.queue.Len()
+			w.protected.Unlock()
 			if n > 0 {
 				// Need to drain the queue before we can shutdown.
 				if n == lastQueueSize {
 					continue
 				}
 				lastQueueSize = n
+				w.protected.Lock()
+				minTx := w.protected.queue[0].tx
+				w.protected.Unlock()
 				lastIdx, err := w.log.LastIndex()
 				logOpts := []any{
 					"msg", "WAL received shutdown request; waiting for log request queue to drain",
 					"queueSize", n,
-					"minTx", (*w.queue)[0].tx,
+					"minTx", minTx,
 					"lastIndex", lastIdx,
 				}
 				if err != nil {
@@ -221,9 +229,12 @@ func (w *FileWAL) run(ctx context.Context) {
 				)
 			}
 			nextTx++
-			w.mtx.Lock()
-			for w.queue.Len() > 0 {
-				if minTx := (*w.queue)[0].tx; minTx != nextTx {
+			w.protected.Lock()
+			// truncateTx will be non-zero if the log should be truncated after
+			// a successful call to StoreLogs.
+			truncateTx := uint64(0)
+			for w.protected.queue.Len() > 0 {
+				if minTx := w.protected.queue[0].tx; minTx != nextTx {
 					if minTx < nextTx {
 						// The next entry must be dropped otherwise progress
 						// will never be made. Log a warning given this could
@@ -233,15 +244,19 @@ func (w *FileWAL) run(ctx context.Context) {
 							"expected", nextTx,
 							"found", minTx,
 						)
-						_ = heap.Pop(w.queue)
+						_ = heap.Pop(&w.protected.queue)
 					}
 					break
 				}
-				r := heap.Pop(w.queue).(*logRequest)
+				r := heap.Pop(&w.protected.queue).(*logRequest)
 				batch = append(batch, r)
+				if r.tx == w.protected.truncateTx {
+					truncateTx = r.tx
+					w.protected.truncateTx = 0
+				}
 				nextTx++
 			}
-			w.mtx.Unlock()
+			w.protected.Unlock()
 			if len(batch) == 0 {
 				continue
 			}
@@ -267,6 +282,19 @@ func (w *FileWAL) run(ctx context.Context) {
 				)
 			}
 
+			if truncateTx != 0 {
+				w.metrics.lastTruncationAt.Set(float64(truncateTx))
+				w.metrics.walTruncations.Inc()
+
+				level.Debug(w.logger).Log("msg", "truncating WAL", "tx", truncateTx)
+				if err := w.log.TruncateFront(truncateTx); err != nil {
+					level.Error(w.logger).Log("msg", "failed to truncate WAL", "tx", truncateTx, "err", err)
+					w.metrics.walTruncationsFailed.Inc()
+				} else {
+					level.Debug(w.logger).Log("msg", "truncated WAL", "tx", truncateTx)
+				}
+			}
+
 			for _, r := range batch {
 				w.logRequestPool.Put(r)
 			}
@@ -274,19 +302,15 @@ func (w *FileWAL) run(ctx context.Context) {
 	}
 }
 
+// Truncate queues a truncation of the WAL at the given tx. Note that the
+// truncation will be performed asynchronously. A nil error does not indicate
+// a successful truncation.
 func (w *FileWAL) Truncate(tx uint64) error {
-	w.metrics.lastTruncationAt.Set(float64(tx))
-	w.metrics.walTruncations.Inc()
-
-	level.Debug(w.logger).Log("msg", "truncating WAL", "tx", tx)
-	err := w.log.TruncateFront(tx)
-	if err != nil {
-		level.Error(w.logger).Log("msg", "failed to truncate WAL", "tx", tx, "err", err)
-		w.metrics.walTruncationsFailed.Inc()
-		return err
+	w.protected.Lock()
+	defer w.protected.Unlock()
+	if tx > w.protected.truncateTx {
+		w.protected.truncateTx = tx
 	}
-	level.Debug(w.logger).Log("msg", "truncated WAL", "tx", tx)
-
 	return nil
 }
 
@@ -313,9 +337,9 @@ func (w *FileWAL) Log(tx uint64, record *walpb.Record) error {
 		return err
 	}
 
-	w.mtx.Lock()
-	heap.Push(w.queue, r)
-	w.mtx.Unlock()
+	w.protected.Lock()
+	heap.Push(&w.protected.queue, r)
+	w.protected.Unlock()
 
 	return nil
 }
@@ -370,9 +394,9 @@ func (w *FileWAL) LogRecord(tx uint64, table string, record arrow.Record) error
 		return err
 	}
 
-	w.mtx.Lock()
-	heap.Push(w.queue, r)
-	w.mtx.Unlock()
+	w.protected.Lock()
+	heap.Push(&w.protected.queue, r)
+	w.protected.Unlock()
 
 	return nil
 }


Test patch:
diff --git a/wal/wal_test.go b/wal/wal_test.go
--- a/wal/wal_test.go
+++ b/wal/wal_test.go
@@ -1,9 +1,11 @@
 package wal
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
@@ -155,3 +157,47 @@ func TestUnexpectedTxn(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, lastIndex, uint64(2))
 }
+
+func TestWALTruncate(t *testing.T) {
+	dir := t.TempDir()
+	w, err := Open(
+		log.NewNopLogger(),
+		prometheus.NewRegistry(),
+		dir,
+	)
+	require.NoError(t, err)
+	defer w.Close()
+	w.RunAsync()
+
+	for i := uint64(0); i < 10; i++ {
+		require.NoError(t, w.Log(i, &walpb.Record{
+			Entry: &walpb.Entry{
+				EntryType: &walpb.Entry_Write_{
+					Write: &walpb.Entry_Write{
+						Data:      []byte(fmt.Sprintf("test-data-%d", i)),
+						TableName: "test-table",
+					},
+				},
+			},
+		}))
+	}
+	require.NoError(t, w.Truncate(9))
+
+	// Wait for the WAL to asynchronously log and truncate.
+	require.Eventually(t, func() bool {
+		tx, _ := w.FirstIndex()
+		return tx == 9
+	}, time.Second, 10*time.Millisecond)
+
+	numRecords := 0
+	require.NoError(
+		t,
+		w.Replay(0, func(tx uint64, r *walpb.Record) error {
+			numRecords++
+			require.Equal(t, uint64(9), tx)
+			require.Equal(t, []byte("test-data-9"), r.Entry.GetWrite().Data)
+			return nil
+		}),
+	)
+	require.Equal(t, 1, numRecords)
+}