add meta service group fallback ratio (tikv#393)

* add fallback ratio

Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com>

* fix typo

Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com>

* Update server/config/config.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* fix lint

Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com>

* fix lint

Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com>

* Update server/config/config.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* handle comment

Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com>

* clean

Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com>

* Update pkg/keyspace/meta_service_group.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update server/config/config.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update pkg/keyspace/meta_service_group.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: AmoebaProtozoa

pkg/keyspace/keyspace.go

@@ -83,6 +83,8 @@ type Config interface {
 	GetAutoAssignMetaServiceGroups() bool
 	SetMetaServiceGroups(map[string]string)
 	GetMetaServiceGroups() map[string]string
+	GetMetaServiceGroupsFallbackRatio() float64
+	SetMetaServiceGroupsFallbackRatio(float64)
 }
 
 // Manager manages keyspace related data.
@@ -244,6 +246,7 @@ func (manager *Manager) UpdateConfig(cfg Config) error {
 		manager.mgm.updateConfig(
 			cfg.GetAutoAssignMetaServiceGroups(),
 			cfg.GetMetaServiceGroups(),
+			cfg.GetMetaServiceGroupsFallbackRatio(),
 		)
 	}
 	return manager.SetGlobalSafePointV2()
@@ -285,9 +288,16 @@ func (manager *Manager) CreateKeyspace(request *CreateKeyspaceRequest) (*keyspac
 	if assignToMetaServiceGroup {
 		metaServiceGroup, err := manager.mgm.AssignToGroup(1)
 		if err != nil {
+			log.Error("[keyspace] failed to assign keyspace to meta-service group",
+				zap.Uint32("keyspace-id", newID),
+				zap.String("name", request.Name),
+				zap.Error(err),
+			)
 			return nil, err
 		}
-		request.Config[MetaServiceGroupIDKey] = metaServiceGroup
+		if metaServiceGroup != "" {
+			request.Config[MetaServiceGroupIDKey] = metaServiceGroup
+		}
 	}
 
 	// Create a disabled keyspace meta for tikv-server to get the config on keyspace split.

pkg/keyspace/keyspace_test.go

@@ -53,14 +53,15 @@ func TestKeyspaceTestSuite(t *testing.T) {
 }
 
 type mockConfig struct {
-	PreAlloc                    []string
-	WaitRegionSplit             bool
-	DisableRawKVRegionSplit     bool
-	WaitRegionSplitTimeout      typeutil.Duration
-	CheckRegionSplitInterval    typeutil.Duration
-	EnableGlobalSafePointV2     bool
-	AutoAssignMetaServiceGroups bool
-	MetaServiceGroups           map[string]string
+	PreAlloc                       []string
+	WaitRegionSplit                bool
+	DisableRawKVRegionSplit        bool
+	WaitRegionSplitTimeout         typeutil.Duration
+	CheckRegionSplitInterval       typeutil.Duration
+	EnableGlobalSafePointV2        bool
+	AutoAssignMetaServiceGroups    bool
+	MetaServiceGroups              map[string]string
+	MetaServiceGroupsFallbackRatio float64
 }
 
 func (m *mockConfig) SetEnableGlobalSafePointV2(isEnable bool) {
@@ -107,16 +108,28 @@ func (m *mockConfig) SetMetaServiceGroups(metaServiceGroups map[string]string) {
 	m.MetaServiceGroups = metaServiceGroups
 }
 
+func (m *mockConfig) GetMetaServiceGroupsFallbackRatio() float64 {
+	return m.MetaServiceGroupsFallbackRatio
+}
+
+func (m *mockConfig) SetMetaServiceGroupsFallbackRatio(ratio float64) {
+	m.MetaServiceGroupsFallbackRatio = ratio
+}
+
 func (suite *keyspaceTestSuite) SetupTest() {
 	re := suite.Require()
 	suite.ctx, suite.cancel = context.WithCancel(context.Background())
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
 	allocator := mockid.NewIDAllocator()
+	cfg := mockConfig{
+		AutoAssignMetaServiceGroups: true,
+		MetaServiceGroups:           mockMetaServiceGroups(),
+	}
 	kgm := NewKeyspaceGroupManager(suite.ctx, store, nil)
-	mgm := NewMetaServiceGroupManager(suite.ctx, store, true, mockMetaServiceGroups())
+	mgm := NewMetaServiceGroupManager(suite.ctx, store, &cfg)
 	mustEnableMetaServiceGroups(re, mgm, mockMetaServiceGroups())
 	var err error
-	suite.manager, err = NewKeyspaceManager(suite.ctx, store, nil, allocator, &mockConfig{}, kgm, mgm)
+	suite.manager, err = NewKeyspaceManager(suite.ctx, store, nil, allocator, &cfg, kgm, mgm)
 	re.NoError(err)
 	re.NoError(kgm.Bootstrap(suite.ctx))
 	re.NoError(suite.manager.Bootstrap())

pkg/keyspace/meta_service_group.go

@@ -17,10 +17,13 @@ package keyspace
 import (
 	"context"
 	"math"
+	"math/rand"
 
+	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/syncutil"
+	"go.uber.org/zap"
 )
 
 // MetaServiceGroupManager manages external meta-service groups.
@@ -30,20 +33,21 @@ type MetaServiceGroupManager struct {
 	syncutil.RWMutex
 	autoAssign        bool
 	metaServiceGroups map[string]string
+	fallbackRatio     float64
 }
 
 // NewMetaServiceGroupManager creates a new MetaServiceGroupManager.
 func NewMetaServiceGroupManager(
 	ctx context.Context,
 	store endpoint.MetaServiceGroupStorage,
-	autoAssign bool,
-	metaServiceGroups map[string]string,
+	config Config,
 ) *MetaServiceGroupManager {
 	return &MetaServiceGroupManager{
 		ctx:               ctx,
 		store:             store,
-		autoAssign:        autoAssign,
-		metaServiceGroups: metaServiceGroups,
+		autoAssign:        config.GetAutoAssignMetaServiceGroups(),
+		metaServiceGroups: config.GetMetaServiceGroups(),
+		fallbackRatio:     config.GetMetaServiceGroupsFallbackRatio(),
 	}
 }
 
@@ -100,6 +104,13 @@ func (m *MetaServiceGroupManager) PatchStatus(groupID string, patch *MetaService
 func (m *MetaServiceGroupManager) AssignToGroup(count int) (string, error) {
 	m.RLock()
 	defer m.RUnlock()
+	if roll := rand.Float64(); roll < m.fallbackRatio {
+		log.Info("[keyspace] fallback meta-service group assignment to PD due to fallback ratio",
+			zap.Float64("roll", roll),
+			zap.Float64("fallback ratio", m.fallbackRatio),
+		)
+		return "", nil
+	}
 	var (
 		assignedGroup       string
 		assignedGroupStatus *endpoint.MetaServiceGroupStatus
@@ -119,7 +130,10 @@ func (m *MetaServiceGroupManager) AssignToGroup(count int) (string, error) {
 			}
 		}
 		if assignedGroup == "" {
-			return errNoAvailableMetaServiceGroups
+			log.Warn("[keyspace] fallback meta-service group assignment to PD due to no available meta-service group",
+				zap.Any("meta-service groups status", statusMap),
+			)
+			return nil
 		}
 		assignedGroupStatus.AssignmentCount += count
 		return m.store.SaveMetaServiceGroupStatus(txn, assignedGroup, assignedGroupStatus)
@@ -187,9 +201,10 @@ func (m *MetaServiceGroupManager) GetAutoAssign() bool {
 }
 
 // updateConfig updates currently available meta-service groups.
-func (m *MetaServiceGroupManager) updateConfig(autoAssign bool, metaServiceGroups map[string]string) {
+func (m *MetaServiceGroupManager) updateConfig(autoAssign bool, metaServiceGroups map[string]string, fallbackRatio float64) {
 	m.Lock()
 	defer m.Unlock()
 	m.autoAssign = autoAssign
 	m.metaServiceGroups = metaServiceGroups
+	m.fallbackRatio = fallbackRatio
 }

pkg/keyspace/meta_service_group_test.go

@@ -45,7 +45,11 @@ func mockMetaServiceGroups() map[string]string {
 func (suite *metaServiceGroupTestSuite) SetupTest() {
 	suite.ctx, suite.cancel = context.WithCancel(context.Background())
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	suite.manager = NewMetaServiceGroupManager(suite.ctx, store, true, mockMetaServiceGroups())
+	cfg := mockConfig{
+		AutoAssignMetaServiceGroups: true,
+		MetaServiceGroups:           mockMetaServiceGroups(),
+	}
+	suite.manager = NewMetaServiceGroupManager(suite.ctx, store, &cfg)
 }
 
 func (suite *metaServiceGroupTestSuite) TearDownTest() {
@@ -62,9 +66,10 @@ func (suite *metaServiceGroupTestSuite) TestInitialState() {
 		re.Zero(status.AssignmentCount)
 		re.False(status.Enabled)
 	}
-	// Assign should return error due to no available groups.
-	_, err = suite.manager.AssignToGroup(1)
-	re.Error(err)
+	// Assignment should fallback to PD
+	group, err := suite.manager.AssignToGroup(1)
+	re.NoError(err)
+	re.Empty(group)
 }
 
 func (suite *metaServiceGroupTestSuite) TestAssignToGroup() {
@@ -165,7 +170,7 @@ func (suite *metaServiceGroupTestSuite) TestUpdateEndpoints() {
 	newMap := map[string]string{
 		"foo": "foo.bar.local",
 	}
-	suite.manager.updateConfig(true, newMap)
+	suite.manager.updateConfig(true, newMap, 0)
 	config := map[string]string{MetaServiceGroupIDKey: "foo"}
 	suite.manager.AttachEndpoints(config)
 	re.Equal("foo.bar.local", config[MetaServiceGroupAddressesKey], "should read from updated metaServiceGroups map")
@@ -191,7 +196,7 @@ func (suite *metaServiceGroupTestSuite) TestUpdateEndpointsAndUpdateAssignment()
 	// Add a new group "etcd-group-3"
 	newMap := mockMetaServiceGroups()
 	newMap["etcd-group-3"] = "etcd-group-3.tidb-serverless.cluster.svc.local"
-	suite.manager.updateConfig(true, newMap)
+	suite.manager.updateConfig(true, newMap, 0)
 
 	// Move the assignment from the originally assigned group to "etcd-group-3"
 	err = suite.manager.UpdateAssignment(assigned, "etcd-group-3")
@@ -212,3 +217,44 @@ func (suite *metaServiceGroupTestSuite) TestUpdateEndpointsAndUpdateAssignment()
 		re.Equal(0, statusMap[groupID].AssignmentCount, "other original groups should remain at 0")
 	}
 }
+
+func (suite *metaServiceGroupTestSuite) TestFallbackRatio() {
+	re := suite.Require()
+	for groupID := range mockMetaServiceGroups() {
+		enable := true
+		re.NoError(suite.manager.PatchStatus(groupID, &MetaServiceGroupStatusPatch{
+			Enabled: &enable,
+		}))
+	}
+	// Fallback ratio defaults to 0, all keyspace should be assigned to a group.
+	for range 10 {
+		assigned, err := suite.manager.AssignToGroup(1)
+		re.NoError(err)
+		re.NotEmpty(assigned, "expected AssignToGroup to return a non-empty group")
+	}
+	// Fallback ratio set to 0.5, some keyspace should be assigned to a group, some should fallback.
+	batchSize := 100
+	expectedFallbackRatio := 0.5
+	suite.manager.fallbackRatio = expectedFallbackRatio
+	totalFallback := 0
+	for range batchSize {
+		assigned, err := suite.manager.AssignToGroup(1)
+		re.NoError(err)
+		if assigned == "" {
+			totalFallback++
+		}
+	}
+	re.InDelta(expectedFallbackRatio, float64(totalFallback)/float64(batchSize), 0.1)
+
+	// Fallback ratio set to 1, all keyspace should fallback to pd.
+	suite.manager.fallbackRatio = 1.0
+	totalFallback = 0
+	for range batchSize {
+		assigned, err := suite.manager.AssignToGroup(1)
+		re.NoError(err)
+		if assigned == "" {
+			totalFallback++
+		}
+	}
+	re.Equal(batchSize, totalFallback, "all keyspace should fallback to pd when fallback ratio is 1.0")
+}

pkg/keyspace/util.go

@@ -91,8 +91,7 @@ var (
 	// ErrUnsupportedOperationInKeyspace is used to indicate this is an unsupported operation.
 	ErrUnsupportedOperationInKeyspace = errors.New("it's a unsupported operation")
 
-	errNoAvailableMetaServiceGroups = errors.New("no available meta-service groups")
-	errUnknownMetaServiceGroup      = errors.New("unknown meta-service group")
+	errUnknownMetaServiceGroup = errors.New("unknown meta-service group")
 
 	// stateTransitionTable lists all allowed next state for the given current state.
 	// Note that transit from any state to itself is allowed for idempotence.

server/config/config.go

@@ -900,6 +900,9 @@ type KeyspaceConfig struct {
 	AutoAssignMetaServiceGroups bool `toml:"auto-assign-meta-service-groups" json:"auto-assign-meta-service-groups,string"`
 	// MetaServiceGroups is the available external meta-service groups.
 	MetaServiceGroups map[string]string `toml:"meta-service-groups" json:"meta-service-groups"`
+	// MetaServiceGroupsFallbackRatio specifies the probability that a newly created keyspace will not be assigned a meta-service group.
+	// Valid range is [0, 1], where 0 means no fallback and 1 means always fallback to pd only.
+	MetaServiceGroupsFallbackRatio float64 `toml:"meta-service-groups-fallback-ratio" json:"meta-service-groups-fallback-ratio"`
 }
 
 // Validate checks if keyspace config falls within acceptable range.
@@ -916,6 +919,9 @@ func (c *KeyspaceConfig) Validate() error {
 			return errors.New("[keyspace] meta-service group addresses cannot be empty")
 		}
 	}
+	if c.MetaServiceGroupsFallbackRatio < 0 || c.MetaServiceGroupsFallbackRatio > 1 {
+		return errors.New("[keyspace] meta-service-groups-fallback-ratio should be in range [0, 1]")
+	}
 	return nil
 }
 
@@ -999,3 +1005,13 @@ func (c *KeyspaceConfig) GetMetaServiceGroups() map[string]string {
 func (c *KeyspaceConfig) SetMetaServiceGroups(metaServiceGroups map[string]string) {
 	c.MetaServiceGroups = metaServiceGroups
 }
+
+// GetMetaServiceGroupsFallbackRatio returns the fallback ratio for meta-service groups.
+func (c *KeyspaceConfig) GetMetaServiceGroupsFallbackRatio() float64 {
+	return c.MetaServiceGroupsFallbackRatio
+}
+
+// SetMetaServiceGroupsFallbackRatio sets the fallback ratio for meta-service groups.
+func (c *KeyspaceConfig) SetMetaServiceGroupsFallbackRatio(fallbackRatio float64) {
+	c.MetaServiceGroupsFallbackRatio = fallbackRatio
+}

server/server.go

@@ -513,7 +513,7 @@ func (s *Server) startServer(ctx context.Context) error {
 	if s.IsAPIServiceMode() {
 		s.keyspaceGroupManager = keyspace.NewKeyspaceGroupManager(s.ctx, s.storage, s.client)
 	}
-	s.metaServiceGroupManager = keyspace.NewMetaServiceGroupManager(s.ctx, s.storage, s.cfg.Keyspace.AutoAssignMetaServiceGroups, s.cfg.Keyspace.MetaServiceGroups)
+	s.metaServiceGroupManager = keyspace.NewMetaServiceGroupManager(s.ctx, s.storage, &s.cfg.Keyspace)
 	s.keyspaceManager, err = keyspace.NewKeyspaceManager(s.ctx, s.storage, s.cluster, keyspaceIDAllocator, &s.cfg.Keyspace, s.keyspaceGroupManager, s.metaServiceGroupManager)
 	if err != nil {
 		return err