keyspace: make TSO keyspace group auto split configurable

conf/config.toml

@@ -215,3 +215,15 @@
 ## Example:
 ## pre-alloc = ["admin", "user1", "user2"]
 # pre-alloc = []
+
+## enable-tso-keyspace-group-auto-split controls whether TSO keyspace groups
+## are automatically split when their keyspace count exceeds the threshold.
+# enable-tso-keyspace-group-auto-split = true
+
+## tso-keyspace-group-auto-split-threshold is the keyspace count threshold for
+## automatically splitting a TSO keyspace group.
+# tso-keyspace-group-auto-split-threshold = 40000
+
+## tso-keyspace-group-auto-split-patrol-interval is the interval for checking
+## TSO keyspace group size.
+# tso-keyspace-group-auto-split-patrol-interval = "15m"

pkg/keyspace/tso_keyspace_group.go

@@ -22,6 +22,7 @@
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"go.etcd.io/etcd/api/v3/mvccpb"
@@ -65,6 +66,32 @@
 	opDelete
 )
 
+type keyspaceGroupAutoSplitConfig struct {
+	enabled                     bool
+	keyspaceCountSplitThreshold int
+	patrolInterval              time.Duration
+}
+
+var defaultKeyspaceGroupAutoSplitConfig = keyspaceGroupAutoSplitConfig{
+	enabled:                     true,
+	keyspaceCountSplitThreshold: defaultKeyspaceCountSplitThreshold,
+	patrolInterval:              autoSplitKeyspaceGroupPatrolInterval,
+}
+
+// KeyspaceGroupManagerOption configures a GroupManager.
+type KeyspaceGroupManagerOption func(*GroupManager)
+
+// WithKeyspaceGroupAutoSplitConfig configures auto-splitting by keyspace count.
+func WithKeyspaceGroupAutoSplitConfig(
+	enabled bool,
+	keyspaceCountSplitThreshold int,
+	patrolInterval time.Duration,
+) KeyspaceGroupManagerOption {
+	return func(m *GroupManager) {
+		m.updateKeyspaceGroupAutoSplitConfig(enabled, keyspaceCountSplitThreshold, patrolInterval, false)
+	}
+}
+
 // GroupManager is the manager of keyspace group related data.
 type GroupManager struct {
 	ctx    context.Context
@@ -88,27 +115,36 @@
 	serviceRegistryMap map[string]string
 	// tsoNodesWatcher is the watcher for the registered tso servers.
 	tsoNodesWatcher *etcdutil.LoopWatcher
+
+	autoSplitConfig        atomic.Value
+	autoSplitConfigChanged chan struct{}
 }
 
 // NewKeyspaceGroupManager creates a Manager of keyspace group related data.
 func NewKeyspaceGroupManager(
 	ctx context.Context,
 	store endpoint.KeyspaceGroupStorage,
 	client *clientv3.Client,
+	opts ...KeyspaceGroupManagerOption,
 ) *GroupManager {
 	ctx, cancel := context.WithCancel(ctx)
 	groups := make(map[endpoint.UserKind]*indexedHeap)
 	for i := range endpoint.UserKindCount {
 		groups[i] = newIndexedHeap(int(mcs.MaxKeyspaceGroupCountInUse))
 	}
 	m := &GroupManager{
-		ctx:                ctx,
-		cancel:             cancel,
-		store:              store,
-		groups:             groups,
-		client:             client,
-		nodesBalancer:      balancer.GenByPolicy[string](defaultBalancerPolicy),
-		serviceRegistryMap: make(map[string]string),
+		ctx:                    ctx,
+		cancel:                 cancel,
+		store:                  store,
+		groups:                 groups,
+		client:                 client,
+		nodesBalancer:          balancer.GenByPolicy[string](defaultBalancerPolicy),
+		serviceRegistryMap:     make(map[string]string),
+		autoSplitConfigChanged: make(chan struct{}, 1),
+	}
+	m.autoSplitConfig.Store(defaultKeyspaceGroupAutoSplitConfig)
+	for _, opt := range opts {
+		opt(m)
 	}
 
 	// If the etcd client is not nil, start the watch loop for the registered tso servers.
@@ -120,6 +156,48 @@
 	return m
 }
 
