From 74aa891cdd601ac7fd52ab4db9a7adf001e7b280 Mon Sep 17 00:00:00 2001
From: MoCuishle28 <32541204+MoCuishle28@users.noreply.github.com>
Date: Tue, 22 Nov 2022 14:47:58 +0800
Subject: [PATCH] br: replace GetTS with GetTSWithRetry (#38948)

ref pingcap/tidb#36910
---
 br/pkg/restore/client.go      |  4 ++--
 br/pkg/restore/client_test.go | 21 ++++++++++-----------
 br/pkg/task/restore.go        |  2 +-
 br/pkg/task/stream.go         |  2 +-
 4 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/br/pkg/restore/client.go b/br/pkg/restore/client.go
index 81215e675c8c0..fb50cc321e87a 100644
--- a/br/pkg/restore/client.go
+++ b/br/pkg/restore/client.go
@@ -1427,7 +1427,7 @@ func (rc *Client) updateMetaAndLoadStats(ctx context.Context, input <-chan *Crea
 			}
 
 			// Not need to return err when failed because of update analysis-meta
-			restoreTS, err := rc.GetTS(ctx)
+			restoreTS, err := rc.GetTSWithRetry(ctx)
 			if err != nil {
 				log.Error("getTS failed", zap.Error(err))
 			} else {
@@ -2399,7 +2399,7 @@ func (rc *Client) RunGCRowsLoader(ctx context.Context) {
 func (rc *Client) InsertGCRows(ctx context.Context) error {
 	close(rc.deleteRangeQueryCh)
 	rc.deleteRangeQueryWaitGroup.Wait()
-	ts, err := rc.GetTS(ctx)
+	ts, err := rc.GetTSWithRetry(ctx)
 	if err != nil {
 		return errors.Trace(err)
 	}
diff --git a/br/pkg/restore/client_test.go b/br/pkg/restore/client_test.go
index c43c3c80c2452..afd3c802cdc3b 100644
--- a/br/pkg/restore/client_test.go
+++ b/br/pkg/restore/client_test.go
@@ -332,18 +332,17 @@ type fakePDClient struct {
 	pd.Client
 	stores []*metapb.Store
 
-	notLeader bool
+	notLeader  bool
+	retryTimes *int
 }
 
-var retryTimes int
-
 func (fpdc fakePDClient) GetAllStores(context.Context, ...pd.GetStoreOption) ([]*metapb.Store, error) {
 	return append([]*metapb.Store{}, fpdc.stores...), nil
 }
 
 func (fpdc fakePDClient) GetTS(ctx context.Context) (int64, int64, error) {
-	retryTimes++
-	if retryTimes >= 3 { // the mock PD leader switched successfully
+	(*fpdc.retryTimes)++
+	if *fpdc.retryTimes >= 3 { // the mock PD leader switched successfully
 		fpdc.notLeader = false
 	}
 
@@ -355,24 +354,24 @@ func (fpdc fakePDClient) GetTS(ctx context.Context) (int64, int64, error) {
 
 func TestGetTSWithRetry(t *testing.T) {
 	t.Run("PD leader is healthy:", func(t *testing.T) {
-		retryTimes = -1000
-		pDClient := fakePDClient{notLeader: false}
+		retryTimes := -1000
+		pDClient := fakePDClient{notLeader: false, retryTimes: &retryTimes}
 		client := restore.NewRestoreClient(pDClient, nil, defaultKeepaliveCfg, false)
 		_, err := client.GetTSWithRetry(context.Background())
 		require.NoError(t, err)
 	})
 
 	t.Run("PD leader failure:", func(t *testing.T) {
-		retryTimes = -1000
-		pDClient := fakePDClient{notLeader: true}
+		retryTimes := -1000
+		pDClient := fakePDClient{notLeader: true, retryTimes: &retryTimes}
 		client := restore.NewRestoreClient(pDClient, nil, defaultKeepaliveCfg, false)
 		_, err := client.GetTSWithRetry(context.Background())
 		require.Error(t, err)
 	})
 
 	t.Run("PD leader switch successfully", func(t *testing.T) {
-		retryTimes = 0
-		pDClient := fakePDClient{notLeader: true}
+		retryTimes := 0
+		pDClient := fakePDClient{notLeader: true, retryTimes: &retryTimes}
 		client := restore.NewRestoreClient(pDClient, nil, defaultKeepaliveCfg, false)
 		_, err := client.GetTSWithRetry(context.Background())
 		require.NoError(t, err)
diff --git a/br/pkg/task/restore.go b/br/pkg/task/restore.go
index cf1ce4682a09f..c16ada97ceb6b 100644
--- a/br/pkg/task/restore.go
+++ b/br/pkg/task/restore.go
@@ -555,7 +555,7 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 	g.Record(summary.RestoreDataSize, archiveSize)
 	//restore from tidb will fetch a general Size issue https://github.com/pingcap/tidb/issues/27247
 	g.Record("Size", archiveSize)
-	restoreTS, err := client.GetTS(ctx)
+	restoreTS, err := client.GetTSWithRetry(ctx)
 	if err != nil {
 		return errors.Trace(err)
 	}
diff --git a/br/pkg/task/stream.go b/br/pkg/task/stream.go
index 86036c2d1ffe2..cd0d1ed664f69 100644
--- a/br/pkg/task/stream.go
+++ b/br/pkg/task/stream.go
@@ -1129,7 +1129,7 @@ func restoreStream(
 	}
 	defer client.Close()
 
-	currentTS, err := client.GetTS(ctx)
+	currentTS, err := client.GetTSWithRetry(ctx)
 	if err != nil {
 		return errors.Trace(err)
 	}