client: refine some client code (tikv#4205)

* Refine some client code Signed-off-by: JmPotato <ghzpotato@gmail.com> * ref tikv#3149 Signed-off-by: JmPotato <ghzpotato@gmail.com> Co-authored-by: Ti Chi Robot <ti-community-prow-bot@tidb.io> Signed-off-by: hzh0425 <642256541@qq.com>
hzh0425 · Nov 16, 2021 · 5ff1c0f · 5ff1c0f
1 parent 701244d
commit 5ff1c0f
Showing 1 changed file with 8 additions and 13 deletions.
diff --git a/client/client.go b/client/client.go
@@ -284,16 +284,7 @@ func (c *client) updateTSODispatcher() {
 	c.allocators.Range(func(dcLocationKey, _ interface{}) bool {
 		dcLocation := dcLocationKey.(string)
 		if !c.checkTSODispatcher(dcLocation) {
-			log.Info("[pd] create tso dispatcher", zap.String("dc-location", dcLocation))
 			c.createTSODispatcher(dcLocation)
-			dispatcher, _ := c.tsoDispatcher.Load(dcLocation)
-			dispatcherCtx := dispatcher.(*tsoDispatcher).dispatcherCtx
-			tsoRequestCh := dispatcher.(*tsoDispatcher).tsoRequestCh
-			// Each goroutine is responsible for handling the tso stream request for its dc-location.
-			// The only case that will make the dispatcher goroutine exit
-			// is that the loopCtx is done, otherwise there is no circumstance
-			// this goroutine should exit.
-			go c.handleDispatcher(dispatcherCtx, dcLocation, tsoRequestCh)
 		}
 		return true
 	})
@@ -490,6 +481,7 @@ func (c *client) checkAllocator(dispatcherCtx context.Context, forwardCancel con
 		failpoint.Inject("unreachableNetwork", func() {
 			resp.Status = healthpb.HealthCheckResponse_UNKNOWN
 		})
+		healthCancel()
 		if err == nil && resp.GetStatus() == healthpb.HealthCheckResponse_SERVING {
 			// create a stream of the original allocator
 			cctx, cancel := context.WithCancel(dispatcherCtx)
@@ -503,11 +495,9 @@ func (c *client) checkAllocator(dispatcherCtx context.Context, forwardCancel con
 					cancel: cancel,
 					stream: stream,
 				}
-				healthCancel()
 				return
 			}
 		}
-		healthCancel()
 		select {
 		case <-dispatcherCtx.Done():
 			return
@@ -534,7 +524,13 @@ func (c *client) createTSODispatcher(dcLocation string) {
 		dispatcherCancel: dispatcherCancel,
 		tsoRequestCh:     make(chan *tsoRequest, maxMergeTSORequests),
 	}
+	// Each goroutine is responsible for handling the tso stream request for its dc-location.
+	// The only case that will make the dispatcher goroutine exit
+	// is that the loopCtx is done, otherwise there is no circumstance
+	// this goroutine should exit.
+	go c.handleDispatcher(dispatcherCtx, dcLocation, dispatcher.tsoRequestCh)
 	c.tsoDispatcher.Store(dcLocation, dispatcher)
+	log.Info("[pd] tso dispatcher created", zap.String("dc-location", dcLocation))
 }
 
 type streamCh chan struct {
@@ -914,11 +910,10 @@ func (c *client) followerClient() (pdpb.PDClient, string) {
 		}
 		healthCtx, healthCancel := context.WithTimeout(c.ctx, c.timeout)
 		resp, err := healthpb.NewHealthClient(cc).Check(healthCtx, &healthpb.HealthCheckRequest{Service: ""})
+		healthCancel()
 		if err == nil && resp.GetStatus() == healthpb.HealthCheckResponse_SERVING {
-			healthCancel()
 			return pdpb.NewPDClient(cc), addr
 		}
-		healthCancel()
 	}
 	return nil, ""
 }