refactor(ingestor): WAL Replay controller: Use singleflight to ensure only one inflight flush is happening (#19831)

Author: Segflow

pkg/ingester/replay_controller.go

@@ -1,11 +1,10 @@
 package ingester
 
 import (
-	"sync"
-
 	"github.com/dustin/go-humanize"
 	"github.com/go-kit/log/level"
 	"go.uber.org/atomic"
+	"golang.org/x/sync/singleflight"
 
 	util_log "github.com/grafana/loki/v3/pkg/util/log"
 )
@@ -49,17 +48,16 @@ type replayController struct {
 	currentBytes atomic.Int64
 	cfg          WALConfig
 	metrics      *ingesterMetrics
-	cond         *sync.Cond
-	isFlushing   atomic.Bool
-	flusher      Flusher
+
+	flusher Flusher
+	flushSF singleflight.Group
 }
 
 // flusher is expected to reduce pressure via calling Sub
 func newReplayController(metrics *ingesterMetrics, cfg WALConfig, flusher Flusher) *replayController {
 	return &replayController{
 		cfg:     cfg,
 		metrics: metrics,
-		cond:    sync.NewCond(&sync.Mutex{}),
 		flusher: flusher,
 	}
 }
@@ -79,51 +77,41 @@ func (c *replayController) Cur() int {
 }
 
 func (c *replayController) Flush() {
-	if c.isFlushing.CompareAndSwap(false, true) {
-		c.metrics.recoveryIsFlushing.Set(1)
-		prior := c.currentBytes.Load()
-		level.Debug(util_log.Logger).Log(
-			"msg", "replay flusher pre-flush",
-			"bytes", humanize.Bytes(uint64(prior)),
-		)
-
-		c.flusher.Flush()
-
-		after := c.currentBytes.Load()
-		level.Debug(util_log.Logger).Log(
-			"msg", "replay flusher post-flush",
-			"bytes", humanize.Bytes(uint64(after)),
-		)
-
-		c.isFlushing.Store(false)
-		c.metrics.recoveryIsFlushing.Set(0)
-
-		// Broadcast after lock is acquired to prevent race conditions with cpu scheduling
-		// where the flush code could finish before the goroutine which initiated it gets to call
-		// c.cond.Wait()
-		c.cond.L.Lock()
-		c.cond.Broadcast()
-		c.cond.L.Unlock()
-	}
+	// Use singleflight to ensure only one flush happens at a time
+	_, _, _ = c.flushSF.Do("flush", func() (interface{}, error) {
+		c.flush()
+		return nil, nil
+	})
+}
+
+func (c *replayController) flush() {
+	c.metrics.recoveryIsFlushing.Set(1)
+	prior := c.currentBytes.Load()
+	level.Debug(util_log.Logger).Log(
+		"msg", "replay flusher pre-flush",
+		"bytes", humanize.Bytes(uint64(prior)),
+	)
+
+	c.flusher.Flush()
+
+	after := c.currentBytes.Load()
+	level.Debug(util_log.Logger).Log(
+		"msg", "replay flusher post-flush",
+		"bytes", humanize.Bytes(uint64(after)),
+	)
+
+	c.metrics.recoveryIsFlushing.Set(0)
 }
 
 // WithBackPressure is expected to call replayController.Add in the passed function to increase the managed byte count.
 // It will call the function as long as there is expected room before the memory cap and will then flush data intermittently
 // when needed.
 func (c *replayController) WithBackPressure(fn func() error) error {
-	// Account for backpressure and wait until there's enough memory to continue replaying the WAL
-	c.cond.L.Lock()
-
 	// use 90% as a threshold since we'll be adding to it.
 	for c.Cur() > int(c.cfg.ReplayMemoryCeiling)*9/10 {
 		// too much backpressure, flush
-		go c.Flush()
-		c.cond.Wait()
+		c.Flush()
 	}
 
-	// Don't hold the lock while executing the provided function.
-	// This ensures we can run functions concurrently.
-	c.cond.L.Unlock()
-
 	return fn()
 }

pkg/ingester/replay_controller_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 
 	"github.com/grafana/loki/v3/pkg/util/constants"
 )
@@ -75,5 +76,62 @@ func TestReplayController(t *testing.T) {
 		"WithBackPressure", // add 50, total 50
 	}
 	require.Equal(t, expected, ops)
+}
+
+// Test to ensure only one flush happens at a time when multiple goroutines call WithBackPressure
+func TestReplayControllerConcurrentFlushes(t *testing.T) {
+	t.Run("multiple goroutines wait for single flush", func(t *testing.T) {
+		var rc *replayController
+
+		var flushesStarted atomic.Int32
+		var flushInProgress atomic.Bool
+
+		flusher := newDumbFlusher(func() {
+			// Check if a flush is already in progress, fail if there is one.
+			if !flushInProgress.CompareAndSwap(false, true) {
+				t.Error("Multiple flushes running concurrently!")
+			}
+
+			flushesStarted.Add(1)
+
+			time.Sleep(100 * time.Millisecond) // Simulate a slow flush
+
+			rc.Sub(200)
+
+			flushInProgress.Store(false)
+		})
+
+		rc = newReplayController(nilMetrics(), WALConfig{ReplayMemoryCeiling: 100}, flusher)
+
+		// Fill to trigger flush condition (90% of 100 = 90 bytes threshold)
+		rc.Add(95)
+
+		// Launch multiple goroutines simultaneously
+		start := make(chan struct{})
+		var wg sync.WaitGroup
+		numGoroutines := 5
+
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				<-start
+
+				err := rc.WithBackPressure(func() error {
+					rc.Add(20) // Each would trigger flush condition
+					return nil
+				})
+				require.NoError(t, err)
+			}()
+		}
+
+		// Start all goroutines at the same time
+		close(start)
+
+		wg.Wait()
 
+		// All goroutines should have shared a single flush
+		require.Equal(t, int32(1), flushesStarted.Load(),
+			"Singleflight should coalesce all flush requests into one")
+	})
 }