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.

Pull request description:
wal: perform truncation in run
Previously, we would synchronously truncate even though we asynchronously wrote
to the WAL. This caused truncations to fail with high likelihood given that the
transaction to truncate at was likely not yet logged.

Closes #431 

Patch:
diff --git a/wal/wal.go b/wal/wal.go
index 7ced8f5..392981b 100644
--- 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
index 583dc64..cd380fb 100644
--- 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)
+}