client/tso: decouple tsoClient and tsoDispatcher (#8134)

ref #8047

- Decouple `tsoClient` and `tsoDispatcher` to appropriately categorize different methods.
- Update `connectionCtxsUpdater` to support the Local TSO.

Signed-off-by: JmPotato <[email protected]>

client/client.go

@@ -301,7 +301,7 @@ func (k *serviceModeKeeper) close() {
 		fallthrough
 	case pdpb.ServiceMode_PD_SVC_MODE:
 		if k.tsoClient != nil {
-			k.tsoClient.Close()
+			k.tsoClient.close()
 		}
 	case pdpb.ServiceMode_UNKNOWN_SVC_MODE:
 	}
@@ -651,11 +651,11 @@ func (c *client) resetTSOClientLocked(mode pdpb.ServiceMode) {
 		log.Warn("[pd] intend to switch to unknown service mode, just return")
 		return
 	}
-	newTSOCli.Setup()
+	newTSOCli.setup()
 	// Replace the old TSO client.
 	oldTSOClient := c.tsoClient
 	c.tsoClient = newTSOCli
-	oldTSOClient.Close()
+	oldTSOClient.close()
 	// Replace the old TSO service discovery if needed.
 	oldTSOSvcDiscovery := c.tsoSvcDiscovery
 	// If newTSOSvcDiscovery is nil, that's expected, as it means we are switching to PD service mode and

client/tso_batch_controller.go

@@ -139,9 +139,11 @@ func (tbc *tsoBatchController) adjustBestBatchSize() {
 func (tbc *tsoBatchController) finishCollectedRequests(physical, firstLogical int64, suffixBits uint32, err error) {
 	for i := 0; i < tbc.collectedRequestCount; i++ {
 		tsoReq := tbc.collectedRequests[i]
+		// Retrieve the request context before the request is done to trace without race.
+		requestCtx := tsoReq.requestCtx
 		tsoReq.physical, tsoReq.logical = physical, tsoutil.AddLogical(firstLogical, int64(i), suffixBits)
-		defer trace.StartRegion(tsoReq.requestCtx, "pdclient.tsoReqDequeue").End() // nolint
 		tsoReq.tryDone(err)
+		trace.StartRegion(requestCtx, "pdclient.tsoReqDequeue").End()
 	}
 	// Prevent the finished requests from being processed again.
 	tbc.collectedRequestCount = 0