Repository: polarsignals/frostdb (go)

Issue:
wal: time out Close
I've found on many instances that the wal package is stuck closing because the next minimum transaction is not the next expected transaction, causing an endless loop waiting for the WAL to drain:
https://github.com/polarsignals/frostdb/blob/336457ecf250723103a7ede2a011377280143206/wal/wal.go#L185-L204

I think it would be nice to put a timeout (maybe after a second of no new records logged) so that the WAL closes and logs an error even in this invalid state. It might be nice to also have metrics on this invalid close, because it's usually indicative of a bug.

Having this hard shutdown helps when running frostdb on pre-emptible machines that have a max time to shutdown. If the WAL doesn't shutdown in time, the whole machine is shut down, skipping other steps of a graceful shutdown which are nice to have (and would even obviate this invalid WAL state, e.g. if we take a snapshot or persist the data).

Pull request description:
wal: time out Close 
Fixes #486 

Patch:
diff --git a/wal/wal.go b/wal/wal.go
index 2ce7949..ce5db01 100644
--- a/wal/wal.go
+++ b/wal/wal.go
@@ -63,6 +63,7 @@ type fileWALMetrics struct {
 	lastTruncationAt      prometheus.Gauge
 	walRepairs            prometheus.Counter
 	walRepairsLostRecords prometheus.Counter
+	walCloseTimeouts      prometheus.Counter
 }
 
 const dirPerms = os.FileMode(0o750)
@@ -89,8 +90,9 @@ type FileWAL struct {
 		nextTx uint64
 	}
 
-	cancel     func()
-	shutdownCh chan struct{}
+	cancel       func()
+	shutdownCh   chan struct{}
+	closeTimeout time.Duration
 }
 
 type logRequest struct {
@@ -157,6 +159,7 @@ func Open(
 				return &bytes.Buffer{}
 			},
 		},
+		closeTimeout: 1 * time.Second,
 		metrics: &fileWALMetrics{
 			failedLogs: promauto.With(reg).NewCounter(prometheus.CounterOpts{
 				Name: "failed_logs_total",
@@ -174,6 +177,10 @@ func Open(
 				Name: "repairs_lost_records_total",
 				Help: "The number of WAL records lost due to WAL repairs (truncations)",
 			}),
+			walCloseTimeouts: promauto.With(reg).NewCounter(prometheus.CounterOpts{
+				Name: "close_timeouts_total",
+				Help: "The number of times the WAL failed to close due to a timeout",
+			}),
 		},
 		shutdownCh: make(chan struct{}),
 	}
@@ -190,6 +197,9 @@ func (w *FileWAL) run(ctx context.Context) {
 	batch := make([]*logRequest, 0, 128)
 	// lastQueueSize is only used on shutdown to reduce debug logging verbosity.
 	lastQueueSize := 0
+	// lastBatchWrite is used to determine when to force a close of the WAL.
+	lastBatchWrite := time.Now()
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -197,6 +207,15 @@ func (w *FileWAL) run(ctx context.Context) {
 			n := w.protected.queue.Len()
 			w.protected.Unlock()
 			if n > 0 {
+				// force the WAL to close after some a timeout.
+				if w.closeTimeout > 0 && time.Since(lastBatchWrite) > w.closeTimeout {
+					w.metrics.walCloseTimeouts.Inc()
+					level.Error(w.logger).Log(
+						"msg", "WAL timed out attempting to close",
+					)
+					return
+				}
+
 				// Need to drain the queue before we can shutdown.
 				if n == lastQueueSize {
 					continue
@@ -303,6 +322,8 @@ func (w *FileWAL) run(ctx context.Context) {
 			for _, r := range batch {
 				w.logRequestPool.Put(r)
 			}
+
+			lastBatchWrite = time.Now()
 		}
 	}
 }


Test patch:
diff --git a/wal/wal_test.go b/wal/wal_test.go
index 17e090c..fd7311d 100644
--- a/wal/wal_test.go
+++ b/wal/wal_test.go
@@ -269,3 +269,32 @@ func TestWALTruncate(t *testing.T) {
 		require.Equal(t, 1, numRecords)
 	})
 }
+
+func TestWALCloseTimeout(t *testing.T) {
+	dir := t.TempDir()
+	w, err := Open(
+		log.NewNopLogger(),
+		prometheus.NewRegistry(),
+		dir,
+	)
+	require.NoError(t, err)
+
+	w.RunAsync()
+
+	// This will cause the WAL to enter a state where it will not close
+	// b/c it was expecting the next transaction to be 1.
+	require.NoError(t, w.Log(2, &walpb.Record{
+		Entry: &walpb.Entry{
+			EntryType: &walpb.Entry_Write_{
+				Write: &walpb.Entry_Write{
+					Data:      []byte("test-data"),
+					TableName: "test-table",
+				},
+			},
+		},
+	}))
+
+	// This should not block forever, otherwise the test will fail by timeout
+	err = w.Close()
+	require.NoError(t, err)
+}