client/tso: fix the bug that TSO may hang when switching mode

client/client.go

@@ -782,23 +782,42 @@
 	req := tsoReqPool.Get().(*tsoRequest)
 	req.requestCtx = ctx
 	req.clientCtx = c.ctx
-	tsoClient := c.getTSOClient()
 	req.start = time.Now()
 	req.dcLocation = dcLocation
 
-	if tsoClient == nil {
-		req.done <- errs.ErrClientGetTSO.FastGenByArgs("tso client is nil")
-		return req
+	if err := c.dispatchTSORequestWithRetry(req); err != nil {
+		req.done <- err
 	}
+	return req
+}
+
+const (
+	dispatchRetryDelay = 50 * time.Millisecond
+	dispatchRetryCount = 2
+)
 
-	if err := tsoClient.dispatchRequest(ctx, dcLocation, req); err != nil {
-		// Wait for a while and try again
-		time.Sleep(50 * time.Millisecond)
-		if err = tsoClient.dispatchRequest(ctx, dcLocation, req); err != nil {
-			req.done <- err
+func (c *client) dispatchTSORequestWithRetry(req *tsoRequest) error {
+	var (
+		retryable bool
+		err       error
+	)
+	for i := 0; i < dispatchRetryCount; i++ {
+		// Do not delay for the first time.
+		if i > 0 {
+			time.Sleep(dispatchRetryDelay)
+		}
+		// Get the tsoClient each time, as it may be initialized or switched during the process.
+		tsoClient := c.getTSOClient()
+		if tsoClient == nil {
+			err = errs.ErrClientGetTSO.FastGenByArgs("tso client is nil")
+			continue
+		}
+		retryable, err = tsoClient.dispatchRequest(req)
+		if !retryable {
+			break
 		}
 	}
-	return req
+	return err
 }
 
 func (c *client) GetTS(ctx context.Context) (physical int64, logical int64, err error) {

client/tso_client.go

@@ -80,7 +80,7 @@ type tsoClient struct {
 
 	// tsoDispatcher is used to dispatch different TSO requests to
 	// the corresponding dc-location TSO channel.
-	tsoDispatcher sync.Map // Same as map[string]chan *tsoRequest
+	tsoDispatcher sync.Map // Same as map[string]*tsoDispatcher
 	// dc-location -> deadline
 	tsDeadline sync.Map // Same as map[string]chan deadline
 	// dc-location -> *tsoInfo while the tsoInfo is the last TSO info

client/tso_dispatcher.go

@@ -73,22 +73,31 @@
 	}
 }
 
-func (c *tsoClient) dispatchRequest(ctx context.Context, dcLocation string, request *tsoRequest) error {
-	dispatcher, ok := c.tsoDispatcher.Load(dcLocation)
+func (c *tsoClient) dispatchRequest(request *tsoRequest) (bool, error) {
+	dispatcher, ok := c.tsoDispatcher.Load(request.dcLocation)
 	if !ok {
-		err := errs.ErrClientGetTSO.FastGenByArgs(fmt.Sprintf("unknown dc-location %s to the client", dcLocation))
-		log.Error("[tso] dispatch tso request error", zap.String("dc-location", dcLocation), errs.ZapError(err))
+		err := errs.ErrClientGetTSO.FastGenByArgs(fmt.Sprintf("unknown dc-location %s to the client", request.dcLocation))
+		log.Error("[tso] dispatch tso request error", zap.String("dc-location", request.dcLocation), errs.ZapError(err))
 		c.svcDiscovery.ScheduleCheckMemberChanged()
-		return err
+		// New dispatcher could be created in the meantime, which is retryable.
+		return true, err
 	}
 
 	defer trace.StartRegion(request.requestCtx, "pdclient.tsoReqEnqueue").End()
 	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case dispatcher.(*tsoDispatcher).tsoBatchController.tsoRequestCh <- request:
+	case <-request.requestCtx.Done():
+		// Caller cancelled the request, no need to retry.
+		return false, request.requestCtx.Err()
+	case <-request.clientCtx.Done():
+		// Client is closed, no need to retry.
+		return false, request.clientCtx.Err()
+	case <-c.ctx.Done():
+		// tsoClient is closed due to the PD service mode switch, which is retryable.
+		return true, c.ctx.Err()
+	default:
+		dispatcher.(*tsoDispatcher).tsoBatchController.tsoRequestCh <- request
 	}
-	return nil
+	return false, nil
 }
 
 // TSFuture is a future which promises to return a TSO.

tests/integrations/client/client_test.go

@@ -1062,9 +1062,23 @@ func TestCloseClient(t *testing.T) {
 	defer cluster.Destroy()
 	endpoints := runServer(re, cluster)
 	cli := setupCli(re, ctx, endpoints)
-	cli.GetTSAsync(context.TODO())
+	ts := cli.GetTSAsync(context.TODO())
 	time.Sleep(time.Second)
 	cli.Close()
+	physical, logical, err := ts.Wait()
+	if err == nil {
+		re.Greater(physical, int64(0))
+		re.Greater(logical, int64(0))
+	} else {
+		re.ErrorIs(err, context.Canceled)
+		re.Zero(physical)
+		re.Zero(logical)
+	}
+	ts = cli.GetTSAsync(context.TODO())
+	physical, logical, err = ts.Wait()
+	re.ErrorIs(err, context.Canceled)
+	re.Zero(physical)
+	re.Zero(logical)
 }
 
 type idAllocator struct {