+// UpdateKeyspaceGroupAutoSplitConfig updates the auto-split config used by the patrol loop.
+func (m *GroupManager) UpdateKeyspaceGroupAutoSplitConfig(
+	enabled bool,
+	keyspaceCountSplitThreshold int,
+	patrolInterval time.Duration,
+) {
+	m.updateKeyspaceGroupAutoSplitConfig(enabled, keyspaceCountSplitThreshold, patrolInterval, true)
+}
+
+func (m *GroupManager) updateKeyspaceGroupAutoSplitConfig(
+	enabled bool,
+	keyspaceCountSplitThreshold int,
+	patrolInterval time.Duration,
+	notify bool,
+) {
+	if keyspaceCountSplitThreshold <= 0 {
+		keyspaceCountSplitThreshold = defaultKeyspaceCountSplitThreshold
+	}
+	if patrolInterval <= 0 {
+		patrolInterval = autoSplitKeyspaceGroupPatrolInterval
+	}
+	m.autoSplitConfig.Store(keyspaceGroupAutoSplitConfig{
+		enabled:                     enabled,
+		keyspaceCountSplitThreshold: keyspaceCountSplitThreshold,
+		patrolInterval:              patrolInterval,
+	})
+	if notify {
+		select {
+		case m.autoSplitConfigChanged <- struct{}{}:
+		default:
+		}
+	}
+}
+
+func (m *GroupManager) getKeyspaceGroupAutoSplitConfig() keyspaceGroupAutoSplitConfig {
+	cfg, ok := m.autoSplitConfig.Load().(keyspaceGroupAutoSplitConfig)
+	if !ok {
+		return defaultKeyspaceGroupAutoSplitConfig
+	}
+	return cfg
+}
+
 // Bootstrap saves default keyspace group info and init group mapping in the memory.
 func (m *GroupManager) Bootstrap(ctx context.Context) error {
 	// Force the membership restriction that the default keyspace must belong to default keyspace group.
@@ -229,14 +307,18 @@
 }
 
 // patrolKeyspaceGroupSizeForAutoSplit periodically checks all tso keyspace groups.
-// If a group's keyspace count exceeds defaultKeyspaceCountSplitThreshold,
+// If a group's keyspace count exceeds the configured threshold,
 // 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)
+	cfg := m.getKeyspaceGroupAutoSplitConfig()
+	ticker := time.NewTicker(cfg.patrolInterval)
 	defer ticker.Stop()
-	log.Info("start to patrol keyspace group size for auto-split")
+	log.Info("start to patrol keyspace group size for auto-split",
+		zap.Bool("enabled", cfg.enabled),
+		zap.Int("keyspace-count-split-threshold", cfg.keyspaceCountSplitThreshold),
+		zap.Duration("patrol-interval", cfg.patrolInterval))
 	for {
 		select {
 		case <-m.ctx.Done():
@@ -245,6 +327,13 @@
 		case <-ctx.Done():
 			log.Info("the raftcluster is closed, stop patrolling keyspace group size for auto-split")
 			return
+		case <-m.autoSplitConfigChanged:
+			cfg = m.getKeyspaceGroupAutoSplitConfig()
+			ticker.Reset(cfg.patrolInterval)
+			log.Info("updated keyspace group auto-split config",
+				zap.Bool("enabled", cfg.enabled),
+				zap.Int("keyspace-count-split-threshold", cfg.keyspaceCountSplitThreshold),
+				zap.Duration("patrol-interval", cfg.patrolInterval))
 		case <-ticker.C:
 			m.doPatrolKeyspaceGroupSizeForAutoSplit(ctx)
 		}
@@ -260,6 +349,10 @@
 		return
 	default:
 	}
+	cfg := m.getKeyspaceGroupAutoSplitConfig()
+	if !cfg.enabled {
+		return
+	}
 	groups, err := m.store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0)
 	if err != nil {
 		log.Error("auto-split patrol failed to load all keyspace groups",
@@ -275,7 +368,7 @@
 			zap.Uint32("max-keyspace-group-count-in-use", mcs.MaxKeyspaceGroupCountInUse))
 		return
 	}
-	threshold := defaultKeyspaceCountSplitThreshold
+	threshold := cfg.keyspaceCountSplitThreshold
 	failpoint.Inject("autoSplitKeyspaceGroupThreshold", func() {
 		threshold = 5
 	})

pkg/keyspace/tso_keyspace_group_test.go

@@ -794,6 +794,64 @@ func (suite *keyspaceGroupTestSuite) TestDoPatrolKeyspaceGroupSizeForAutoSplitBe
 	re.Nil(kg1)
 }
 
