feat(resourcemanager): add service limit for keyspace (#9354)

ref #9296

Add a keyspace-level service limiter to control the maximum request rate for each keyspace.
The service limit is an unburstable RU per second rate limit that applies across all resource groups within a keyspace.

Signed-off-by: JmPotato <github@ipotato.me>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

Author: JmPotato

pkg/mcs/resourcemanager/server/apis/v1/api.go

@@ -99,6 +99,12 @@ func (s *Service) RegisterRouter() {
 	configEndpoint.DELETE("/group/:name", s.deleteResourceGroup)
 	configEndpoint.GET("/controller", s.getControllerConfig)
 	configEndpoint.POST("/controller", s.setControllerConfig)
+	// Without keyspace name, it will get/set the service limit of the null keyspace.
+	configEndpoint.POST("/keyspace/service-limit", s.setKeyspaceServiceLimit)
+	configEndpoint.GET("/keyspace/service-limit", s.getKeyspaceServiceLimit)
+	// With keyspace name, it will get/set the service limit of the given keyspace.
+	configEndpoint.POST("/keyspace/service-limit/:keyspace_name", s.setKeyspaceServiceLimit)
+	configEndpoint.GET("/keyspace/service-limit/:keyspace_name", s.getKeyspaceServiceLimit)
 }
 
 func (s *Service) handler() http.Handler {
@@ -239,3 +245,62 @@ func (s *Service) setControllerConfig(c *gin.Context) {
 	}
 	c.String(http.StatusOK, "Success!")
 }
+
+// KeyspaceServiceLimitRequest is the request body for setting the service limit of the keyspace.
+type KeyspaceServiceLimitRequest struct {
+	ServiceLimit float64 `json:"service_limit"`
+}
+
+// SetKeyspaceServiceLimit
+//
+//	@Tags		ResourceManager
+//	@Summary	Set the service limit of the keyspace. If the keyspace is valid, the service limit will be set.
+//	@Param		keyspace_name		path	string	true	"Keyspace name"
+//	@Param		service_limit	body		object	true	"json params, keyspaceServiceLimitRequest"
+//	@Success	200				{string}	string	"Success!"
+//	@Failure	400				{string}	error
+//	@Router		/config/keyspace/service-limit/{keyspace_name} [post]
+func (s *Service) setKeyspaceServiceLimit(c *gin.Context) {
+	keyspaceName := c.Param("keyspace_name")
+	keyspaceIDValue, err := s.manager.GetKeyspaceIDByName(c, keyspaceName)
+	if err != nil {
+		c.String(http.StatusBadRequest, err.Error())
+		return
+	}
+	var req KeyspaceServiceLimitRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.String(http.StatusBadRequest, err.Error())
+		return
+	}
+	if req.ServiceLimit < 0 {
+		c.String(http.StatusBadRequest, "service_limit must be non-negative")
+		return
+	}
+	s.manager.SetKeyspaceServiceLimit(keyspaceIDValue.GetValue(), req.ServiceLimit)
+	c.String(http.StatusOK, "Success!")
+}
+
+// GetKeyspaceServiceLimit
+//
+//	@Tags		ResourceManager
+//	@Summary	Get the service limit of the keyspace. If the keyspace name is empty, it will return the service limit of the null keyspace.
+//	@Param		keyspace_name	path		string	true	"Keyspace name"
+//	@Success	200				{string}	json	format	of	rmserver.serviceLimiter
+//	@Failure	400				{string}	error
+//	@Failure	404				{string}	error
+//	@Router		/config/keyspace/service-limit/{keyspace_name} [get]
+func (s *Service) getKeyspaceServiceLimit(c *gin.Context) {
+	keyspaceName := c.Param("keyspace_name")
+	keyspaceIDValue, err := s.manager.GetKeyspaceIDByName(c, keyspaceName)
+	if err != nil {
+		c.String(http.StatusBadRequest, err.Error())
+		return
+	}
+	keyspaceID := keyspaceIDValue.GetValue()
+	limiter := s.manager.GetKeyspaceServiceLimiter(keyspaceID)
+	if limiter == nil {
+		c.String(http.StatusNotFound, fmt.Sprintf("keyspace manager not found with keyspace name: %s, id: %d", keyspaceName, keyspaceID))
+		return
+	}
+	c.IndentedJSON(http.StatusOK, limiter)
+}

pkg/mcs/resourcemanager/server/grpc_service.go

