fix(livestore): stop periodic WAL flush goroutines before shutdown flush

CHANGELOG.md

@@ -82,6 +82,7 @@
 * [BUGFIX] live-store: fixed unsuccessful deregistering from membership/partition rings during shutdown [#6848](https://github.com/grafana/tempo/pull/6848) (@zhxiaogg)
 * [BUGFIX] fix: respect context cancellation when reading WAL block iterator [#6928](https://github.com/grafana/tempo/pull/6928) (@zhxiaogg)
 * [BUGFIX] Complete lifecycler shutdown on errors [#6906](https://github.com/grafana/tempo/pull/6906) (@javiermolinar)
+* [BUGFIX] livestore: fix concurrent WAL writes from periodic and shutdown flushes [#6972](https://github.com/grafana/tempo/pull/6972) (@zhxiaogg)
 
 ### 3.0 Cleanup
 

modules/livestore/live_store.go

@@ -152,6 +152,9 @@ type LiveStore struct {
 	lagCancel           context.CancelFunc
 	readyErr            atomic.Pointer[error] // nil when ready to serve queries
 	lastRecordTimeNanos atomic.Int64          // stores timestamp of last consumed record as UnixNano, -1 means not set
+
+	cutToWalStop chan struct{}  // closed to stop perTenantCutToWalLoop goroutines before shutdown flush
+	cutToWalWg   sync.WaitGroup // tracks active perTenantCutToWalLoop goroutines
 }
 
 func New(cfg Config, overridesService overrides.Interface, completeBlockFlusher completeBlockFlusher, logger log.Logger, reg prometheus.Registerer) (*LiveStore, error) {
@@ -180,6 +183,7 @@ func New(cfg Config, overridesService overrides.Interface, completeBlockFlusher
 		completeBlockLifecycle: completeBlockLifecycle,
 		completeQueues:         flushqueues.New[*completeOp](metricCompleteQueueLength),
 		startupComplete:        make(chan struct{}),
+		cutToWalStop:           make(chan struct{}),
 	}
 
 	// Initialize ready state to starting
@@ -451,6 +455,10 @@ func (s *LiveStore) stopping(error) error {
 		stopErr = errors.Join(stopErr, err)
 	}
 
+	level.Info(s.logger).Log("msg", "stopping periodic WAL flush goroutines")
+	s.stopAllCutToWalLoops()
+	level.Info(s.logger).Log("msg", "periodic WAL flush goroutines stopped")
+
 	// Flush all data to disk.
 	level.Info(s.logger).Log("msg", "cutting all instances to WAL")
 	s.cutAllInstancesToWal()
@@ -690,9 +698,7 @@ func (s *LiveStore) getOrCreateInstance(tenantID string) (*instance, error) {
 
 	s.instances[tenantID] = inst
 
-	s.runInBackground(func() {
-		s.perTenantCutToWalLoop(inst)
-	})
+	s.startPerTenantCutToWalLoop(inst)
 	s.runInBackground(func() {
 		s.perTenantCleanupLoop(inst)
 	})

modules/livestore/live_store_background.go

@@ -175,19 +175,38 @@ func (s *LiveStore) retryCompleteOp(op *completeOp, span oteltrace.Span, msg str
 	}()
 }
 
-func (s *LiveStore) perTenantCutToWalLoop(instance *instance) {
-	// ticker
-	ticker := time.NewTicker(s.cfg.InstanceFlushPeriod)
-	defer ticker.Stop()
+func (s *LiveStore) startPerTenantCutToWalLoop(inst *instance) {
+	s.cutToWalWg.Add(1)
+	go func() {
+		defer s.cutToWalWg.Done()
 
-	for {
+		// Wait for startup to finish; also listen on cutToWalStop so we can
+		// exit if shutdown happens before startup completes.
 		select {
-		case <-ticker.C:
-			s.cutOneInstanceToWal(s.ctx, instance, false)
-		case <-s.ctx.Done():
+		case <-s.startupComplete:
+		case <-s.cutToWalStop:
 			return
 		}
-	}
+
+		ticker := time.NewTicker(s.cfg.InstanceFlushPeriod)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				s.cutOneInstanceToWal(s.ctx, inst, false)
+			case <-s.cutToWalStop:
+				return
+			case <-s.ctx.Done():
+				return
+			}
+		}
+	}()
+}
+
+func (s *LiveStore) stopAllCutToWalLoops() {
+	close(s.cutToWalStop)
+	s.cutToWalWg.Wait()
 }
 
 func (s *LiveStore) perTenantCleanupLoop(inst *instance) {

modules/livestore/live_store_test.go

@@ -148,8 +148,9 @@ func TestLiveStorePushBytesRejectsWhenStopping(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, liveStore)
 
-	// Transition to stopping state, then verify writes are rejected.
-	_ = liveStore.stopping(nil)
+	// Stop the service so writes are rejected.
+	err = services.StopAndAwaitTerminated(context.Background(), liveStore)
+	require.NoError(t, err)
 
 	id := test.ValidTraceID(nil)
 	expectedTrace := test.MakeTrace(1, id)
@@ -700,7 +701,7 @@ func TestLiveStoreShutdownWithPendingCompletions(t *testing.T) {
 	requireTraceInLiveStore(t, liveStore, expectedID, expectedTrace)
 	requireInstanceState(t, inst, instanceState{liveTraces: 1, walBlocks: 0, completeBlocks: 0})
 
-	require.NoError(t, liveStore.stopping(nil))
+	require.NoError(t, services.StopAndAwaitTerminated(context.Background(), liveStore))
 }
 
 func TestLiveStoreQueryMethodsBeforeStarted(t *testing.T) {
@@ -843,8 +844,8 @@ func TestLiveStoreQueryMethodsAfterStopping(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, liveStore)
 
-	// Error expected from Kafka reader shutdown; we only care about query behavior after stopping begins.
-	_ = liveStore.stopping(nil)
+	// Stop the service so queries are rejected.
+	_ = services.StopAndAwaitTerminated(context.Background(), liveStore)
 
 	ctx := user.InjectOrgID(context.Background(), testTenantID)
 
@@ -876,8 +877,8 @@ func TestLiveStoreQueryMethodsAfterStoppingWithFailOnHighLag(t *testing.T) {
 
 	liveStore.cfg.FailOnHighLag = true
 
-	// Error expected from Kafka reader shutdown; we only care about query behavior after stopping begins.
-	_ = liveStore.stopping(nil)
+	// Stop the service so queries are rejected.
+	_ = services.StopAndAwaitTerminated(context.Background(), liveStore)
 
 	ctx := user.InjectOrgID(context.Background(), testTenantID)