client, pkg: handle canceled contexts in select paths

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")

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,24 +225,39 @@ 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
 		}
 	}
 
 	// 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
 		}

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)
+}

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
 	}
 }

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

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())
+}