@@ -174,15 +174,27 @@ func (s *Service) AcquireTokenBuckets(stream rmpb.ResourceManager_AcquireTokenBu
 		if err := s.checkServing(); err != nil {
 			return err
 		}
-		targetPeriodMs := request.GetTargetRequestPeriodMs()
-		clientUniqueID := request.GetClientUniqueId()
-		resps := &rmpb.TokenBucketsResponse{}
+		var (
+			targetPeriodMs = request.GetTargetRequestPeriodMs()
+			clientUniqueID = request.GetClientUniqueId()
+			resps          = &rmpb.TokenBucketsResponse{}
+			logFields      = make([]zap.Field, 2)
+		)
 		for _, req := range request.Requests {
+			keyspaceID := extractKeyspaceID(req.GetKeyspaceId())
 			resourceGroupName := req.GetResourceGroupName()
+			logFields[0] = zap.Uint32("keyspace-id", keyspaceID)
+			logFields[1] = zap.String("resource-group", resourceGroupName)
+			// Get keyspace resource group manager to apply service limit later.
+			krgm := s.manager.getKeyspaceResourceGroupManager(keyspaceID)
+			if krgm == nil {
+				log.Warn("keyspace resource group manager not found", logFields...)
+				continue
+			}
 			// Get the resource group from manager to acquire token buckets.
-			rg := s.manager.GetMutableResourceGroup(extractKeyspaceID(req.GetKeyspaceId()), resourceGroupName)
+			rg := s.manager.GetMutableResourceGroup(keyspaceID, resourceGroupName)
 			if rg == nil {
-				log.Warn("resource group not found", zap.String("resource-group", resourceGroupName))
+				log.Warn("resource group not found", logFields...)
 				continue
 			}
 			// Send the consumption to update the metrics.
@@ -202,18 +214,19 @@ func (s *Service) AcquireTokenBuckets(stream rmpb.ResourceManager_AcquireTokenBu
 				var tokens *rmpb.GrantedRUTokenBucket
 				for _, re := range req.GetRuItems().GetRequestRU() {
 					if re.Type == rmpb.RequestUnitType_RU {
-						tokens = rg.RequestRU(now, re.Value, targetPeriodMs, clientUniqueID)
+						tokens = rg.RequestRU(now, re.Value, targetPeriodMs, clientUniqueID, krgm.getServiceLimiter())
 					}
 					if tokens == nil {
 						continue
 					}
 					resp.GrantedRUTokens = append(resp.GrantedRUTokens, tokens)
 				}
 			case rmpb.GroupMode_RawMode:
-				log.Warn("not supports the resource type", zap.String("resource-group", resourceGroupName), zap.String("mode", rmpb.GroupMode_name[int32(rmpb.GroupMode_RawMode)]))
+				log.Warn("not supports the resource type",
+					append(logFields, zap.String("mode", rmpb.GroupMode_name[int32(rmpb.GroupMode_RawMode)]))...)
 				continue
 			}
-			log.Debug("finish token request from", zap.String("resource-group", resourceGroupName))
+			log.Debug("finish token request from", logFields...)
 			resps.Responses = append(resps.Responses, resp)
 		}
 		if err := stream.Send(resps); err != nil {

pkg/mcs/resourcemanager/server/keyspace_manager.go

@@ -56,13 +56,15 @@ type keyspaceResourceGroupManager struct {
 
 	keyspaceID uint32
 	storage    endpoint.ResourceGroupStorage
+	sl         *serviceLimiter
 }
 
 func newKeyspaceResourceGroupManager(keyspaceID uint32, storage endpoint.ResourceGroupStorage) *keyspaceResourceGroupManager {
 	return &keyspaceResourceGroupManager{
 		groups:     make(map[string]*ResourceGroup),
 		keyspaceID: keyspaceID,
 		storage:    storage,
+		sl:         newServiceLimiter(keyspaceID, 0),
 	}
 }
 
@@ -237,3 +239,17 @@ func (krgm *keyspaceResourceGroupManager) persistResourceGroupRunningState() {
 		krgm.RUnlock()
 	}
 }
+
+func (krgm *keyspaceResourceGroupManager) setServiceLimiter(serviceLimit float64) {
+	krgm.RLock()
+	sl := krgm.sl
+	krgm.RUnlock()
+	// Set the new service limit to the limiter.
+	sl.setServiceLimit(serviceLimit)
+}
+
+func (krgm *keyspaceResourceGroupManager) getServiceLimiter() *serviceLimiter {
+	krgm.RLock()
+	defer krgm.RUnlock()
+	return krgm.sl
+}

pkg/mcs/resourcemanager/server/manager.go

@@ -67,6 +67,8 @@ type Manager struct {
 	consumptionDispatcher chan *consumptionItem
 	// cached keyspace name for each keyspace ID.
 	keyspaceNameLookup map[uint32]string
+	// used to get the keyspace ID by name.
+	keyspaceIDLookup map[string]uint32
 }
 
 // ConfigProvider is used to get resource manager config from the given
@@ -83,6 +85,7 @@ func NewManager[T ConfigProvider](srv bs.Server) *Manager {
 		krgms:                 make(map[uint32]*keyspaceResourceGroupManager),
 		consumptionDispatcher: make(chan *consumptionItem, defaultConsumptionChanSize),
 		keyspaceNameLookup:    make(map[uint32]string),
+		keyspaceIDLookup:      make(map[string]uint32),
 	}
 	// The first initialization after the server is started.
 	srv.AddStartCallback(func() {
@@ -116,6 +119,21 @@ func (m *Manager) GetStorage() endpoint.ResourceGroupStorage {
 	return m.storage
 }
 
+// GetKeyspaceServiceLimiter returns the service limit of the keyspace.
+func (m *Manager) GetKeyspaceServiceLimiter(keyspaceID uint32) *serviceLimiter {
+	krgm := m.getKeyspaceResourceGroupManager(keyspaceID)
+	if krgm == nil {
+		return nil
+	}
+	return krgm.getServiceLimiter().Clone()
+}
+
+// SetKeyspaceServiceLimit sets the service limit of the keyspace.
+func (m *Manager) SetKeyspaceServiceLimit(keyspaceID uint32, serviceLimit float64) {
+	// If the keyspace is not found, create a new keyspace resource group manager.
+	m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, true).setServiceLimiter(serviceLimit)
+}
+
 func (m *Manager) getOrCreateKeyspaceResourceGroupManager(keyspaceID uint32, initDefault bool) *keyspaceResourceGroupManager {
 	m.Lock()
 	krgm, ok := m.krgms[keyspaceID]
@@ -369,7 +387,9 @@ func (m *Manager) getKeyspaceNameByID(ctx context.Context, id uint32) (string, e
 		return "", nil
 	}
 	// Try to get the keyspace name from the cache first.
+	m.RLock()
 	name, ok := m.keyspaceNameLookup[id]
+	m.RUnlock()
 	if ok {
 		return name, nil
 	}
@@ -391,10 +411,51 @@ func (m *Manager) getKeyspaceNameByID(ctx context.Context, id uint32) (string, e
 		return "", fmt.Errorf("got an empty keyspace name by id %d", id)
 	}
 	// Update the cache.
-	m.keyspaceNameLookup[id] = loadedName
+	m.updateKeyspaceNameLookup(id, loadedName)
 	return loadedName, nil
 }
 
+func (m *Manager) updateKeyspaceNameLookup(id uint32, name string) {
+	m.Lock()
+	defer m.Unlock()
+	m.keyspaceNameLookup[id] = name
+	m.keyspaceIDLookup[name] = id
+}
+
+// GetKeyspaceIDByName gets the keyspace ID by name.
+func (m *Manager) GetKeyspaceIDByName(ctx context.Context, name string) (*rmpb.KeyspaceIDValue, error) {
+	if len(name) == 0 {
+		return &rmpb.KeyspaceIDValue{Value: constant.NullKeyspaceID}, nil
+	}
+	m.RLock()
+	id, ok := m.keyspaceIDLookup[name]
+	m.RUnlock()
+	if ok {
+		return &rmpb.KeyspaceIDValue{Value: id}, nil
+	}
+	var (
+		loadedID uint32
+		err      error
+	)
+	err = m.storage.RunInTxn(ctx, func(txn kv.Txn) error {
+		ok, loadedID, err = m.storage.LoadKeyspaceID(txn, name)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		log.Error("failed to get the keyspace id", zap.String("keyspace-name", name), zap.Error(err))
+		return nil, err
+	}
+	if !ok {
+		return nil, fmt.Errorf("keyspace not found with name: %s", name)
+	}
+	// Update the cache.
+	m.updateKeyspaceNameLookup(loadedID, name)
+	return &rmpb.KeyspaceIDValue{Value: loadedID}, nil
+}
+
 func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 	defer logutil.LogPanic()
 	defer m.wg.Done()

pkg/mcs/resourcemanager/server/manager_test.go

@@ -164,7 +164,10 @@ func prepareKeyspaceName(ctx context.Context, re *require.Assertions, manager *M
 	}
 	err := manager.storage.RunInTxn(ctx, func(txn kv.Txn) error {
 		err := manager.storage.SaveKeyspaceMeta(txn, keyspaceMeta)
-		return err
+		if err != nil {
+			return err
+		}
+		return manager.storage.SaveKeyspaceID(txn, keyspaceMeta.Id, keyspaceMeta.Name)
 	})
 	re.NoError(err)
 }
@@ -262,3 +265,106 @@ func TestCleanUpTicker(t *testing.T) {
 	re.NoError(err)
 	re.Equal("test_keyspace", keyspaceName)
 }
+
+func TestKeyspaceServiceLimit(t *testing.T) {
+	re := require.New(t)
+
+	storage := storage.NewStorageWithMemoryBackend()
+	m := NewManager[*mockConfigProvider](&mockConfigProvider{})
+	m.storage = storage
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	err := m.Init(ctx)
+	re.NoError(err)
+	// Test the default service limit is 0.0.
+	limiter := m.GetKeyspaceServiceLimiter(constant.NullKeyspaceID)
+	re.NotNil(limiter)
+	re.Equal(0.0, limiter.ServiceLimit)
+	re.Equal(0.0, limiter.AvailableTokens)
+	group := &rmpb.ResourceGroup{
+		Name:     "test_group",
+		Mode:     rmpb.GroupMode_RUMode,
+		Priority: 5,
+		RUSettings: &rmpb.GroupRequestUnitSettings{
+			RU: &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{
+					FillRate:   100,
+					BurstLimit: 200,
+				},
+			},
+		},
+		KeyspaceId: &rmpb.KeyspaceIDValue{Value: 1},
+	}
+	// Test the limiter of the non-existing keyspace is nil.
+	limiter = m.GetKeyspaceServiceLimiter(group.KeyspaceId.Value)
+	re.Nil(limiter)
+	// Test the limiter of the newly created keyspace is 0.0.
+	err = m.AddResourceGroup(group)
+	re.NoError(err)
+	limiter = m.GetKeyspaceServiceLimiter(1)
+	re.Equal(0.0, limiter.ServiceLimit)
+	re.Equal(0.0, limiter.AvailableTokens)
+	// Test set the service limit of the keyspace.
+	m.SetKeyspaceServiceLimit(1, 100.0)
+	limiter = m.GetKeyspaceServiceLimiter(1)
+	re.Equal(100.0, limiter.ServiceLimit)
+	re.Equal(0.0, limiter.AvailableTokens) // When setting from 0 to positive, available tokens remain 0
+	// Test set the service limit of the non-existing keyspace.
+	limiter = m.GetKeyspaceServiceLimiter(2)
+	re.Nil(limiter)
+	m.SetKeyspaceServiceLimit(2, 100.0)
+	limiter = m.GetKeyspaceServiceLimiter(2)
+	re.Equal(100.0, limiter.ServiceLimit)
+	re.Equal(0.0, limiter.AvailableTokens)
+	// Ensure the keyspace resource group manager is initialized correctly.
+	krgm := m.getKeyspaceResourceGroupManager(2)
+	re.NotNil(krgm)
+	re.Equal(uint32(2), krgm.keyspaceID)
+	re.Equal(DefaultResourceGroupName, krgm.getMutableResourceGroup(DefaultResourceGroupName).Name)
+}
+
+func TestKeyspaceNameLookup(t *testing.T) {
+	re := require.New(t)
+	m := prepareManager()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	err := m.Init(ctx)
+	re.NoError(err)
+	// Get the null keyspace ID by an empty name.
+	idValue, err := m.GetKeyspaceIDByName(ctx, "")
+	re.NoError(err)
+	re.NotNil(idValue)
+	re.Equal(constant.NullKeyspaceID, idValue.Value)
+	// Get the non-existing keyspace ID by name.
+	idValue, err = m.GetKeyspaceIDByName(ctx, "non-existing-keyspace")
+	re.Error(err)
+	re.Nil(idValue)
+	// Get the null keyspace name.
+	name, err := m.getKeyspaceNameByID(ctx, constant.NullKeyspaceID)
+	re.NoError(err)
+	re.Empty(name)
+	// Get the non-existing keyspace name.
+	name, err = m.getKeyspaceNameByID(ctx, 1)
+	re.Error(err)
+	re.Empty(name)
+	// Get the keyspace ID by name first, then get the keyspace name by ID.
+	prepareKeyspaceName(ctx, re, m, &rmpb.KeyspaceIDValue{Value: 1}, "test_keyspace")
+	idValue, err = m.GetKeyspaceIDByName(ctx, "test_keyspace")
+	re.NoError(err)
+	re.NotNil(idValue)
+	re.Equal(uint32(1), idValue.Value)
+	name, err = m.getKeyspaceNameByID(ctx, 1)
+	re.NoError(err)
+	re.Equal("test_keyspace", name)
+	// Get the keyspace name by ID first, then get the keyspace ID by name.
+	prepareKeyspaceName(ctx, re, m, &rmpb.KeyspaceIDValue{Value: 2}, "test_keyspace_2")
+	name, err = m.getKeyspaceNameByID(ctx, 2)
+	re.NoError(err)
+	re.Equal("test_keyspace_2", name)
+	idValue, err = m.GetKeyspaceIDByName(ctx, "test_keyspace_2")
+	re.NoError(err)
+	re.NotNil(idValue)
+	re.Equal(uint32(2), idValue.Value)
+}