check label

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

Author: HuSharp

integration_tests/pd_api_test.go

@@ -44,6 +44,7 @@ import (
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
+	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/stretchr/testify/suite"
 	"github.com/tikv/client-go/v2/oracle"
 	"github.com/tikv/client-go/v2/tikv"
@@ -103,9 +104,9 @@ func (c *storeSafeTsMockClient) CloseAddr(addr string) error {
 	return c.Client.CloseAddr(addr)
 }
 
-func (s *apiTestSuite) TestGetStoreMinResolvedTS() {
+func (s *apiTestSuite) TestGetClusterMinResolvedTS() {
 	util.EnableFailpoints()
-	// Try to get the minimum resolved timestamp of the store from PD.
+	// Try to get the minimum resolved timestamp of the cluster from PD.
 	require := s.Require()
 	require.Nil(failpoint.Enable("tikvclient/InjectMinResolvedTS", `return(100)`))
 	mockClient := storeSafeTsMockClient{
@@ -141,6 +142,53 @@ func (s *apiTestSuite) TestGetStoreMinResolvedTS() {
 	require.Equal(uint64(150), s.store.GetMinSafeTS(oracle.GlobalTxnScope))
 }
 
+func (s *apiTestSuite) TestDCLabelClusterMinResolvedTS() {
+	util.EnableFailpoints()
+	// Try to get the minimum resolved timestamp of the cluster from PD.
+	require := s.Require()
+	require.Nil(failpoint.Enable("tikvclient/InjectMinResolvedTS", `return(100)`))
+	mockClient := storeSafeTsMockClient{
+		Client: s.store.GetTiKVClient(),
+	}
+	s.store.SetTiKVClient(&mockClient)
+	var retryCount int
+	for s.store.GetMinSafeTS(oracle.GlobalTxnScope) != 100 {
+		time.Sleep(2 * time.Second)
+		if retryCount > 5 {
+			break
+		}
+		retryCount++
+	}
+	require.Equal(atomic.LoadInt32(&mockClient.requestCount), int32(0))
+	require.Equal(uint64(100), s.store.GetMinSafeTS(oracle.GlobalTxnScope))
+	defer func() {
+		s.Require().Nil(failpoint.Disable("tikvclient/InjectMinResolvedTS"))
+	}()
+
+	// Set DC label for store 1.
+	dcLabel := "testDC"
+	labels := []*metapb.StoreLabel{
+		{
+			Key:   tikv.DCLabelKey,
+			Value: dcLabel,
+		},
+	}
+	s.store.GetRegionCache().SetRegionCacheStore(1, tikvrpc.TiKV, 1, labels)
+
+	// Try to get the minimum resolved timestamp of the store from TiKV.
+	retryCount = 0
+	for s.store.GetMinSafeTS(dcLabel) != 150 {
+		time.Sleep(2 * time.Second)
+		if retryCount > 5 {
+			break
+		}
+		retryCount++
+	}
+
+	require.GreaterOrEqual(atomic.LoadInt32(&mockClient.requestCount), int32(1))
+	require.Equal(uint64(150), s.store.GetMinSafeTS(dcLabel))
+}
+
 func (s *apiTestSuite) TearDownTest() {
 	if s.store != nil {
 		s.Require().Nil(s.store.Close())

tikv/kv.go

@@ -547,24 +547,14 @@ func (s *KVStore) updateSafeTS(ctx context.Context) {
 	tikvClient := s.GetTiKVClient()
 	wg := &sync.WaitGroup{}
 	wg.Add(len(stores))
+	safeTS, txnScopeMap, err := s.buildTxnScopeMap(ctx)
 	for _, store := range stores {
 		storeID := store.StoreID()
 		storeAddr := store.GetAddr()
 		go func(ctx context.Context, wg *sync.WaitGroup, storeID uint64, storeAddr string) {
 			defer wg.Done()
 
-			var (
-				safeTS uint64
-				err    error
-			)
 			storeIDStr := strconv.Itoa(int(storeID))
-			// Try to get the minimum resolved timestamp of the store from PD.
-			if s.pdHttpClient != nil {
-				safeTS, err = s.pdHttpClient.GetStoreMinResolvedTS(ctx, storeID)
-				if err != nil {
-					logutil.BgLogger().Debug("get resolved TS from PD failed", zap.Error(err), zap.Uint64("store-id", storeID))
-				}
-			}
 			// If getting the minimum resolved timestamp from PD failed or returned 0, try to get it from TiKV.
 			if safeTS == 0 || err != nil {
 				resp, err := tikvClient.SendRequest(
@@ -601,18 +591,37 @@ func (s *KVStore) updateSafeTS(ctx context.Context) {
 		}(ctx, wg, storeID, storeAddr)
 	}
 
-	txnScopeMap := make(map[string][]uint64)
+	for txnScope, storeIDs := range txnScopeMap {
+		s.updateMinSafeTS(txnScope, storeIDs)
+	}
+	wg.Wait()
+}
+
+// build txnScopeMap and judge whether it is needed to get safeTS from PD.
+// - if stores label are global, return get cluster min resolved ts from pd.
+// - if contains dc label store, return try to get it from TiKV.
+func (s *KVStore) buildTxnScopeMap(ctx context.Context) (safeTS uint64, txnScopeMap map[string][]uint64, err error) {
+	isGlobal := true
+	txnScopeMap = make(map[string][]uint64)
+	stores := s.regionCache.GetStoresByType(tikvrpc.TiKV)
 	for _, store := range stores {
 		txnScopeMap[oracle.GlobalTxnScope] = append(txnScopeMap[oracle.GlobalTxnScope], store.StoreID())
 
 		if label, ok := store.GetLabelValue(DCLabelKey); ok {
 			txnScopeMap[label] = append(txnScopeMap[label], store.StoreID())
+			isGlobal = false
 		}
 	}
-	for txnScope, storeIDs := range txnScopeMap {
-		s.updateMinSafeTS(txnScope, storeIDs)
+
+	// Try to get the minimum resolved timestamp of the cluster from PD.
+	if s.pdHttpClient != nil && isGlobal {
+		safeTS, err = s.pdHttpClient.GetClusterMinResolvedTS(ctx)
+		if err != nil {
+			logutil.BgLogger().Debug("get resolved TS from PD failed", zap.Error(err))
+		}
 	}
-	wg.Wait()
+
+	return safeTS, txnScopeMap, err
 }
 
 // Variables defines the variables used by TiKV storage.

util/pd.go

@@ -86,12 +86,11 @@ func NewPDHTTPClient(
 	}
 }
 
-// GetStoreMinResolvedTS get store-level min-resolved-ts from pd.
-func (p *PDHTTPClient) GetStoreMinResolvedTS(ctx context.Context, storeID uint64) (uint64, error) {
+// GetClusterMinResolvedTS get cluster-level min-resolved-ts from pd.
+func (p *PDHTTPClient) GetClusterMinResolvedTS(ctx context.Context) (uint64, error) {
 	var err error
 	for _, addr := range p.addrs {
-		query := fmt.Sprintf("%s/%d", storeMinResolvedTSPrefix, storeID)
-		v, e := pdRequest(ctx, addr, query, p.cli, http.MethodGet, nil)
+		v, e := pdRequest(ctx, addr, storeMinResolvedTSPrefix, p.cli, http.MethodGet, nil)
 		if e != nil {
 			logutil.BgLogger().Debug("failed to get min resolved ts", zap.String("addr", addr), zap.Error(e))
 			err = e