+func (suite *keyspaceGroupTestSuite) TestDoPatrolKeyspaceGroupSizeForAutoSplitDisabled() {
+	re := suite.Require()
+	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
+	keyspaces := buildSequentialKeyspaces(0, defaultKeyspaceCountSplitThreshold+1)
+	savePatrolTestKeyspaceGroups(suite.ctx, suite.T(), store, &endpoint.KeyspaceGroup{
+		ID:        constant.DefaultKeyspaceGroupID,
+		UserKind:  endpoint.Basic.String(),
+		Keyspaces: keyspaces,
+		Members:   testKeyspaceGroupMembers(),
+	})
+
+	kgm := NewKeyspaceGroupManager(suite.ctx, store, nil)
+	re.NoError(kgm.Bootstrap(suite.ctx))
+	kgm.UpdateKeyspaceGroupAutoSplitConfig(false, defaultKeyspaceCountSplitThreshold, autoSplitKeyspaceGroupPatrolInterval)
+
+	kgm.doPatrolKeyspaceGroupSizeForAutoSplit(suite.ctx)
+
+	kg0, err := kgm.GetKeyspaceGroupByID(constant.DefaultKeyspaceGroupID)
+	re.NoError(err)
+	re.NotNil(kg0)
+	re.Equal(keyspaces, kg0.Keyspaces)
+	re.False(kg0.IsSplitting())
+	kg1, err := kgm.GetKeyspaceGroupByID(1)
+	re.NoError(err)
+	re.Nil(kg1)
+}
+
+func (suite *keyspaceGroupTestSuite) TestDoPatrolKeyspaceGroupSizeForAutoSplitCustomThreshold() {
+	re := suite.Require()
+	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
+	const threshold = 4
+	keyspaces := buildSequentialKeyspaces(0, threshold+1)
+	savePatrolTestKeyspaceGroups(suite.ctx, suite.T(), store, &endpoint.KeyspaceGroup{
+		ID:        constant.DefaultKeyspaceGroupID,
+		UserKind:  endpoint.Basic.String(),
+		Keyspaces: keyspaces,
+		Members:   testKeyspaceGroupMembers(),
+	})
+
+	kgm := NewKeyspaceGroupManager(suite.ctx, store, nil,
+		WithKeyspaceGroupAutoSplitConfig(true, threshold, autoSplitKeyspaceGroupPatrolInterval))
+	re.NoError(kgm.Bootstrap(suite.ctx))
+
+	kgm.doPatrolKeyspaceGroupSizeForAutoSplit(suite.ctx)
+
+	splitIdx := len(keyspaces) / 2
+	kg0, err := kgm.GetKeyspaceGroupByID(constant.DefaultKeyspaceGroupID)
+	re.NoError(err)
+	re.NotNil(kg0)
+	re.Equal(keyspaces[:splitIdx], kg0.Keyspaces)
+	re.True(kg0.IsSplitSource())
+	kg1, err := kgm.GetKeyspaceGroupByID(1)
+	re.NoError(err)
+	re.NotNil(kg1)
+	re.Equal(keyspaces[splitIdx:], kg1.Keyspaces)
+	re.True(kg1.IsSplitTarget())
+}
+
 func (suite *keyspaceGroupTestSuite) TestDoPatrolKeyspaceGroupSizeForAutoSplitSkipsSplittingAndMergingGroups() {
 	re := suite.Require()
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)

server/config/config.go

