tso: keyspace group auto splitting (#10932)

close #10933\n\nThe TSO keyspace group supports checking the keyspace count automatically every 15 minutes and triggering a split if a certain threshold is reached, preventing the corresponding etcd value from becoming too long.\n\nSigned-off-by: ystaticy <y_static_y@sina.com>

Author: ystaticy

pkg/keyspace/tso_keyspace_group.go

@@ -53,6 +53,11 @@ const (
 	allocNodesToKeyspaceGroupsInterval = 1 * time.Second
 	allocNodesTimeout                  = 1 * time.Second
 	allocNodesInterval                 = 10 * time.Millisecond
+	// defaultKeyspaceCountSplitThreshold is the keyspace count threshold for auto-splitting
+	// a keyspace group. When a group's keyspace count exceeds this value, a new group will be split automatically.
+	defaultKeyspaceCountSplitThreshold = 40000
+	// autoSplitKeyspaceGroupPatrolInterval is the interval for patrolling keyspace group size for auto-split.
+	autoSplitKeyspaceGroupPatrolInterval = 15 * time.Minute
 )
 
 const (
@@ -151,6 +156,8 @@ func (m *GroupManager) Bootstrap(ctx context.Context) error {
 	if m.client != nil {
 		m.wg.Add(1)
 		go m.allocNodesToAllKeyspaceGroups(ctx)
+		m.wg.Add(1)
+		go m.patrolKeyspaceGroupSizeForAutoSplit(ctx)
 	}
 	return nil
 }
@@ -221,6 +228,133 @@ func (m *GroupManager) allocNodesToAllKeyspaceGroups(ctx context.Context) {
 	}
 }
 
+// patrolKeyspaceGroupSizeForAutoSplit periodically checks all tso keyspace groups.
+// If a group's keyspace count exceeds defaultKeyspaceCountSplitThreshold,
+// it automatically splits a new group and moves about half of the keyspaces to the new group.
+func (m *GroupManager) patrolKeyspaceGroupSizeForAutoSplit(ctx context.Context) {
+	defer logutil.LogPanic()
+	defer m.wg.Done()
+	ticker := time.NewTicker(autoSplitKeyspaceGroupPatrolInterval)
+	defer ticker.Stop()
+	log.Info("start to patrol keyspace group size for auto-split")
+	for {
+		select {
+		case <-m.ctx.Done():
+			log.Info("server is closed, stop patrolling keyspace group size for auto-split")
+			return
+		case <-ctx.Done():
+			log.Info("the raftcluster is closed, stop patrolling keyspace group size for auto-split")
+			return
+		case <-ticker.C:
+			m.doPatrolKeyspaceGroupSizeForAutoSplit(ctx)
+		}
+	}
+}
+
+// doPatrolKeyspaceGroupSizeForAutoSplit checks once and performs at most one auto-split.
+func (m *GroupManager) doPatrolKeyspaceGroupSizeForAutoSplit(ctx context.Context) {
+	select {
+	case <-m.ctx.Done():
+		return
+	case <-ctx.Done():
+		return
+	default:
+	}
+	groups, err := m.store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0)
+	if err != nil {
+		log.Error("auto-split patrol failed to load all keyspace groups",
+			zap.Error(err))
+		return
+	}
+	if len(groups) == 0 {
+		return
+	}
+	nextID, ok := findNextAvailableKeyspaceGroupID(groups, mcs.MaxKeyspaceGroupCountInUse)
+	if !ok {
+		log.Warn("no available keyspace group id for auto-split, max id reached",
+			zap.Uint32("max-keyspace-group-count-in-use", mcs.MaxKeyspaceGroupCountInUse))
+		return
+	}
+	threshold := defaultKeyspaceCountSplitThreshold
+	failpoint.Inject("autoSplitKeyspaceGroupThreshold", func() {
+		threshold = 5
+	})
+
+	sortedGroups := make([]*endpoint.KeyspaceGroup, len(groups))
+	copy(sortedGroups, groups)
+	sort.Slice(sortedGroups, func(i, j int) bool {
+		if len(sortedGroups[i].Keyspaces) == len(sortedGroups[j].Keyspaces) {
+			return sortedGroups[i].ID < sortedGroups[j].ID
+		}
+		return len(sortedGroups[i].Keyspaces) > len(sortedGroups[j].Keyspaces)
+	})
+
+	for _, group := range sortedGroups {
+		if group.IsSplitting() || group.IsMerging() {
+			continue
+		}
+		if len(group.Members) < mcs.DefaultKeyspaceGroupReplicaCount {
+			continue
+		}
+		count := len(group.Keyspaces)
+		if count <= threshold {
+			continue
+		}
+		// Split about half of the keyspaces to the new group (excluding protected bootstrap/system keyspace).
+		splitIdx := count / 2
+		keyspacesToMove := make([]uint32, 0, count-splitIdx)
+		for i := splitIdx; i < count; i++ {
+			kid := group.Keyspaces[i]
+			if isProtectedKeyspaceID(kid) {
+				continue
+			}
+			keyspacesToMove = append(keyspacesToMove, kid)
+		}
+		if len(keyspacesToMove) == 0 {
+			log.Warn("no keyspaces to move for auto-split, skip",
+				zap.Uint32("keyspace-group-id", group.ID),
+				zap.Int("keyspace-count", count))
+			continue
+		}
+		err := m.SplitKeyspaceGroupByID(group.ID, nextID, keyspacesToMove)
+		if err != nil {
+			log.Error("failed to auto-split keyspace group by keyspace count",
+				zap.Uint32("source-id", group.ID),
+				zap.Uint32("target-id", nextID),
+				zap.Int("keyspace-count", count),
+				zap.Int("keyspaces-to-move", len(keyspacesToMove)),
+				zap.Error(err))
+			return
+		}
+		log.Info("auto-split keyspace group by keyspace count",
+			zap.Uint32("source-id", group.ID),
+			zap.Uint32("target-id", nextID),
+			zap.Int("keyspace-count", count),
+			zap.Int("keyspaces-moved", len(keyspacesToMove)))
+		return
+	}
+}
+
+// findNextAvailableKeyspaceGroupID returns the smallest unused keyspace group ID in
+// [1, maxCount). Group 0 is reserved for the default keyspace group.
+func findNextAvailableKeyspaceGroupID(groups []*endpoint.KeyspaceGroup, maxCount uint32) (uint32, bool) {
+	if maxCount <= 1 {
+		return 0, false
+	}
+	used := make(map[uint32]struct{}, len(groups))
+	for _, g := range groups {
+		if g.ID < maxCount {
+			used[g.ID] = struct{}{}
+		}
+	}
+	for id := uint32(1); id < maxCount; id++ {
+		if _, exists := used[id]; !exists {
+			return id, true
+		}
+	}
+	return 0, false
+}
+
 func (m *GroupManager) initTSONodesWatcher(client *clientv3.Client) {
 	putFn := func(kv *mvccpb.KeyValue) error {
 		s := &discovery.ServiceRegistryEntry{}