Repository: tikv/pd (2025-08-01T11:02:56Z)

Issue:
Bug Report: Goroutine Leak in TSO Dispatcher When Stream Errors Occur
## Bug Report


### What did you do?

Enable tso proxy,  the pd follower has goroutine leak.
When stream errors during request processing were causing goroutine leaks. Specifically, when a stream error occurred, pending requests were not being properly cleaned up, leaving goroutines in a blocked state waiting for responses that would never arrive.

In environments with unstable network connections or frequent stream errors, this issue would lead to a continuous increase in goroutine count and memory usage over time.

### What did you expect to see?

When stream errors occur, all pending requests should be promptly cancelled, and associated goroutines should be released to prevent resource leaks. The system should maintain stable memory usage even under error conditions.

### What did you see instead?

Pending requests were not being properly cancelled when stream errors occurred, causing goroutines to remain blocked indefinitely. This led to a gradual increase in memory usage and potential performance degradation over time, especially in long-running systems with frequent stream errors.

<img width="962" height="687" alt="Image" src="https://github.com/user-attachments/assets/0fe66ce8-ce2a-450f-8552-6f4525d04f64" />

### What version of PD are you using (`pd-server -V`)?

all version since v5.3.0



Discussion:
Welcome @felixYing! It looks like this is your first issue to tikv/pd 🎉
> What version of PD are you using (pd-server -V)?
> all version since v5.3.0

@felixYing there are quite a few fixes in master even beyond what you find in 8.5.2. I'd recommend to repro it on master if it is still that bad.



Patch:
diff --git a/pkg/utils/tsoutil/tso_dispatcher.go b/pkg/utils/tsoutil/tso_dispatcher.go
index ba3c1909b1c..52cd74ca705 100644
--- a/pkg/utils/tsoutil/tso_dispatcher.go
+++ b/pkg/utils/tsoutil/tso_dispatcher.go
@@ -75,19 +75,21 @@ func (s *TSODispatcher) DispatchRequest(serverCtx context.Context, req Request,
 	key := req.getForwardedHost()
 	val, loaded := s.dispatchChs.Load(key)
 	if !loaded {
-		val = &tsoRequestProxyQueue{requestCh: make(chan Request, maxMergeRequests+1)}
-		val, loaded = s.dispatchChs.LoadOrStore(key, val)
-	}
-	tsoQueue := val.(*tsoRequestProxyQueue)
-	if !loaded {
-		log.Info("start new tso proxy dispatcher", zap.String("forwarded-host", req.getForwardedHost()))
-		tsDeadlineCh := make(chan *TSDeadline, 1)
 		dispatcherCtx, ctxCancel := context.WithCancelCause(serverCtx)
-		tsoQueue.ctx = dispatcherCtx
-		tsoQueue.cancel = ctxCancel
-		go s.dispatch(tsoQueue, tsoProtoFactory, req.getForwardedHost(), req.getClientConn(), tsDeadlineCh, tsoPrimaryWatchers...)
-		go WatchTSDeadline(dispatcherCtx, tsDeadlineCh)
+		tsoQueue := &tsoRequestProxyQueue{
+			ctx:       dispatcherCtx,
+			cancel:    ctxCancel,
+			requestCh: make(chan Request, maxMergeRequests+1),
+		}
+		val, loaded = s.dispatchChs.LoadOrStore(key, tsoQueue)
+		if !loaded {
+			log.Info("start new tso proxy dispatcher", zap.String("forwarded-host", req.getForwardedHost()))
+			tsDeadlineCh := make(chan *TSDeadline, 1)
+			go s.dispatch(tsoQueue, tsoProtoFactory, req.getForwardedHost(), req.getClientConn(), tsDeadlineCh, tsoPrimaryWatchers...)
+			go WatchTSDeadline(tsoQueue.ctx, tsDeadlineCh)
+		}
 	}
+	tsoQueue := val.(*tsoRequestProxyQueue)
 	tsoQueue.requestCh <- req
 	return tsoQueue.ctx
 }