@@ -251,6 +251,10 @@ const (
 	minCheckRegionSplitInterval     = 1 * time.Millisecond
 	maxCheckRegionSplitInterval     = 100 * time.Millisecond
 
+	defaultEnableTSOKeyspaceGroupAutoSplit         = true
+	defaultTSOKeyspaceGroupAutoSplitThreshold      = 40000
+	defaultTSOKeyspaceGroupAutoSplitPatrolInterval = 15 * time.Minute
+
 	defaultEnableSchedulingFallback      = true
 	defaultEnableTSODynamicSwitching     = false
 	defaultEnableResourceManagerFallback = true
@@ -884,6 +888,12 @@ type KeyspaceConfig struct {
 	// MetaServiceGroups is the available external meta-service groups.
 	// The key is the meta-service group name, and the value is the corresponding endpoint.
 	MetaServiceGroups map[string]string `toml:"meta-service-groups" json:"meta-service-groups"`
+	// EnableTSOKeyspaceGroupAutoSplit indicates whether to auto-split TSO keyspace groups by keyspace count.
+	EnableTSOKeyspaceGroupAutoSplit bool `toml:"enable-tso-keyspace-group-auto-split" json:"enable-tso-keyspace-group-auto-split"`
+	// TSOKeyspaceGroupAutoSplitThreshold is the keyspace count threshold for auto-splitting a TSO keyspace group.
+	TSOKeyspaceGroupAutoSplitThreshold int `toml:"tso-keyspace-group-auto-split-threshold" json:"tso-keyspace-group-auto-split-threshold"`
+	// TSOKeyspaceGroupAutoSplitPatrolInterval is the patrol interval for TSO keyspace group auto-split.
+	TSOKeyspaceGroupAutoSplitPatrolInterval typeutil.Duration `toml:"tso-keyspace-group-auto-split-patrol-interval" json:"tso-keyspace-group-auto-split-patrol-interval"`
 }
 
 // Validate checks if keyspace config falls within acceptable range.
@@ -895,6 +905,16 @@ func (c *KeyspaceConfig) Validate() error {
 	if c.CheckRegionSplitInterval.Duration >= c.WaitRegionSplitTimeout.Duration {
 		return errors.New("[keyspace] check-region-split-interval should be less than wait-region-split-timeout")
 	}
+	return c.validateTSOKeyspaceGroupAutoSplit()
+}
+
+func (c *KeyspaceConfig) validateTSOKeyspaceGroupAutoSplit() error {
+	if c.TSOKeyspaceGroupAutoSplitThreshold <= 0 {
+		return errors.New("[keyspace] tso-keyspace-group-auto-split-threshold should be greater than 0")
+	}
+	if c.TSOKeyspaceGroupAutoSplitPatrolInterval.Duration <= 0 {
+		return errors.New("[keyspace] tso-keyspace-group-auto-split-patrol-interval should be greater than 0")
+	}
 	return nil
 }
 
@@ -908,8 +928,20 @@ func (c *KeyspaceConfig) adjust(meta *configutil.ConfigMetaData) error {
 	if !meta.IsDefined("check-region-split-interval") {
 		c.CheckRegionSplitInterval = typeutil.NewDuration(defaultCheckRegionSplitInterval)
 	}
+	if !meta.IsDefined("enable-tso-keyspace-group-auto-split") {
+		c.EnableTSOKeyspaceGroupAutoSplit = defaultEnableTSOKeyspaceGroupAutoSplit
+	}
+	if !meta.IsDefined("tso-keyspace-group-auto-split-threshold") {
+		c.TSOKeyspaceGroupAutoSplitThreshold = defaultTSOKeyspaceGroupAutoSplitThreshold
+	}
+	if !meta.IsDefined("tso-keyspace-group-auto-split-patrol-interval") {
+		c.TSOKeyspaceGroupAutoSplitPatrolInterval = typeutil.NewDuration(defaultTSOKeyspaceGroupAutoSplitPatrolInterval)
+	}
 
-	return AdjustMetaServiceGroups(c.MetaServiceGroups)
+	if err := AdjustMetaServiceGroups(c.MetaServiceGroups); err != nil {
+		return err
+	}
+	return c.validateTSOKeyspaceGroupAutoSplit()
 }
 
 // IsValidMetaServiceGroupID reports whether id is safe to use as a single path
@@ -999,3 +1031,18 @@ func (c *KeyspaceConfig) GetMetaServiceGroups() map[string]string {
 func (c *KeyspaceConfig) SetMetaServiceGroups(metaServiceGroups map[string]string) {
 	c.MetaServiceGroups = metaServiceGroups
 }
+
+// IsTSOKeyspaceGroupAutoSplitEnabled returns whether TSO keyspace group auto-split is enabled.
+func (c *KeyspaceConfig) IsTSOKeyspaceGroupAutoSplitEnabled() bool {
+	return c.EnableTSOKeyspaceGroupAutoSplit
+}
+
+// GetTSOKeyspaceGroupAutoSplitThreshold returns the keyspace count threshold for auto-splitting.
+func (c *KeyspaceConfig) GetTSOKeyspaceGroupAutoSplitThreshold() int {
+	return c.TSOKeyspaceGroupAutoSplitThreshold
+}
+
+// GetTSOKeyspaceGroupAutoSplitPatrolInterval returns the patrol interval for TSO keyspace group auto-split.
+func (c *KeyspaceConfig) GetTSOKeyspaceGroupAutoSplitPatrolInterval() time.Duration {
+	return c.TSOKeyspaceGroupAutoSplitPatrolInterval.Duration
+}