From b554ebea24d0d8dba0e948239354bd90f36019c1 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Tue, 23 Jun 2026 13:41:53 +0800 Subject: [PATCH 1/6] client, pkg: handle canceled contexts in select paths Signed-off-by: Ryan Leung --- client/clients/router/client.go | 1 + client/pkg/batch/batch_controller.go | 55 +++++++++++++++++++ client/pkg/batch/batch_controller_test.go | 41 ++++++++++++++ client/pkg/deadline/watcher.go | 11 ++-- client/resource_manager_client.go | 4 ++ client/servicediscovery/service_discovery.go | 10 ++++ pkg/ratelimit/concurrency_limiter.go | 18 +++++- pkg/ratelimit/concurrency_limiter_test.go | 14 +++++ pkg/ratelimit/runner.go | 3 + pkg/ratelimit/runner_test.go | 21 +++++++ pkg/utils/syncutil/ordered_single_flight.go | 21 +++++-- .../syncutil/ordered_single_flight_test.go | 18 ++++++ 12 files changed, 208 insertions(+), 9 deletions(-) diff --git a/client/clients/router/client.go b/client/clients/router/client.go index 6b4fd946a0e..b5f63198ad6 100644 --- a/client/clients/router/client.go +++ b/client/clients/router/client.go @@ -621,6 +621,7 @@ batchLoop: // Check if the dispatcher is canceled or the timeout timer is triggered. select { case <-ctx.Done(): + c.batchController.FinishCollectedRequests(requestFinisher(nil), ctx.Err()) return case <-timeoutTimer.C: log.Error("[router] router stream connection is not ready until timeout, abort the batch") diff --git a/client/pkg/batch/batch_controller.go b/client/pkg/batch/batch_controller.go index 51205286580..6d9b4c11173 100644 --- a/client/pkg/batch/batch_controller.go +++ b/client/pkg/batch/batch_controller.go @@ -82,6 +82,9 @@ func (bc *Controller[T]) FetchPendingRequests(ctx context.Context, requestCh <-c // TODO: `bc.collectedRequestCount` should never be non-empty here. Consider do assertion here. bc.collectedRequestCount = 0 for { + if errRet = ctx.Err(); errRet != nil { + return errRet + } // If the batch size reaches the maxBatchSize limit but the token haven't arrived yet, don't receive more // requests, and return when token is ready. if bc.collectedRequestCount >= bc.maxBatchSize && !tokenAcquired { @@ -92,6 +95,10 @@ func (bc *Controller[T]) FetchPendingRequests(ctx context.Context, requestCh <-c case <-ctx.Done(): return ctx.Err() case <-tokenCh: + tokenAcquired = true + if errRet = ctx.Err(); errRet != nil { + return errRet + } return nil } } @@ -103,22 +110,34 @@ func (bc *Controller[T]) FetchPendingRequests(ctx context.Context, requestCh <-c case req := <-requestCh: // Start to batch when the first request arrives. bc.pushRequest(req) + if errRet = ctx.Err(); errRet != nil { + return errRet + } // A request arrives but the token is not ready yet. Continue waiting, and also allowing collecting the next // request if it arrives. continue case <-tokenCh: tokenAcquired = true + if errRet = ctx.Err(); errRet != nil { + return errRet + } } } // After the token is ready or it's working without token, // wait for the first request to arrive. if bc.collectedRequestCount == 0 { + if errRet = ctx.Err(); errRet != nil { + return errRet + } select { case <-ctx.Done(): return ctx.Err() case firstRequest := <-requestCh: bc.pushRequest(firstRequest) + if errRet = ctx.Err(); errRet != nil { + return errRet + } } } @@ -131,9 +150,15 @@ func (bc *Controller[T]) FetchPendingRequests(ctx context.Context, requestCh <-c // This loop is for trying best to collect more requests, so we use `bc.maxBatchSize` here. fetchPendingRequestsLoop: for bc.collectedRequestCount < bc.maxBatchSize { + if errRet = ctx.Err(); errRet != nil { + return errRet + } select { case req := <-requestCh: bc.pushRequest(req) + if errRet = ctx.Err(); errRet != nil { + return errRet + } case <-ctx.Done(): return ctx.Err() default: @@ -153,12 +178,21 @@ fetchPendingRequestsLoop: after := time.NewTimer(maxBatchWaitInterval) defer after.Stop() for bc.collectedRequestCount < bc.bestBatchSize { + if errRet = ctx.Err(); errRet != nil { + return errRet + } select { case req := <-requestCh: bc.pushRequest(req) + if errRet = ctx.Err(); errRet != nil { + return errRet + } case <-ctx.Done(): return ctx.Err() case <-after.C: + if errRet = ctx.Err(); errRet != nil { + return errRet + } return nil } } @@ -168,9 +202,15 @@ fetchPendingRequestsLoop: // of `bc.bestBatchSize` because trying best to fetch more requests is necessary so that // we can adjust the `bc.bestBatchSize` dynamically later. for bc.collectedRequestCount < bc.maxBatchSize { + if errRet = ctx.Err(); errRet != nil { + return errRet + } select { case req := <-requestCh: bc.pushRequest(req) + if errRet = ctx.Err(); errRet != nil { + return errRet + } case <-ctx.Done(): return ctx.Err() default: @@ -185,12 +225,21 @@ fetchPendingRequestsLoop: func (bc *Controller[T]) FetchRequestsWithTimer(ctx context.Context, requestCh <-chan T, timer *time.Timer) error { batchingLoop: for bc.collectedRequestCount < bc.maxBatchSize { + if err := ctx.Err(); err != nil { + return err + } select { case <-ctx.Done(): return ctx.Err() case req := <-requestCh: bc.pushRequest(req) + if err := ctx.Err(); err != nil { + return err + } case <-timer.C: + if err := ctx.Err(); err != nil { + return err + } break batchingLoop } } @@ -198,11 +247,17 @@ batchingLoop: // Try to collect more requests in non-blocking way. nonWaitingBatchLoop: for bc.collectedRequestCount < bc.maxBatchSize { + if err := ctx.Err(); err != nil { + return err + } select { case <-ctx.Done(): return ctx.Err() case req := <-requestCh: bc.pushRequest(req) + if err := ctx.Err(); err != nil { + return err + } default: break nonWaitingBatchLoop } diff --git a/client/pkg/batch/batch_controller_test.go b/client/pkg/batch/batch_controller_test.go index 942a608af72..e2b14db862d 100644 --- a/client/pkg/batch/batch_controller_test.go +++ b/client/pkg/batch/batch_controller_test.go @@ -17,6 +17,7 @@ package batch import ( "context" "testing" + "time" "github.com/stretchr/testify/require" "go.uber.org/goleak" @@ -145,3 +146,43 @@ func TestFetchPendingRequests(t *testing.T) { // Drain the requestCh. <-requestCh } + +func TestFetchPendingRequestsWithCanceledContext(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + finishedCount := 0 + bc := NewController[int](testMaxBatchSize, func(_ int, _ int, err error) { + re.ErrorIs(err, context.Canceled) + finishedCount++ + }, nil) + requestCh := make(chan int, 1) + requestCh <- 1 + tokenCh := make(chan struct{}, 1) + tokenCh <- struct{}{} + + err := bc.FetchPendingRequests(ctx, requestCh, tokenCh, 0) + re.ErrorIs(err, context.Canceled) + re.Zero(bc.collectedRequestCount) + re.Zero(finishedCount) + re.Len(requestCh, 1) + re.Len(tokenCh, 1) +} + +func TestFetchRequestsWithTimerWithCanceledContext(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + bc := NewController[int](testMaxBatchSize, nil, nil) + requestCh := make(chan int, 1) + requestCh <- 1 + timer := time.NewTimer(time.Hour) + defer timer.Stop() + + err := bc.FetchRequestsWithTimer(ctx, requestCh, timer) + re.ErrorIs(err, context.Canceled) + re.Zero(bc.collectedRequestCount) + re.Len(requestCh, 1) +} diff --git a/client/pkg/deadline/watcher.go b/client/pkg/deadline/watcher.go index b40857edbfd..e8647776fae 100644 --- a/client/pkg/deadline/watcher.go +++ b/client/pkg/deadline/watcher.go @@ -83,12 +83,11 @@ func (w *Watcher) Start( cancel context.CancelFunc, ) chan struct{} { // Check if the watcher is already canceled. - select { - case <-w.ctx.Done(): + if w.ctx.Err() != nil { return nil - case <-ctx.Done(): + } + if ctx.Err() != nil { return nil - default: } // Initialize the deadline. timer := timerutil.GlobalTimerPool.Get(timeout) @@ -106,6 +105,10 @@ func (w *Watcher) Start( timerutil.GlobalTimerPool.Put(timer) return nil case w.Ch <- d: + if w.ctx.Err() != nil || ctx.Err() != nil { + close(d.done) + return nil + } return d.done } } diff --git a/client/resource_manager_client.go b/client/resource_manager_client.go index 30ed6cf62c3..d31c2f2d956 100644 --- a/client/resource_manager_client.go +++ b/client/resource_manager_client.go @@ -359,6 +359,10 @@ func (c *innerClient) handleResourceTokenDispatcher(dispatcherCtx context.Contex return case firstRequest = <-tbc.tokenRequestCh: } + if err := dispatcherCtx.Err(); err != nil { + firstRequest.done <- errors.WithStack(err) + return + } // Try to get a stream connection. stream, streamCtx = connection.stream, connection.ctx select { diff --git a/client/servicediscovery/service_discovery.go b/client/servicediscovery/service_discovery.go index 5716c24125c..b69a87d5cd6 100644 --- a/client/servicediscovery/service_discovery.go +++ b/client/servicediscovery/service_discovery.go @@ -930,6 +930,11 @@ func (c *serviceDiscovery) getClusterInfo(ctx context.Context, url string, timeo } start := time.Now() defer func() { metrics.InternalCmdDurationGetClusterInfo.Observe(time.Since(start).Seconds()) }() + if err := ctx.Err(); err != nil { + metrics.InternalCmdFailedDurationGetClusterInfo.Observe(time.Since(start).Seconds()) + attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String()) + return nil, errs.ErrClientGetClusterInfo.Wrap(attachErr).GenWithStackByCause() + } key := "GetClusterInfo-" + url r := c.flight.DoChan(key, func() (any, error) { return pdpb.NewPDClient(cc).GetClusterInfo(ctx, &pdpb.GetClusterInfoRequest{}) @@ -966,6 +971,11 @@ func (c *serviceDiscovery) getMembers(ctx context.Context, url string, timeout t } start := time.Now() defer func() { metrics.InternalCmdDurationGetMembers.Observe(time.Since(start).Seconds()) }() + if err := ctx.Err(); err != nil { + metrics.InternalCmdFailedDurationGetMembers.Observe(time.Since(start).Seconds()) + attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String()) + return nil, errs.ErrClientGetMember.Wrap(attachErr).GenWithStackByCause() + } key := "GetMembers-" + url r := c.flight.DoChan(key, func() (any, error) { return pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{}) diff --git a/pkg/ratelimit/concurrency_limiter.go b/pkg/ratelimit/concurrency_limiter.go index e5379bc48cc..9060da2ed89 100644 --- a/pkg/ratelimit/concurrency_limiter.go +++ b/pkg/ratelimit/concurrency_limiter.go @@ -108,7 +108,15 @@ func (l *ConcurrencyLimiter) GetWaitingTasksNum() uint64 { // AcquireToken acquires a token from the limiter. which will block until a token is available or ctx is done, like Timeout. func (l *ConcurrencyLimiter) AcquireToken(ctx context.Context) (*TaskToken, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + l.mu.Lock() + if err := ctx.Err(); err != nil { + l.mu.Unlock() + return nil, err + } if l.current >= l.limit { l.waiting++ l.mu.Unlock() @@ -121,9 +129,17 @@ func (l *ConcurrencyLimiter) AcquireToken(ctx context.Context) (*TaskToken, erro return nil, ctx.Err() case token := <-l.queue: l.mu.Lock() + l.waiting-- + if err := ctx.Err(); err != nil { + l.mu.Unlock() + select { + case l.queue <- token: + default: + } + return nil, err + } token.released = false l.current++ - l.waiting-- l.mu.Unlock() return token, nil } diff --git a/pkg/ratelimit/concurrency_limiter_test.go b/pkg/ratelimit/concurrency_limiter_test.go index 16c31ebe2b9..e47a9d2e5f9 100644 --- a/pkg/ratelimit/concurrency_limiter_test.go +++ b/pkg/ratelimit/concurrency_limiter_test.go @@ -102,6 +102,20 @@ func TestConcurrencyLimiter2(t *testing.T) { require.Equal(t, uint64(1), limiter.GetRunningTasksNum(), "Expected running tasks to be 1") } +func TestConcurrencyLimiterAcquireWithCanceledContext(t *testing.T) { + re := require.New(t) + limiter := NewConcurrencyLimiter(1) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + token, err := limiter.AcquireToken(ctx) + re.ErrorIs(err, context.Canceled) + re.Nil(token) + re.Zero(limiter.GetRunningTasksNum()) + re.Zero(limiter.GetWaitingTasksNum()) +} + func TestConcurrencyLimiterAcquire(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() diff --git a/pkg/ratelimit/runner.go b/pkg/ratelimit/runner.go index d273dd19f2d..d373ba8251f 100644 --- a/pkg/ratelimit/runner.go +++ b/pkg/ratelimit/runner.go @@ -148,6 +148,9 @@ func (cr *ConcurrentRunner) Start(ctx context.Context) { func (cr *ConcurrentRunner) run(ctx context.Context, task *Task, token *TaskToken) { start := time.Now() + if token != nil { + defer cr.limiter.ReleaseToken(token) + } select { case <-ctx.Done(): return diff --git a/pkg/ratelimit/runner_test.go b/pkg/ratelimit/runner_test.go index 25ace41c381..fab8d469b9e 100644 --- a/pkg/ratelimit/runner_test.go +++ b/pkg/ratelimit/runner_test.go @@ -117,4 +117,25 @@ func TestConcurrentRunner(t *testing.T) { lastSubmitted := runner.pendingTasks[len(runner.pendingTasks)-1].submittedAt require.Greater(t, updatedSubmitted, lastSubmitted) }) + + t.Run("ReleaseTokenWhenContextCanceled", func(t *testing.T) { + limiter := NewConcurrencyLimiter(1) + runner := NewConcurrentRunner("test", limiter, time.Minute) + token, err := limiter.AcquireToken(context.Background()) + require.NoError(t, err) + require.Equal(t, uint64(1), limiter.GetRunningTasksNum()) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + called := false + runner.run(ctx, &Task{ + name: "test4", + f: func(context.Context) { + called = true + }, + }, token) + + require.False(t, called) + require.Zero(t, limiter.GetRunningTasksNum()) + }) } diff --git a/pkg/utils/syncutil/ordered_single_flight.go b/pkg/utils/syncutil/ordered_single_flight.go index a4169cbb048..6a08f9203bc 100644 --- a/pkg/utils/syncutil/ordered_single_flight.go +++ b/pkg/utils/syncutil/ordered_single_flight.go @@ -112,6 +112,11 @@ func NewOrderedSingleFlight[TResult any]() *OrderedSingleFlight[TResult] { // Do tries to execute the function `f`, or if possible, wait and reuse the result of an execution triggered by another // goroutine. See comments of OrderedSingleFlight for details. func (s *OrderedSingleFlight[TResult]) Do(ctx context.Context, f func(context.Context) (TResult, error)) (TResult, error) { + var zero TResult + if err := ctx.Err(); err != nil { + return zero, err + } + var currentTask *task[TResult] s.mu.Lock() @@ -122,11 +127,15 @@ func (s *OrderedSingleFlight[TResult]) Do(ctx context.Context, f func(context.Co currentTask.refState.Add(1) s.mu.Unlock() + if err := ctx.Err(); err != nil { + currentTask.cancelOne() + return zero, err + } + select { case <-ctx.Done(): currentTask.cancelOne() - var res TResult - return res, ctx.Err() + return zero, ctx.Err() case <-currentTask.finishCh: // Executed by other goroutines. Reuse the result. return currentTask.result, currentTask.err @@ -146,6 +155,11 @@ func (s *OrderedSingleFlight[TResult]) Do(ctx context.Context, f func(context.Co return currentTask.result, currentTask.err default: } + if err := ctx.Err(); err != nil { + currentTask.cancelOne() + s.tokenCh <- struct{}{} + return zero, err + } // The following code executes exactly once for each instance of the `task`. @@ -174,8 +188,7 @@ func (s *OrderedSingleFlight[TResult]) Do(ctx context.Context, f func(context.Co select { case <-ctx.Done(): currentTask.cancelOne() - var res TResult - return res, ctx.Err() + return zero, ctx.Err() case <-currentTask.finishCh: return currentTask.result, currentTask.err } diff --git a/pkg/utils/syncutil/ordered_single_flight_test.go b/pkg/utils/syncutil/ordered_single_flight_test.go index dfc07d994e5..5508be1efed 100644 --- a/pkg/utils/syncutil/ordered_single_flight_test.go +++ b/pkg/utils/syncutil/ordered_single_flight_test.go @@ -123,6 +123,24 @@ func TestOrderedSingleFlight(t *testing.T) { re.Equal(200, res) } +func TestOrderedSingleFlightCanceledContextDoesNotExecute(t *testing.T) { + re := require.New(t) + + s := NewOrderedSingleFlight[int]() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + called := atomic.Bool{} + res, err := s.Do(ctx, func(context.Context) (int, error) { + called.Store(true) + return 1, nil + }) + re.ErrorIs(err, context.Canceled) + re.Zero(res) + re.False(called.Load()) + re.Equal(int64(0), s.ExecCount()) +} + func TestOrderedSingleFlightCancellation(t *testing.T) { re := require.New(t) From 72663228dcdda3f396316ca8bca87def3181b86f Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Tue, 23 Jun 2026 14:59:56 +0800 Subject: [PATCH 2/6] client: avoid work for canceled internal requests Signed-off-by: Ryan Leung --- client/resource_manager_client.go | 18 ++++++++- client/resource_manager_client_test.go | 21 +++++++++++ client/servicediscovery/service_discovery.go | 20 +++++----- .../service_discovery_test.go | 37 +++++++++++++++++++ 4 files changed, 85 insertions(+), 11 deletions(-) diff --git a/client/resource_manager_client.go b/client/resource_manager_client.go index d31c2f2d956..a0f129fab0d 100644 --- a/client/resource_manager_client.go +++ b/client/resource_manager_client.go @@ -258,13 +258,25 @@ func (c *client) LoadResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, // AcquireTokenBuckets implements the ResourceManagerClient interface. func (c *client) AcquireTokenBuckets(ctx context.Context, request *rmpb.TokenBucketsRequest) ([]*rmpb.TokenBucketResponse, error) { + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + if err := c.inner.ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } req := &tokenRequest{ done: make(chan error, 1), requestCtx: ctx, clientCtx: c.inner.ctx, Request: request, } - c.inner.tokenDispatcher.tokenBatchController.tokenRequestCh <- req + select { + case c.inner.tokenDispatcher.tokenBatchController.tokenRequestCh <- req: + case <-ctx.Done(): + return nil, errors.WithStack(ctx.Err()) + case <-c.inner.ctx.Done(): + return nil, errors.WithStack(c.inner.ctx.Err()) + } grantedTokens, err := req.wait() if err != nil { return nil, err @@ -363,6 +375,10 @@ func (c *innerClient) handleResourceTokenDispatcher(dispatcherCtx context.Contex firstRequest.done <- errors.WithStack(err) return } + if err := firstRequest.requestCtx.Err(); err != nil { + firstRequest.done <- errors.WithStack(err) + continue + } // Try to get a stream connection. stream, streamCtx = connection.stream, connection.ctx select { diff --git a/client/resource_manager_client_test.go b/client/resource_manager_client_test.go index 84674dcca4b..8c0cb75aae8 100644 --- a/client/resource_manager_client_test.go +++ b/client/resource_manager_client_test.go @@ -373,3 +373,24 @@ func TestTryResourceManagerConnectUsesRMForTokenAndFallbackToPD(t *testing.T) { require.EqualValues(t, 1, pdServer.tokenCount.Load()) }) } + +func TestAcquireTokenBucketsWithCanceledContextDoesNotEnqueue(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + tokenRequestCh := make(chan *tokenRequest, 1) + tokenRequestCh <- &tokenRequest{done: make(chan error, 1)} + cli := &client{ + inner: &innerClient{ + ctx: ctx, + tokenDispatcher: &tokenDispatcher{ + tokenBatchController: newTokenBatchController(tokenRequestCh), + }, + }, + } + + requestCtx, requestCancel := context.WithCancel(context.Background()) + requestCancel() + _, err := cli.AcquireTokenBuckets(requestCtx, &rmpb.TokenBucketsRequest{}) + require.ErrorIs(t, err, context.Canceled) + require.Len(t, tokenRequestCh, 1) +} diff --git a/client/servicediscovery/service_discovery.go b/client/servicediscovery/service_discovery.go index b69a87d5cd6..fd004cd66e7 100644 --- a/client/servicediscovery/service_discovery.go +++ b/client/servicediscovery/service_discovery.go @@ -924,17 +924,17 @@ func (c *serviceDiscovery) updateMember() error { func (c *serviceDiscovery) getClusterInfo(ctx context.Context, url string, timeout time.Duration) (*pdpb.GetClusterInfoResponse, error) { ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - cc, err := c.GetOrCreateGRPCConn(url) - if err != nil { - return nil, err - } start := time.Now() defer func() { metrics.InternalCmdDurationGetClusterInfo.Observe(time.Since(start).Seconds()) }() if err := ctx.Err(); err != nil { metrics.InternalCmdFailedDurationGetClusterInfo.Observe(time.Since(start).Seconds()) - attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String()) + attachErr := errors.Errorf("error:%s target:%s", err, url) return nil, errs.ErrClientGetClusterInfo.Wrap(attachErr).GenWithStackByCause() } + cc, err := c.GetOrCreateGRPCConn(url) + if err != nil { + return nil, err + } key := "GetClusterInfo-" + url r := c.flight.DoChan(key, func() (any, error) { return pdpb.NewPDClient(cc).GetClusterInfo(ctx, &pdpb.GetClusterInfoRequest{}) @@ -965,17 +965,17 @@ func (c *serviceDiscovery) getClusterInfo(ctx context.Context, url string, timeo func (c *serviceDiscovery) getMembers(ctx context.Context, url string, timeout time.Duration) (*pdpb.GetMembersResponse, error) { ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - cc, err := c.GetOrCreateGRPCConn(url) - if err != nil { - return nil, err - } start := time.Now() defer func() { metrics.InternalCmdDurationGetMembers.Observe(time.Since(start).Seconds()) }() if err := ctx.Err(); err != nil { metrics.InternalCmdFailedDurationGetMembers.Observe(time.Since(start).Seconds()) - attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String()) + attachErr := errors.Errorf("error:%s target:%s", err, url) return nil, errs.ErrClientGetMember.Wrap(attachErr).GenWithStackByCause() } + cc, err := c.GetOrCreateGRPCConn(url) + if err != nil { + return nil, err + } key := "GetMembers-" + url r := c.flight.DoChan(key, func() (any, error) { return pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{}) diff --git a/client/servicediscovery/service_discovery_test.go b/client/servicediscovery/service_discovery_test.go index 295c4592cdf..96dbb5c62e8 100644 --- a/client/servicediscovery/service_discovery_test.go +++ b/client/servicediscovery/service_discovery_test.go @@ -452,3 +452,40 @@ func TestGRPCDialOption(t *testing.T) { re.Error(err) re.Greater(time.Since(start), 500*time.Millisecond) } + +func TestCanceledGetClusterInfoAndMembersDoesNotCreateConn(t *testing.T) { + const url = "http://127.0.0.1:1" + newCanceledCtx := func() context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + } + + t.Run("get-cluster-info", func(t *testing.T) { + re := require.New(t) + cli := &serviceDiscovery{ + ctx: context.Background(), + option: opt.NewOption(), + } + resp, err := cli.getClusterInfo(newCanceledCtx(), url, time.Second) + re.Error(err) + re.Contains(err.Error(), context.Canceled.Error()) + re.Nil(resp) + _, ok := cli.clientConns.Load(url) + re.False(ok) + }) + + t.Run("get-members", func(t *testing.T) { + re := require.New(t) + cli := &serviceDiscovery{ + ctx: context.Background(), + option: opt.NewOption(), + } + resp, err := cli.getMembers(newCanceledCtx(), url, time.Second) + re.Error(err) + re.Contains(err.Error(), context.Canceled.Error()) + re.Nil(resp) + _, ok := cli.clientConns.Load(url) + re.False(ok) + }) +} From 5d363b636ee4f62cb0922b763d319b2fe2527ad4 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Tue, 23 Jun 2026 15:37:13 +0800 Subject: [PATCH 3/6] client: close request cancellation gaps Signed-off-by: Ryan Leung --- client/resource_manager_client.go | 4 ++- client/resource_manager_client_test.go | 32 ++++++++++++++++++++ client/servicediscovery/service_discovery.go | 10 ++++++ 3 files changed, 45 insertions(+), 1 deletion(-) diff --git a/client/resource_manager_client.go b/client/resource_manager_client.go index a0f129fab0d..5361fdec439 100644 --- a/client/resource_manager_client.go +++ b/client/resource_manager_client.go @@ -434,7 +434,9 @@ func (c *innerClient) processTokenRequests(stream rmpb.ResourceManager_AcquireTo return err } if resp.GetError() != nil { - return errors.Errorf("[resource_manager] %s", resp.GetError().Message) + err := errors.Errorf("[resource_manager] %s", resp.GetError().Message) + t.done <- err + return err } t.TokenBuckets = resp.GetResponses() t.done <- nil diff --git a/client/resource_manager_client_test.go b/client/resource_manager_client_test.go index 8c0cb75aae8..855687df411 100644 --- a/client/resource_manager_client_test.go +++ b/client/resource_manager_client_test.go @@ -131,6 +131,7 @@ type testRMServer struct { modifyCount atomic.Int32 deleteCount atomic.Int32 tokenCount atomic.Int32 + tokenError atomic.Bool } func (s *testRMServer) ListResourceGroups(context.Context, *rmpb.ListResourceGroupsRequest) (*rmpb.ListResourceGroupsResponse, error) { @@ -176,6 +177,14 @@ func (s *testRMServer) AcquireTokenBuckets(stream rmpb.ResourceManager_AcquireTo return err } s.tokenCount.Add(1) + if s.tokenError.Load() { + if err := stream.Send(&rmpb.TokenBucketsResponse{ + Error: &rmpb.Error{Message: "token error"}, + }); err != nil { + return err + } + continue + } resp := &rmpb.TokenBucketsResponse{ Responses: make([]*rmpb.TokenBucketResponse, 0, len(req.GetRequests())), } @@ -394,3 +403,26 @@ func TestAcquireTokenBucketsWithCanceledContextDoesNotEnqueue(t *testing.T) { require.ErrorIs(t, err, context.Canceled) require.Len(t, tokenRequestCh, 1) } + +func TestAcquireTokenBucketsReturnsServerError(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + pdAddr, pdServer, pdCleanup := startTestRMServer(t, "pd") + t.Cleanup(pdCleanup) + pdServer.tokenError.Store(true) + + inner := newInnerClientForRMRouteTest(t, ctx, pdAddr) + inner.createTokenDispatcher() + t.Cleanup(func() { + inner.tokenDispatcher.dispatcherCancel() + inner.wg.Wait() + }) + cli := &client{inner: inner} + + _, err := cli.AcquireTokenBuckets(ctx, &rmpb.TokenBucketsRequest{ + Requests: []*rmpb.TokenBucketRequest{{ResourceGroupName: "test-group"}}, + }) + require.ErrorContains(t, err, "token error") + require.EqualValues(t, 1, pdServer.tokenCount.Load()) +} diff --git a/client/servicediscovery/service_discovery.go b/client/servicediscovery/service_discovery.go index fd004cd66e7..9b4a204a93e 100644 --- a/client/servicediscovery/service_discovery.go +++ b/client/servicediscovery/service_discovery.go @@ -935,6 +935,11 @@ func (c *serviceDiscovery) getClusterInfo(ctx context.Context, url string, timeo if err != nil { return nil, err } + if err := ctx.Err(); err != nil { + metrics.InternalCmdFailedDurationGetClusterInfo.Observe(time.Since(start).Seconds()) + attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String()) + return nil, errs.ErrClientGetClusterInfo.Wrap(attachErr).GenWithStackByCause() + } key := "GetClusterInfo-" + url r := c.flight.DoChan(key, func() (any, error) { return pdpb.NewPDClient(cc).GetClusterInfo(ctx, &pdpb.GetClusterInfoRequest{}) @@ -976,6 +981,11 @@ func (c *serviceDiscovery) getMembers(ctx context.Context, url string, timeout t if err != nil { return nil, err } + if err := ctx.Err(); err != nil { + metrics.InternalCmdFailedDurationGetMembers.Observe(time.Since(start).Seconds()) + attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String()) + return nil, errs.ErrClientGetMember.Wrap(attachErr).GenWithStackByCause() + } key := "GetMembers-" + url r := c.flight.DoChan(key, func() (any, error) { return pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{}) From e01fcd25f6ca69cee7bc936f40b64f7599fa2968 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Tue, 23 Jun 2026 16:45:45 +0800 Subject: [PATCH 4/6] pkg: keep concurrency limiter tokens bounded Signed-off-by: Ryan Leung --- pkg/ratelimit/concurrency_limiter.go | 48 +++++++++++++---------- pkg/ratelimit/concurrency_limiter_test.go | 38 ++++++++++++++++++ 2 files changed, 65 insertions(+), 21 deletions(-) diff --git a/pkg/ratelimit/concurrency_limiter.go b/pkg/ratelimit/concurrency_limiter.go index 9060da2ed89..a1b92059272 100644 --- a/pkg/ratelimit/concurrency_limiter.go +++ b/pkg/ratelimit/concurrency_limiter.go @@ -29,12 +29,12 @@ type ConcurrencyLimiter struct { // statistic maxLimit uint64 - queue chan *TaskToken + queue chan struct{} // wake-up signals for waiting tasks } // NewConcurrencyLimiter creates a new ConcurrencyLimiter. func NewConcurrencyLimiter(limit uint64) *ConcurrencyLimiter { - return &ConcurrencyLimiter{limit: limit, queue: make(chan *TaskToken, limit)} + return &ConcurrencyLimiter{limit: limit, queue: make(chan struct{}, limit)} } const unlimit = uint64(0) @@ -117,50 +117,56 @@ func (l *ConcurrencyLimiter) AcquireToken(ctx context.Context) (*TaskToken, erro l.mu.Unlock() return nil, err } - if l.current >= l.limit { - l.waiting++ + if l.limit == unlimit || l.current < l.limit { + l.current++ l.mu.Unlock() - // block the waiting task on the caller goroutine + return &TaskToken{}, nil + } + l.waiting++ + l.mu.Unlock() + + // Block the waiting task on the caller goroutine. + for { select { case <-ctx.Done(): l.mu.Lock() l.waiting-- l.mu.Unlock() return nil, ctx.Err() - case token := <-l.queue: + case <-l.queue: l.mu.Lock() - l.waiting-- if err := ctx.Err(); err != nil { + l.waiting-- l.mu.Unlock() - select { - case l.queue <- token: - default: - } return nil, err } - token.released = false - l.current++ + if l.limit == unlimit || l.current < l.limit { + l.waiting-- + l.current++ + l.mu.Unlock() + return &TaskToken{}, nil + } l.mu.Unlock() - return token, nil } } - l.current++ - token := &TaskToken{} - l.mu.Unlock() - return token, nil } // ReleaseToken releases the token. func (l *ConcurrencyLimiter) ReleaseToken(token *TaskToken) { l.mu.Lock() - defer l.mu.Unlock() if token.released { + l.mu.Unlock() return } token.released = true l.current-- - if len(l.queue) < int(l.limit) { - l.queue <- token + shouldWakeUp := l.waiting > 0 && (l.limit == unlimit || l.current < l.limit) + l.mu.Unlock() + if shouldWakeUp { + select { + case l.queue <- struct{}{}: + default: + } } } diff --git a/pkg/ratelimit/concurrency_limiter_test.go b/pkg/ratelimit/concurrency_limiter_test.go index e47a9d2e5f9..04aecb9619a 100644 --- a/pkg/ratelimit/concurrency_limiter_test.go +++ b/pkg/ratelimit/concurrency_limiter_test.go @@ -116,6 +116,44 @@ func TestConcurrencyLimiterAcquireWithCanceledContext(t *testing.T) { re.Zero(limiter.GetWaitingTasksNum()) } +func TestConcurrencyLimiterDoesNotReuseStaleQueuedToken(t *testing.T) { + re := require.New(t) + limiter := NewConcurrencyLimiter(1) + + token1, err := limiter.AcquireToken(context.Background()) + re.NoError(err) + limiter.ReleaseToken(token1) + + token2, err := limiter.AcquireToken(context.Background()) + re.NoError(err) + re.Equal(uint64(1), limiter.GetRunningTasksNum()) + + type acquireResult struct { + token *TaskToken + err error + } + ctx, cancel := context.WithCancel(context.Background()) + resultCh := make(chan acquireResult, 1) + go func() { + token, err := limiter.AcquireToken(ctx) + resultCh <- acquireResult{token: token, err: err} + }() + + select { + case result := <-resultCh: + re.Failf("unexpected token acquisition", "token: %v, err: %v", result.token, result.err) + case <-time.After(100 * time.Millisecond): + } + + cancel() + result := <-resultCh + re.ErrorIs(result.err, context.Canceled) + re.Nil(result.token) + re.Equal(uint64(1), limiter.GetRunningTasksNum()) + re.Zero(limiter.GetWaitingTasksNum()) + limiter.ReleaseToken(token2) +} + func TestConcurrencyLimiterAcquire(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() From bf704ab7bfdf6a54b5d1a3bbe6f352633a6f5674 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Tue, 23 Jun 2026 17:50:48 +0800 Subject: [PATCH 5/6] pkg: forward concurrency limiter wake-up on cancel Signed-off-by: Ryan Leung --- pkg/ratelimit/concurrency_limiter.go | 7 +++ pkg/ratelimit/concurrency_limiter_test.go | 56 +++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/pkg/ratelimit/concurrency_limiter.go b/pkg/ratelimit/concurrency_limiter.go index a1b92059272..eff5657830c 100644 --- a/pkg/ratelimit/concurrency_limiter.go +++ b/pkg/ratelimit/concurrency_limiter.go @@ -137,7 +137,14 @@ func (l *ConcurrencyLimiter) AcquireToken(ctx context.Context) (*TaskToken, erro l.mu.Lock() if err := ctx.Err(); err != nil { l.waiting-- + shouldWakeUp := l.waiting > 0 && (l.limit == unlimit || l.current < l.limit) l.mu.Unlock() + if shouldWakeUp { + select { + case l.queue <- struct{}{}: + default: + } + } return nil, err } if l.limit == unlimit || l.current < l.limit { diff --git a/pkg/ratelimit/concurrency_limiter_test.go b/pkg/ratelimit/concurrency_limiter_test.go index 04aecb9619a..af33221c9b2 100644 --- a/pkg/ratelimit/concurrency_limiter_test.go +++ b/pkg/ratelimit/concurrency_limiter_test.go @@ -154,6 +154,62 @@ func TestConcurrencyLimiterDoesNotReuseStaleQueuedToken(t *testing.T) { limiter.ReleaseToken(token2) } +func TestConcurrencyLimiterForwardsWakeUpFromCanceledWaiter(t *testing.T) { + re := require.New(t) + limiter := NewConcurrencyLimiter(1) + + token, err := limiter.AcquireToken(context.Background()) + re.NoError(err) + + type acquireResult struct { + token *TaskToken + err error + } + acquire := func(ctx context.Context, resultCh chan<- acquireResult) { + token, err := limiter.AcquireToken(ctx) + resultCh <- acquireResult{token: token, err: err} + } + + canceledCtx, cancelCanceledWaiter := context.WithCancel(context.Background()) + defer cancelCanceledWaiter() + canceledResultCh := make(chan acquireResult, 1) + go acquire(canceledCtx, canceledResultCh) + re.Eventually(func() bool { + return limiter.GetWaitingTasksNum() == 1 + }, time.Second, 10*time.Millisecond) + + survivingCtx, cancelSurvivingWaiter := context.WithCancel(context.Background()) + defer cancelSurvivingWaiter() + survivingResultCh := make(chan acquireResult, 1) + go acquire(survivingCtx, survivingResultCh) + re.Eventually(func() bool { + return limiter.GetWaitingTasksNum() == 2 + }, time.Second, 10*time.Millisecond) + + // Force the first waiter to consume the wake-up before it observes cancellation. + limiter.mu.Lock() + token.released = true + limiter.current-- + limiter.queue <- struct{}{} + cancelCanceledWaiter() + limiter.mu.Unlock() + + canceledResult := <-canceledResultCh + re.ErrorIs(canceledResult.err, context.Canceled) + re.Nil(canceledResult.token) + + select { + case survivingResult := <-survivingResultCh: + re.NoError(survivingResult.err) + re.NotNil(survivingResult.token) + limiter.ReleaseToken(survivingResult.token) + case <-time.After(time.Second): + re.Fail("surviving waiter did not acquire the released token") + } + re.Zero(limiter.GetRunningTasksNum()) + re.Zero(limiter.GetWaitingTasksNum()) +} + func TestConcurrencyLimiterAcquire(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() From a575cf8c412bbab8bdce71568329f9b7b0f15ccc Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 25 Jun 2026 15:29:53 +0800 Subject: [PATCH 6/6] client: finish token request on canceled stream Signed-off-by: Ryan Leung --- client/resource_manager_client.go | 1 + 1 file changed, 1 insertion(+) diff --git a/client/resource_manager_client.go b/client/resource_manager_client.go index 5361fdec439..04f221fef5e 100644 --- a/client/resource_manager_client.go +++ b/client/resource_manager_client.go @@ -405,6 +405,7 @@ func (c *innerClient) handleResourceTokenDispatcher(dispatcherCtx context.Contex } select { case <-streamCtx.Done(): + firstRequest.done <- errors.WithStack(streamCtx.Err()) connection.reset() log.Info("[resource_manager] token stream is canceled") continue