logpuller: fix potential shutdown hang in paused region event push

logservice/logpuller/region_event_sink.go

@@ -26,16 +26,21 @@ import (
 
 // regionEventSink delivers region events to dynstream and owns push-side flow control.
 type regionEventSink struct {
-	ctx context.Context
-	ds  dynstream.DynamicStream[int, SubscriptionID, regionEvent, *subscribedSpan, *regionEventHandler]
-	// the following three fields are used to manage feedback from ds and notify other goroutines
-	mu     sync.Mutex
-	cond   *sync.Cond
+	// mu/cond coordinate the paused push path with pause/resume and shutdown signals.
+	mu   sync.Mutex
+	cond *sync.Cond
+	// paused tracks whether region event pushing is temporarily held back by feedback.
 	paused atomic.Bool
+	// stopped marks the sink as shutting down so blocked pushers can exit instead of waiting for resume.
+	stopped atomic.Bool
+
+	// ds owns the dynstream used to deliver region events and receive flow-control feedback.
+	ds dynstream.DynamicStream[int, SubscriptionID, regionEvent, *subscribedSpan, *regionEventHandler]
 }
 
-func newRegionEventSink(ctx context.Context, failureHandler *regionFailureHandler) *regionEventSink {
-	sink := &regionEventSink{ctx: ctx}
+func newRegionEventSink(failureHandler *regionFailureHandler) *regionEventSink {
+	sink := &regionEventSink{}
+	sink.cond = sync.NewCond(&sink.mu)
 
 	option := dynstream.NewOption()
 	// Note: it is max batch size of the kv sent from tikv(not committed rows)
@@ -51,7 +56,6 @@ func newRegionEventSink(ctx context.Context, failureHandler *regionFailureHandle
 	)
 	ds.Start()
 	sink.ds = ds
-	sink.cond = sync.NewCond(&sink.mu)
 	return sink
 }
 
@@ -73,6 +77,9 @@ func (s *regionEventSink) Wake(subID SubscriptionID) {
 }
 
 func (s *regionEventSink) Push(subID SubscriptionID, event regionEvent) {
+	if s.stopped.Load() {
+		return
+	}
 	// fast path
 	if !s.paused.Load() {
 		s.ds.Push(subID, event)
@@ -81,36 +88,31 @@ func (s *regionEventSink) Push(subID SubscriptionID, event regionEvent) {
 
 	// slow path: wait until paused is false
 	s.mu.Lock()
-	for s.paused.Load() {
-		select {
-		case <-s.ctx.Done():
-			s.mu.Unlock()
-			return
-		default:
-			s.cond.Wait()
-		}
+	for s.paused.Load() && !s.stopped.Load() {
+		s.cond.Wait()
 	}
+	stopped := s.stopped.Load()
 	s.mu.Unlock()
+
+	if stopped {
+		return
+	}
 	s.ds.Push(subID, event)
 }
 
 func (s *regionEventSink) Run(ctx context.Context) error {
 	for {
 		select {
 		case <-ctx.Done():
+			s.stop()
 			return nil
 		case feedback := <-s.ds.Feedback():
 			switch feedback.FeedbackType {
 			case dynstream.PauseArea:
-				s.mu.Lock()
-				s.paused.Store(true)
-				s.mu.Unlock()
+				s.pause()
 				log.Info("subscription client pause push region event")
 			case dynstream.ResumeArea:
-				s.mu.Lock()
-				s.paused.Store(false)
-				s.cond.Broadcast()
-				s.mu.Unlock()
+				s.resume()
 				log.Info("subscription client resume push region event")
 			case dynstream.ReleasePath, dynstream.ResumePath:
 				// Ignore it, because it is no need to pause and resume a path in puller.
@@ -148,9 +150,35 @@ func (s *regionEventSink) UpdateMetrics() {
 }
 
 func (s *regionEventSink) Close() {
+	s.stop()
+	s.ds.Close()
+}
+
+func (s *regionEventSink) pause() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.stopped.Load() || s.paused.Load() {
+		return
+	}
+	s.paused.Store(true)
+}
+
+func (s *regionEventSink) resume() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if !s.paused.Load() {
+		return
+	}
+	s.paused.Store(false)
+	s.cond.Broadcast()
+}
+
+func (s *regionEventSink) stop() {
+	if !s.stopped.CompareAndSwap(false, true) {
+		return
+	}
 	s.mu.Lock()
 	s.paused.Store(false)
 	s.cond.Broadcast()
 	s.mu.Unlock()
-	s.ds.Close()
 }

logservice/logpuller/region_event_sink_test.go

@@ -71,16 +71,20 @@ func (s *mockRegionEventSinkStream) GetMetrics() dynstream.Metrics[int, Subscrip
 	return s.metrics
 }
 
+func newTestRegionEventSink(
+	ds dynstream.DynamicStream[int, SubscriptionID, regionEvent, *subscribedSpan, *regionEventHandler],
+) *regionEventSink {
+	sink := &regionEventSink{ds: ds}
+	sink.cond = sync.NewCond(&sink.mu)
+	return sink
+}
+
 func TestRegionEventSinkRunPausesAndResumesPush(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	ds := newMockRegionEventSinkStream()
-	sink := &regionEventSink{
-		ctx: ctx,
-		ds:  ds,
-	}
-	sink.cond = sync.NewCond(&sink.mu)
+	sink := newTestRegionEventSink(ds)
 
 	runErrCh := make(chan error, 1)
 	go func() {
@@ -152,10 +156,8 @@ func TestRegionEventSinkUpdateMetrics(t *testing.T) {
 		).Set(456)
 
 		sink := &regionEventSink{
-			ctx: context.Background(),
-			ds:  ds,
+			ds: ds,
 		}
-		sink.cond = sync.NewCond(&sink.mu)
 		sink.UpdateMetrics()
 
 		require.Equal(t, float64(11), testutil.ToFloat64(metricSubscriptionClientDSChannelSize))
@@ -191,10 +193,8 @@ func TestRegionEventSinkUpdateMetrics(t *testing.T) {
 		}
 
 		sink := &regionEventSink{
-			ctx: context.Background(),
-			ds:  ds,
+			ds: ds,
 		}
-		sink.cond = sync.NewCond(&sink.mu)
 		sink.UpdateMetrics()
 
 		require.Equal(t, float64(33), testutil.ToFloat64(metricSubscriptionClientDSChannelSize))
@@ -213,3 +213,49 @@ func TestRegionEventSinkUpdateMetrics(t *testing.T) {
 		)))
 	})
 }
+
+func TestRegionEventSinkRunCancelUnblocksPush(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	ds := newMockRegionEventSinkStream()
+	sink := newTestRegionEventSink(ds)
+
+	runErrCh := make(chan error, 1)
+	go func() {
+		runErrCh <- sink.Run(ctx)
+	}()
+
+	ds.feedbackCh <- dynstream.Feedback[int, SubscriptionID, *subscribedSpan]{
+		FeedbackType: dynstream.PauseArea,
+	}
+	require.Eventually(t, sink.paused.Load, time.Second, 10*time.Millisecond)
+
+	pushDone := make(chan struct{})
+	go func() {
+		sink.Push(SubscriptionID(1), regionEvent{resolvedTs: 100})
+		close(pushDone)
+	}()
+
+	select {
+	case <-pushDone:
+		t.Fatal("Push should block while the sink is paused")
+	case <-time.After(100 * time.Millisecond):
+	}
+	require.Equal(t, int32(0), ds.pushCount.Load())
+
+	cancel()
+
+	select {
+	case <-pushDone:
+	case <-time.After(time.Second):
+		t.Fatal("Push should be unblocked by Run context cancellation")
+	}
+	require.Equal(t, int32(0), ds.pushCount.Load())
+
+	select {
+	case err := <-runErrCh:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("Run should exit after context cancellation")
+	}
+}

logservice/logpuller/region_request_worker_test.go

@@ -157,7 +157,7 @@ func newDispatchResolvedTsTestWorker(regionCount int) (*regionRequestWorker, *mo
 			metrics: sharedClientMetrics{
 				batchResolvedSize: prometheus.ObserverFunc(func(float64) {}),
 			},
-			eventSink: &regionEventSink{ds: ds},
+			eventSink: newTestRegionEventSink(ds),
 		},
 	}
 	worker.requestedRegions.subscriptions = map[SubscriptionID]regionFeedStates{

logservice/logpuller/subscription_client.go

@@ -180,7 +180,7 @@ func NewSubscriptionClient(
 	}
 	subClient.ctx, subClient.cancel = context.WithCancel(context.Background())
 	subClient.failureHandler = newRegionFailureHandler(subClient)
-	subClient.eventSink = newRegionEventSink(subClient.ctx, subClient.failureHandler)
+	subClient.eventSink = newRegionEventSink(subClient.failureHandler)
 	subClient.spanRegistry = newSpanRegistry(subClient.pd, subClient.pdClock)
 
 	subClient.initMetrics()

logservice/logpuller/subscription_client_test.go

@@ -343,7 +343,7 @@ func TestStopTaskUsesSubscribedSpanFilterLoop(t *testing.T) {
 
 func TestOnRegionFailQueuesCanceledErrorCache(t *testing.T) {
 	client := &subscriptionClient{
-		eventSink: &regionEventSink{ds: &mockDynamicStream{}},
+		eventSink: newTestRegionEventSink(&mockDynamicStream{}),
 	}
 	client.spanRegistry = newSpanRegistry(nil, nil)
 	client.failureHandler = newRegionFailureHandler(client)
@@ -417,11 +417,7 @@ func (s *mockDynamicStream) GetMetrics() dynstream.Metrics[int, SubscriptionID]
 }
 
 func TestPushRegionEventToDSUnblocksOnClose(t *testing.T) {
-	sink := &regionEventSink{
-		ctx: context.Background(),
-		ds:  &mockDynamicStream{},
-	}
-	sink.cond = sync.NewCond(&sink.mu)
+	sink := newTestRegionEventSink(&mockDynamicStream{})
 	client := &subscriptionClient{
 		eventSink:       sink,
 		regionTaskQueue: priorityqueue.New[PriorityTask](),