client: add lock in tso.Request to avoid data race

client/clients/tso/client.go

@@ -174,13 +174,15 @@ func (c *Cli) scheduleUpdateTSOConnectionCtxs() {
 func (c *Cli) GetTSORequest(ctx context.Context) *Request {
 	req := c.tsoReqPool.Get().(*Request)
 	// Set needed fields in the request before using it.
+	req.mu.Lock()
 	req.start = time.Now()
 	req.pool = c.tsoReqPool
 	req.requestCtx = ctx
 	req.clientCtx = c.ctx
 	req.physical = 0
 	req.logical = 0
 	req.streamID = ""
+	req.mu.Unlock()
 	return req
 }
 

client/clients/tso/dispatcher.go

@@ -475,9 +475,11 @@ func (td *tsoDispatcher) processRequests(
 func tsoRequestFinisher(physical, firstLogical int64, streamID string) batch.FinisherFunc[*Request] {
 	return func(idx int, tsoReq *Request, err error) {
 		// Retrieve the request context before the request is done to trace without race.
+		tsoReq.mu.Lock()
 		requestCtx := tsoReq.requestCtx
 		tsoReq.physical, tsoReq.logical = physical, firstLogical+int64(idx)
 		tsoReq.streamID = streamID
+		tsoReq.mu.Unlock()
 		tsoReq.TryDone(err)
 		trace.StartRegion(requestCtx, "pdclient.tsoReqDequeue").End()
 	}

client/clients/tso/request.go

@@ -38,6 +38,7 @@ var (
 
 // Request is a TSO request.
 type Request struct {
+	mu         sync.RWMutex
 	requestCtx context.Context
 	clientCtx  context.Context
 	done       chan error
@@ -57,11 +58,15 @@ func (req *Request) IsFrom(pool *sync.Pool) bool {
 	if req == nil {
 		return false
 	}
+	req.mu.RLock()
+	defer req.mu.RUnlock()
 	return req.pool == pool
 }
 
 // TryDone tries to send the result to the channel, it will not block.
 func (req *Request) TryDone(err error) {
+	req.mu.RLock()
+	defer req.mu.RUnlock()
 	select {
 	case req.done <- err:
 	default:
@@ -75,29 +80,38 @@ func (req *Request) Wait() (physical int64, logical int64, err error) {
 
 // waitCtx waits for the TSO result with specified ctx, while not using req.requestCtx.
 func (req *Request) waitCtx(ctx context.Context) (physical int64, logical int64, err error) {
+	req.mu.RLock()
+	startAt := req.start
+	clientCtx := req.clientCtx
+	pool := req.pool
+	requestCtxForTrace := req.requestCtx
+	req.mu.RUnlock()
+
 	// If tso command duration is observed very high, the reason could be it
 	// takes too long for Wait() be called.
 	start := time.Now()
-	metrics.CmdDurationTSOAsyncWait.Observe(start.Sub(req.start).Seconds())
+	metrics.CmdDurationTSOAsyncWait.Observe(start.Sub(startAt).Seconds())
 	select {
 	case err = <-req.done:
-		defer req.pool.Put(req)
-		defer trace.StartRegion(req.requestCtx, "pdclient.tsoReqDone").End()
+		defer pool.Put(req)
+		defer trace.StartRegion(requestCtxForTrace, "pdclient.tsoReqDone").End()
 		err = errors.WithStack(err)
 		now := time.Now()
 		if err != nil {
 			metrics.CmdFailedDurationTSOWait.Observe(now.Sub(start).Seconds())
-			metrics.CmdFailedDurationTSO.Observe(now.Sub(req.start).Seconds())
+			metrics.CmdFailedDurationTSO.Observe(now.Sub(startAt).Seconds())
 			return 0, 0, err
 		}
+		req.mu.RLock()
 		physical, logical = req.physical, req.logical
+		req.mu.RUnlock()
 		metrics.CmdDurationTSOWait.Observe(now.Sub(start).Seconds())
-		metrics.CmdDurationTSO.Observe(now.Sub(req.start).Seconds())
+		metrics.CmdDurationTSO.Observe(now.Sub(startAt).Seconds())
 		return
 	case <-ctx.Done():
 		return 0, 0, errors.WithStack(ctx.Err())
-	case <-req.clientCtx.Done():
-		return 0, 0, errors.WithStack(req.clientCtx.Err())
+	case <-clientCtx.Done():
+		return 0, 0, errors.WithStack(clientCtx.Err())
 	}
 }