Implement the keyspace group merge checker

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

errors.toml

@@ -766,6 +766,11 @@ error = '''
 the keyspace group id is invalid, %s
 '''
 
+["PD:tso:ErrKeyspaceGroupIsMerging"]
+error = '''
+the keyspace group %d is merging
+'''
+
 ["PD:tso:ErrKeyspaceGroupNotInitialized"]
 error = '''
 the keyspace group %d isn't initialized

pkg/errs/errno.go

@@ -55,6 +55,7 @@ var (
 	ErrKeyspaceGroupNotInitialized      = errors.Normalize("the keyspace group %d isn't initialized", errors.RFCCodeText("PD:tso:ErrKeyspaceGroupNotInitialized"))
 	ErrKeyspaceNotAssigned              = errors.Normalize("the keyspace %d isn't assigned to any keyspace group", errors.RFCCodeText("PD:tso:ErrKeyspaceNotAssigned"))
 	ErrGetMinTS                         = errors.Normalize("get min ts failed, %s", errors.RFCCodeText("PD:tso:ErrGetMinTS"))
+	ErrKeyspaceGroupIsMerging           = errors.Normalize("the keyspace group %d is merging", errors.RFCCodeText("PD:tso:ErrKeyspaceGroupIsMerging"))
 )
 
 // member errors

pkg/tso/allocator_manager.go

@@ -1406,3 +1406,16 @@ func (am *AllocatorManager) GetLeaderAddr() string {
 	}
 	return leaderAddrs[0]
 }
+
+// Construct the timestampOracle path prefix, which is:
+//  1. for the default keyspace group:
+//     "" in /pd/{cluster_id}/timestamp
+//  2. for the non-default keyspace groups:
+//     {group}/gta in /ms/{cluster_id}/tso/{group}/gta/timestamp
+func (am *AllocatorManager) getKeyspaceGroupTSPath(groupID uint32) string {
+	tsPath := ""
+	if am.kgID != mcsutils.DefaultKeyspaceGroupID {
+		tsPath = path.Join(fmt.Sprintf("%05d", groupID), globalTSOAllocatorEtcdPrefix)
+	}
+	return tsPath
+}

pkg/tso/global_allocator.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"path"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -89,16 +88,6 @@ func NewGlobalTSOAllocator(
 	am *AllocatorManager,
 	startGlobalLeaderLoop bool,
 ) Allocator {
-	// Construct the timestampOracle path prefix, which is:
-	// 1. for the default keyspace group:
-	//     "" in /pd/{cluster_id}/timestamp
-	// 2. for the non-default keyspace groups:
-	//     {group}/gta in /ms/{cluster_id}/tso/{group}/gta/timestamp
-	tsPath := ""
-	if am.kgID != mcsutils.DefaultKeyspaceGroupID {
-		tsPath = path.Join(fmt.Sprintf("%05d", am.kgID), globalTSOAllocatorEtcdPrefix)
-	}
-
 	ctx, cancel := context.WithCancel(ctx)
 	gta := &GlobalTSOAllocator{
 		ctx:    ctx,
@@ -107,8 +96,7 @@ func NewGlobalTSOAllocator(
 		member: am.member,
 		timestampOracle: &timestampOracle{
 			client:                 am.member.GetLeadership().GetClient(),
-			rootPath:               am.rootPath,
-			tsPath:                 tsPath,
+			tsPath:                 am.getKeyspaceGroupTSPath(am.kgID),
 			storage:                am.storage,
 			saveInterval:           am.saveInterval,
 			updatePhysicalInterval: am.updatePhysicalInterval,

pkg/tso/keyspace_group_manager.go

@@ -34,13 +34,15 @@ import (
 	"github.com/tikv/pd/pkg/mcs/discovery"
 	mcsutils "github.com/tikv/pd/pkg/mcs/utils"
 	"github.com/tikv/pd/pkg/member"
+	"github.com/tikv/pd/pkg/slice"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
 	"github.com/tikv/pd/pkg/utils/logutil"
 	"github.com/tikv/pd/pkg/utils/memberutil"
 	"github.com/tikv/pd/pkg/utils/tsoutil"
+	"github.com/tikv/pd/pkg/utils/typeutil"
 	"go.etcd.io/etcd/clientv3"
 	"go.etcd.io/etcd/mvcc/mvccpb"
 	"go.uber.org/zap"
@@ -50,6 +52,9 @@ const (
 	keyspaceGroupsElectionPath = mcsutils.KeyspaceGroupsKey + "/election"
 	// primaryKey is the key for keyspace group primary election.
 	primaryKey = "primary"
+	// mergingCheckInterval is the interval for merging check to see if the keyspace groups
+	// merging process could be moved forward.
+	mergingCheckInterval = 5 * time.Second
 )
 
 type state struct {
@@ -381,12 +386,9 @@ func (kgm *KeyspaceGroupManager) Close() {
 }
 
 func (kgm *KeyspaceGroupManager) isAssignedToMe(group *endpoint.KeyspaceGroup) bool {
-	for _, member := range group.Members {
-		if member.Address == kgm.tsoServiceID.ServiceAddr {
-			return true
-		}
-	}
-	return false
+	return slice.AnyOf(group.Members, func(i int) bool {
+		return group.Members[i].Address == kgm.tsoServiceID.ServiceAddr
+	})
 }
 
 // updateKeyspaceGroup applies the given keyspace group. If the keyspace group is just assigned to
@@ -410,6 +412,12 @@ func (kgm *KeyspaceGroupManager) updateKeyspaceGroup(group *endpoint.KeyspaceGro
 		return
 	}
 
+	// If this host owns a replica of the keyspace group which is the merge target, it should
+	// run the merging checker.
+	if group.IsMergeTarget() {
+		go kgm.mergingChecker(kgm.ctx, group.ID, group.MergeState.MergeList)
+	}
+
 	// If this host is already assigned a replica of this keyspace group, i.e., the election member
 	// is already initialized, just update the meta.
 	if oldAM, oldGroup := kgm.getKeyspaceGroupMeta(group.ID); oldAM != nil {
@@ -729,6 +737,10 @@ func (kgm *KeyspaceGroupManager) HandleTSORequest(
 	if err != nil {
 		return pdpb.Timestamp{}, curKeyspaceGroupID, err
 	}
+	err = kgm.checkTSOMerge(curKeyspaceGroupID)
+	if err != nil {
+		return pdpb.Timestamp{}, curKeyspaceGroupID, err
+	}
 	ts, err = am.HandleRequest(dcLocation, count)
 	return ts, curKeyspaceGroupID, err
 }
@@ -881,3 +893,170 @@ func (kgm *KeyspaceGroupManager) finishSplitKeyspaceGroup(id uint32) error {
 	kgm.kgs[id] = splitGroup
 	return nil
 }
+
+func (kgm *KeyspaceGroupManager) finishMergeKeyspaceGroup(id uint32) error {
+	kgm.Lock()
+	defer kgm.Unlock()
+	// Check if the keyspace group is in the merging state.
+	mergeTarget := kgm.kgs[id]
+	if !mergeTarget.IsMergeTarget() {
+		return nil
+	}
+	// Check if the HTTP client is initialized.
+	if kgm.httpClient == nil {
+		return nil
+	}
+	statusCode, err := apiutil.DoDelete(
+		kgm.httpClient,
+		kgm.cfg.GeBackendEndpoints()+keyspaceGroupsAPIPrefix+fmt.Sprintf("/%d/merge", id))
+	if err != nil {
+		return err
+	}
+	if statusCode != http.StatusOK {
+		log.Warn("failed to finish merging keyspace group",
+			zap.Uint32("keyspace-group-id", id),
+			zap.Int("status-code", statusCode))
+		return errs.ErrSendRequest.FastGenByArgs()
+	}
+	// Pre-update the split keyspace group split state in memory.
+	mergeTarget.MergeState = nil
+	kgm.kgs[id] = mergeTarget
+	return nil
+}
+
+// mergingChecker is used to check if the keyspace group is in merge state, and if so, it will
+// make sure the newly merged TSO keep consistent with the original ones.
+func (kgm *KeyspaceGroupManager) mergingChecker(ctx context.Context, mergeTargetID uint32, mergeList []uint32) {
+	log.Info("start to merge the keyspace group",
+		zap.Uint32("merge-target-id", mergeTargetID),
+		zap.Any("merge-list", mergeList))
+	checkTicker := time.NewTicker(mergingCheckInterval)
+	defer checkTicker.Stop()
+	// Prepare the merge map.
+	mergeMap := make(map[uint32]struct{}, len(mergeList))
+	for _, id := range mergeList {
+		mergeMap[id] = struct{}{}
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			log.Info("merging checker is closed",
+				zap.Uint32("merge-target-id", mergeTargetID),
+				zap.Any("merge-list", mergeList))
+			return
+		case <-checkTicker.C:
+		}
+		// Check if current TSO node is the merge target TSO primary node.
+		am, err := kgm.GetAllocatorManager(mergeTargetID)
+		if err != nil {
+			log.Error("failed to get the merge target allocator manager",
+				zap.Uint32("keyspace-group-id", mergeTargetID),
+				zap.Any("merge-list", mergeList),
+				zap.Error(err))
+			continue
+		}
+		// If the current TSO node is not the merge target TSO primary node,
+		// we still need to keep this loop running to avoid unexpected primary changes.
+		if !am.IsLeader() {
+			log.Debug("current tso node is not the merge target primary",
+				zap.Uint32("merge-target-id", mergeTargetID),
+				zap.Any("merge-list", mergeList))
+			continue
+		}
+		// Check if the keyspace group primaries in the merge map are all gone.
+		if len(mergeMap) != 0 {
+			err := kgm.tsoSvcStorage.RunInTxn(ctx, func(txn kv.Txn) error {
+				for id := range mergeMap {
+					leaderPath := path.Join(kgm.primaryPathBuilder.getKeyspaceGroupIDPath(id), primaryKey)
+					val, err := txn.Load(leaderPath)
+					if err != nil {
+						return err
+					}
+					if len(val) == 0 {
+						delete(mergeMap, id)
+					}
+				}
+				return nil
+			})
+			if err != nil {
+				log.Error("failed to check if the keyspace group primaries in the merge list are all gone",
+					zap.Uint32("merge-target-id", mergeTargetID),
+					zap.Any("merge-list", mergeList),
+					zap.Any("remaining", mergeMap),
+					zap.Error(err))
+				continue
+			}
+		}
+		if len(mergeMap) > 0 {
+			continue
+		}
+		// All the keyspace group primaries in the merge list are gone,
+		// update the newly merged TSO to make sure it is greater than the original ones.
+		var mergedTS time.Time
+		for _, id := range mergeList {
+			ts, err := kgm.tsoSvcStorage.LoadTimestamp(am.getKeyspaceGroupTSPath(id))
+			if err != nil || ts == typeutil.ZeroTime {
+				log.Error("failed to load the keyspace group TSO",
+					zap.Uint32("merge-target-id", mergeTargetID),
+					zap.Any("merge-list", mergeList),
+					zap.Uint32("merge-id", id),
+					zap.Time("ts", ts),
+					zap.Error(err))
+				mergedTS = typeutil.ZeroTime
+				break
+			}
+			if ts.After(mergedTS) {
+				mergedTS = ts
+			}
+		}
+		if mergedTS == typeutil.ZeroTime {
+			continue
+		}
+		// Update the newly merged TSO.
+		// TODO: support the Local TSO Allocator.
+		allocator, err := am.GetAllocator(GlobalDCLocation)
+		if err != nil {
+			log.Error("failed to get the allocator",
+				zap.Uint32("merge-target-id", mergeTargetID),
+				zap.Any("merge-list", mergeList),
+				zap.Error(err))
+			continue
+		}
+		err = allocator.SetTSO(
+			tsoutil.GenerateTS(tsoutil.GenerateTimestamp(mergedTS, 1)),
+			true, true)
+		if err != nil {
+			log.Error("failed to update the newly merged TSO",
+				zap.Uint32("merge-target-id", mergeTargetID),
+				zap.Any("merge-list", mergeList),
+				zap.Time("merged-ts", mergedTS),
+				zap.Error(err))
+			continue
+		}
+		// Finish the merge.
+		err = kgm.finishMergeKeyspaceGroup(mergeTargetID)
+		if err != nil {
+			log.Error("failed to finish the merge",
+				zap.Uint32("merge-target-id", mergeTargetID),
+				zap.Any("merge-list", mergeList),
+				zap.Error(err))
+			continue
+		}
+		log.Info("finished merging keyspace group",
+			zap.Uint32("merge-target-id", mergeTargetID),
+			zap.Any("merge-list", mergeList),
+			zap.Time("merged-ts", mergedTS))
+		return
+	}
+}
+
+func (kgm *KeyspaceGroupManager) checkTSOMerge(
+	keyspaceGroupID uint32,
+) error {
+	_, group := kgm.getKeyspaceGroupMeta(keyspaceGroupID)
+	if !group.IsMerging() {
+		return nil
+	}
+	return errs.ErrKeyspaceGroupIsMerging.FastGenByArgs(keyspaceGroupID)
+}

pkg/tso/local_allocator.go

@@ -72,7 +72,6 @@ func NewLocalTSOAllocator(
 		leadership:       leadership,
 		timestampOracle: &timestampOracle{
 			client:                 leadership.GetClient(),
-			rootPath:               am.rootPath,
 			tsPath:                 tsPath,
 			storage:                am.storage,
 			saveInterval:           am.saveInterval,

pkg/tso/tso.go

@@ -60,8 +60,7 @@ type tsoObject struct {
 
 // timestampOracle is used to maintain the logic of TSO.
 type timestampOracle struct {
-	client   *clientv3.Client
-	rootPath string
+	client *clientv3.Client
 	// When tsPath is empty, it means that it is a global timestampOracle.
 	tsPath  string
 	storage endpoint.TSOStorage