Merge branch 'master' into fix-flush

.github/workflows/check.yaml

@@ -8,21 +8,11 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 20
     steps:
-      - uses: actions/setup-go@v3
-        with:
-          go-version: '1.21'
       - name: Checkout code
-        uses: actions/checkout@v3
-      - name: Restore cache
-        uses: actions/cache@v3
+        uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
         with:
-          path: |
-            ~/go/pkg/mod
-            ~/.cache/go-build
-            **/.dashboard_download_cache
-          key: ${{ runner.os }}-golang-${{ hashFiles('**/go.sum') }}
-          restore-keys: |
-            ${{ runner.os }}-golang
+          go-version: '1.21'
       - name: Make Check
         run: |
           SWAGGER=1 make build

.github/workflows/label.yaml

@@ -7,7 +7,7 @@ jobs:
   add_labels:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/github-script@v4
+      - uses: actions/github-script@v7
         name: Add labels
         with:
           script: |

.github/workflows/pd-docker-image.yaml

@@ -15,10 +15,10 @@ jobs:
     strategy:
       fail-fast: true
     steps:
-      - uses: actions/setup-go@v3
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
         with:
           go-version: '1.21'
-      - name: Checkout code
-        uses: actions/checkout@v3
       - name: Make
         run: make docker-image

.github/workflows/pd-tests.yaml

@@ -29,20 +29,11 @@ jobs:
     outputs:
       job-total: 13
     steps:
-      - uses: actions/setup-go@v3
-        with:
-          go-version: '1.21'
       - name: Checkout code
-        uses: actions/checkout@v3
-      - name: Restore cache
-        uses: actions/cache@v3
+        uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
         with:
-          path: |
-            ~/go/pkg/mod
-            ~/.cache/go-build
-            **/.tools
-            **/.dashboard_download_cache
-          key: ${{ runner.os }}-go-${{ matrix.worker_id }}-${{ hashFiles('**/go.sum') }}
+          go-version: '1.21'
       - name: Make Test
         env:
           WORKER_ID: ${{ matrix.worker_id }}
@@ -53,20 +44,21 @@ jobs:
           mv covprofile covprofile_$WORKER_ID
           sed -i "/failpoint_binding/d" covprofile_$WORKER_ID
       - name: Upload coverage result ${{ matrix.worker_id }}
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
-          name: cover-reports
+          name: cover-reports-${{ matrix.worker_id }}
           path: covprofile_${{ matrix.worker_id }}
   report-coverage:
     needs: chunks
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Download chunk report
-        uses: actions/download-artifact@v2
+        uses: actions/download-artifact@v4
         with:
-          name: cover-reports
+          pattern: cover-reports-*
+          merge-multiple: true
       - name: Merge
         env:
           TOTAL_JOBS: ${{needs.chunks.outputs.job-total}}

.github/workflows/tso-consistency-test.yaml

@@ -8,10 +8,10 @@ jobs:
   tso-consistency-test:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/setup-go@v3
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
         with:
           go-version: '1.21'
-      - name: Checkout code
-        uses: actions/checkout@v3
       - name: Make TSO Consistency Test
         run: make test-tso-consistency

.github/workflows/tso-function-test.yaml

@@ -21,10 +21,10 @@ jobs:
   tso-function-test:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/setup-go@v3
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
         with:
           go-version: '1.21'
-      - name: Checkout code
-        uses: actions/checkout@v3
       - name: Make TSO Function Test
         run: make test-tso-function

client/client.go

@@ -798,34 +798,19 @@ func (c *client) GetLocalTSAsync(ctx context.Context, dcLocation string) TSFutur
 		defer span.Finish()
 	}
 
-	req := c.getTSORequest(ctx, dcLocation)
-	if err := c.dispatchTSORequestWithRetry(req); err != nil {
-		req.tryDone(err)
-	}
-	return req
-}
-
-func (c *client) getTSORequest(ctx context.Context, dcLocation string) *tsoRequest {
-	req := tsoReqPool.Get().(*tsoRequest)
-	// Set needed fields in the request before using it.
-	req.start = time.Now()
-	req.clientCtx = c.ctx
-	req.requestCtx = ctx
-	req.physical = 0
-	req.logical = 0
-	req.dcLocation = dcLocation
-	return req
+	return c.dispatchTSORequestWithRetry(ctx, dcLocation)
 }
 
 const (
 	dispatchRetryDelay = 50 * time.Millisecond
 	dispatchRetryCount = 2
 )
 
