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..04f221fef5e 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 @@ -359,6 +371,14 @@ func (c *innerClient) handleResourceTokenDispatcher(dispatcherCtx context.Contex return case firstRequest = <-tbc.tokenRequestCh: } + if err := dispatcherCtx.Err(); err != nil { + 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 { @@ -385,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 @@ -414,7 +435,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 84674dcca4b..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())), } @@ -373,3 +382,47 @@ 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) +} + +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 5716c24125c..9b4a204a93e 100644 --- a/client/servicediscovery/service_discovery.go +++ b/client/servicediscovery/service_discovery.go @@ -924,12 +924,22 @@ 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() + 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", err, url) + return nil, errs.ErrClientGetClusterInfo.Wrap(attachErr).GenWithStackByCause() + } 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()) + 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{}) @@ -960,12 +970,22 @@ 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() + 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", err, url) + return nil, errs.ErrClientGetMember.Wrap(attachErr).GenWithStackByCause() + } 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()) + 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/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) + }) +} diff --git a/pkg/ratelimit/concurrency_limiter.go b/pkg/ratelimit/concurrency_limiter.go index e5379bc48cc..eff5657830c 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) @@ -108,43 +108,72 @@ 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 l.current >= l.limit { - l.waiting++ + if err := ctx.Err(); err != nil { + l.mu.Unlock() + return nil, err + } + 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() - token.released = false - l.current++ - l.waiting-- + 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 { + 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 16c31ebe2b9..af33221c9b2 100644 --- a/pkg/ratelimit/concurrency_limiter_test.go +++ b/pkg/ratelimit/concurrency_limiter_test.go @@ -102,6 +102,114 @@ 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 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 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() 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)