@@ -101,7 +103,13 @@ func (s *TSODispatcher) dispatch(
 	tsoPrimaryWatchers ...*etcdutil.LoopWatcher) {
 	defer logutil.LogPanic()
 	dispatcherCtx := tsoQueue.ctx
-	defer s.dispatchChs.Delete(forwardedHost)
+	// Note: We use clearPendingRequests in a defer statement to ensure proper cleanup when an error occurs.
+	// This function not only deletes the queue from the dispatcher but also clears all pending requests
+	// to prevent goroutine leakage and ensure that all waiting goroutines are notified and can exit gracefully.
+	var err error
+	defer func() {
+		s.clearPendingRequests(tsoQueue, forwardedHost, err)
+	}()
 
 	forwardStream, cancel, err := tsoProtoFactory.createForwardStream(tsoQueue.ctx, clientConn)
 	failpoint.Inject("canNotCreateForwardStream", func() {
@@ -121,7 +129,8 @@ func (s *TSODispatcher) dispatch(
 	}
 	defer cancel()
 
-	requests := make([]Request, maxMergeRequests+1)
+	// make(chan Request, maxMergeRequests+1) + 1
+	requests := make([]Request, maxMergeRequests+2)
 	needUpdateServicePrimaryAddr := len(tsoPrimaryWatchers) > 0 && tsoPrimaryWatchers[0] != nil
 	noProxyRequestsTimer := time.NewTimer(tsoProxyStreamIdleTimeout)
 	for {
@@ -179,7 +188,9 @@ func (s *TSODispatcher) processRequests(forwardStream stream, requests []Request
 		return err
 	}
 	s.tsoProxyHandleDuration.Observe(time.Since(start).Seconds())
-	s.tsoProxyBatchSize.Observe(float64(count))
+	if s.tsoProxyBatchSize != nil {
+		s.tsoProxyBatchSize.Observe(float64(count))
+	}
 	// Split the response
 	ts := resp.GetTimestamp()
 	physical, logical := ts.GetPhysical(), ts.GetLogical()
@@ -202,6 +213,40 @@ func (*TSODispatcher) finishRequest(requests []Request, physical, firstLogical i
 	return nil
 }
 
+// clearPendingRequests clears all pending requests in the queue to prevent goroutine leakage.
+// This method should be called when an error occurs to ensure that all waiting goroutines
+// are notified and can exit gracefully.
+func (s *TSODispatcher) clearPendingRequests(tsoQueue *tsoRequestProxyQueue, forwardedHost string, err error) {
+	// Delete the queue from the dispatcher to prevent new requests from being accepted
+	s.dispatchChs.Delete(forwardedHost)
+	defer tsoQueue.cancel(err)
+
+	// Clear all pending requests in the queue
+	// We don't close the channel here to avoid panic in other goroutines that might try to send to it
+	// Instead, we drain the channel
+	for {
+		select {
+		case <-tsoQueue.requestCh:
+			// We can't directly notify the request of the error since the Request interface
+			// doesn't have an onError method. The goroutines waiting for these requests
+			// will eventually be notified through the context cancellation.
+		default:
+			// No more requests in the channel
+			return
+		}
+	}
+}
+
+// Stop the TSODispatcher and clears all pending requests
+// Used for test only.
+func (s *TSODispatcher) Stop() {
+	s.dispatchChs.Range(func(_, value any) bool {
+		tsoQueue := value.(*tsoRequestProxyQueue)
+		tsoQueue.cancel(errors.New("TSODispatcher stopped"))
+		return true
+	})
+}
+
 // TSDeadline is used to watch the deadline of each tso request.
 type TSDeadline struct {
 	timer  *time.Timer


Test patch:
diff --git a/pkg/utils/tsoutil/tso_dispatcher_test.go b/pkg/utils/tsoutil/tso_dispatcher_test.go
new file mode 100644
index 00000000000..91890b10a71
--- /dev/null
+++ b/pkg/utils/tsoutil/tso_dispatcher_test.go
@@ -0,0 +1,194 @@
+// Copyright 2025 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tsoutil
+
+import (
+	"context"
+	"errors"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/suite"
+	"go.uber.org/goleak"
+	"google.golang.org/grpc"
+
+	"github.com/pingcap/kvproto/pkg/pdpb"
+
+	"github.com/tikv/pd/pkg/utils/testutil"
+)
+
+// mockStream is a mock stream used to test TSO request processing
+type mockStream struct {
+	sync.Mutex
+	// delay is used to simulate a long-running processRequest call
+	delay     time.Duration
+	succCount int // Counter for successfully processed requests
+}
+
+// process simulates the process of handling TSO requests
+func (m *mockStream) process(clusterID uint64, _ uint32, _ uint32, _ uint32) (tsoResp, error) {
+	m.Lock()
+	defer m.Unlock()
+
+	m.succCount++
+	// Simulate an error after every 10 successful requests
+	if m.succCount == 10 {
+		// Simulate stream error
+		err := errors.New("simulated stream error")
+		m.succCount = 0
+		return nil, err
+	}
+
+	// If delay is set, sleep for a while to simulate processing time
+	if m.delay > 0 {
+		time.Sleep(m.delay)
+	}
+
+	// Simulate a successful response
+	return &pdpb.TsoResponse{
+		Header: &pdpb.ResponseHeader{
+			ClusterId: clusterID,
+		},
+		Timestamp: &pdpb.Timestamp{
+			Physical: time.Now().UnixMilli(), // Use current time as physical timestamp
+			Logical:  1,                      // Logical timestamp is fixed at 1
+		},
+	}, nil
+}
+
+type mockProtoFactory struct {
+	stream *mockStream
+}
+
+func (m *mockProtoFactory) createForwardStream(_ context.Context, _ *grpc.ClientConn) (stream, context.CancelFunc, error) {
+	if m.stream == nil {
+		return nil, nil, errors.New("no stream available")
+	}
+	return m.stream, func() {}, nil
+}
+
+type mockRequest struct {
+	forwardedHost string
+	clientConn    *grpc.ClientConn
+	count         uint32
+	doneCh        chan struct{}
+	err           error
+}
+
+func (m *mockRequest) getForwardedHost() string {
+	return m.forwardedHost
+}
+
+func (m *mockRequest) getClientConn() *grpc.ClientConn {
+	return m.clientConn
+}
+
+func (m *mockRequest) getCount() uint32 {
+	return m.count
+}
+
+func (m *mockRequest) process(forwardStream stream, count uint32) (tsoResp, error) {
+	if m.err != nil {
+		return nil, m.err
+	}
+	return forwardStream.process(0, count, 0, 0)
+}
+
+func (m *mockRequest) postProcess(countSum int64, _ int64, _ int64) (int64, error) {
+	close(m.doneCh)
+	return countSum, m.err
+}
+
+type tsoDispatcherTestSuite struct {
+	suite.Suite
+	dispatcher *TSODispatcher
+}
+
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
+func TestTSODispatcherTestSuite(t *testing.T) {
+	suite.Run(t, new(tsoDispatcherTestSuite))
+}
+
+func (suite *tsoDispatcherTestSuite) SetupTest() {
+	tsoProxyHandleDuration := prometheus.NewHistogram(prometheus.HistogramOpts{
+		Name: "test_tso_proxy_handle_duration_seconds",
+		Help: "Histogram of TSO proxy handle duration",
+	})
+	suite.dispatcher = NewTSODispatcher(tsoProxyHandleDuration, nil)
+}
+
+func (suite *tsoDispatcherTestSuite) TearDownTest() {
+	suite.dispatcher.Stop()
+	runtime.GC() // Force garbage collection to clean up any lingering goroutines
+}
+
+func (suite *tsoDispatcherTestSuite) TestGoroutineLeakOnStreamError() {
+	re := suite.Require()
+
+	// Create a mock tso stream
+	mockStream := &mockStream{
+		delay: 10 * time.Millisecond,
+	}
+	protoFactory := &mockProtoFactory{
+		stream: mockStream,
+	}
+
+	ctx := context.Background()
+	forwardedHost := "test-host"
+	var reqDispatchCount atomic.Int64
+	var wg sync.WaitGroup
+	// mock 11000 goroutines as tso proxy client
+	for range 11000 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			dipatchReqDoneCh := make(chan struct{})
+			for range 20 {
+				req := &mockRequest{
+					forwardedHost: forwardedHost,
+					clientConn:    nil,
+					count:         1,
+					doneCh:        make(chan struct{}),
+				}
+
+				go func() {
+					// if there has goroutine leak, it will block here
+					suite.dispatcher.DispatchRequest(ctx, req, protoFactory)
+					dipatchReqDoneCh <- struct{}{}
+				}()
+
+				select {
+				case <-dipatchReqDoneCh:
+					reqDispatchCount.Add(1)
+				case <-time.After(20 * time.Second):
+					// If the request takes too long, we assume there is a goroutine leak
+				}
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	re.Equal(int64(11000*20), reqDispatchCount.Load(), "The number of dispatched requests should match the total requests sent")
+
+	suite.dispatcher.Stop()
+}