-func (c *client) dispatchTSORequestWithRetry(req *tsoRequest) error {
+func (c *client) dispatchTSORequestWithRetry(ctx context.Context, dcLocation string) TSFuture {
 	var (
 		retryable bool
 		err       error
+		req       *tsoRequest
 	)
 	for i := 0; i < dispatchRetryCount; i++ {
 		// Do not delay for the first time.
@@ -838,12 +823,22 @@ func (c *client) dispatchTSORequestWithRetry(req *tsoRequest) error {
 			err = errs.ErrClientGetTSO.FastGenByArgs("tso client is nil")
 			continue
 		}
+		// Get a new request from the pool if it's nil or not from the current pool.
+		if req == nil || req.pool != tsoClient.tsoReqPool {
+			req = tsoClient.getTSORequest(ctx, dcLocation)
+		}
 		retryable, err = tsoClient.dispatchRequest(req)
 		if !retryable {
 			break
 		}
 	}
-	return err
+	if err != nil {
+		if req == nil {
+			return newTSORequestFastFail(err)
+		}
+		req.tryDone(err)
+	}
+	return req
 }
 
 func (c *client) GetTS(ctx context.Context) (physical int64, logical int64, err error) {

client/tso_client.go

@@ -43,33 +43,6 @@ type TSOClient interface {
 	GetMinTS(ctx context.Context) (int64, int64, error)
 }
 
-type tsoRequest struct {
-	start      time.Time
-	clientCtx  context.Context
-	requestCtx context.Context
-	done       chan error
-	physical   int64
-	logical    int64
-	dcLocation string
-}
-
-var tsoReqPool = sync.Pool{
-	New: func() any {
-		return &tsoRequest{
-			done:     make(chan error, 1),
-			physical: 0,
-			logical:  0,
-		}
-	},
-}
-
-func (req *tsoRequest) tryDone(err error) {
-	select {
-	case req.done <- err:
-	default:
-	}
-}
-
 type tsoClient struct {
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -84,6 +57,8 @@ type tsoClient struct {
 	// tso allocator leader is switched.
 	tsoAllocServingURLSwitchedCallback []func()
 
+	// tsoReqPool is the pool to recycle `*tsoRequest`.
+	tsoReqPool *sync.Pool
 	// tsoDispatcher is used to dispatch different TSO requests to
 	// the corresponding dc-location TSO channel.
 	tsoDispatcher sync.Map // Same as map[string]*tsoDispatcher
@@ -104,11 +79,20 @@ func newTSOClient(
 ) *tsoClient {
 	ctx, cancel := context.WithCancel(ctx)
 	c := &tsoClient{
-		ctx:                       ctx,
-		cancel:                    cancel,
-		option:                    option,
-		svcDiscovery:              svcDiscovery,
-		tsoStreamBuilderFactory:   factory,
+		ctx:                     ctx,
+		cancel:                  cancel,
+		option:                  option,
+		svcDiscovery:            svcDiscovery,
+		tsoStreamBuilderFactory: factory,
+		tsoReqPool: &sync.Pool{
+			New: func() any {
+				return &tsoRequest{
+					done:     make(chan error, 1),
+					physical: 0,
+					logical:  0,
+				}
+			},
+		},
 		checkTSDeadlineCh:         make(chan struct{}),
 		checkTSODispatcherCh:      make(chan struct{}, 1),
 		updateTSOConnectionCtxsCh: make(chan struct{}, 1),
@@ -155,6 +139,19 @@ func (c *tsoClient) Close() {
 	log.Info("tso client is closed")
 }
 
+func (c *tsoClient) getTSORequest(ctx context.Context, dcLocation string) *tsoRequest {
+	req := c.tsoReqPool.Get().(*tsoRequest)
+	// Set needed fields in the request before using it.
+	req.start = time.Now()
+	req.pool = c.tsoReqPool
+	req.requestCtx = ctx
+	req.clientCtx = c.ctx
+	req.physical = 0
+	req.logical = 0
+	req.dcLocation = dcLocation
+	return req
+}
+
 // GetTSOAllocators returns {dc-location -> TSO allocator leader URL} connection map
 func (c *tsoClient) GetTSOAllocators() *sync.Map {
 	return &c.tsoAllocators

client/tso_dispatcher.go

@@ -115,38 +115,6 @@ func (c *tsoClient) dispatchRequest(request *tsoRequest) (bool, error) {
 	return false, nil
 }
 
-// TSFuture is a future which promises to return a TSO.
-type TSFuture interface {
-	// Wait gets the physical and logical time, it would block caller if data is not available yet.
-	Wait() (int64, int64, error)
-}
-
-func (req *tsoRequest) Wait() (physical int64, logical int64, err error) {
-	// If tso command duration is observed very high, the reason could be it
-	// takes too long for Wait() be called.
-	start := time.Now()
-	cmdDurationTSOAsyncWait.Observe(start.Sub(req.start).Seconds())
-	select {
-	case err = <-req.done:
-		defer trace.StartRegion(req.requestCtx, "pdclient.tsoReqDone").End()
-		err = errors.WithStack(err)
-		defer tsoReqPool.Put(req)
-		if err != nil {
-			cmdFailDurationTSO.Observe(time.Since(req.start).Seconds())
-			return 0, 0, err
-		}
-		physical, logical = req.physical, req.logical
-		now := time.Now()
-		cmdDurationWait.Observe(now.Sub(start).Seconds())
-		cmdDurationTSO.Observe(now.Sub(req.start).Seconds())
-		return
-	case <-req.requestCtx.Done():
-		return 0, 0, errors.WithStack(req.requestCtx.Err())
-	case <-req.clientCtx.Done():
-		return 0, 0, errors.WithStack(req.clientCtx.Err())
-	}
-}
-
 func (c *tsoClient) updateTSODispatcher() {
 	// Set up the new TSO dispatcher and batch controller.
 	c.GetTSOAllocators().Range(func(dcLocationKey, _ any) bool {