resource_manager: implement the token-consuming by client

client/resource_group/controller/controller.go

@@ -410,7 +410,7 @@ func (c *ResourceGroupsController) sendTokenBucketRequests(ctx context.Context,
 	}()
 }
 
-// OnRequestWait is used to check whether resource group has enough tokens. It maybe needs wait some time.
+// OnRequestWait is used to check whether resource group has enough tokens. It maybe needs to wait some time.
 func (c *ResourceGroupsController) OnRequestWait(
 	ctx context.Context, resourceGroupName string, info RequestInfo,
 ) (*rmpb.Consumption, error) {
@@ -837,13 +837,13 @@ func (gc *groupCostController) modifyTokenCounter(counter *tokenCounter, bucket
 	counter.inDegradedMode = false
 	var cfg tokenBucketReconfigureArgs
 	cfg.NewBurst = bucket.GetSettings().GetBurstLimit()
-	// when trickleTimeMs equals zero, server has enough tokens and does not need to
+	// When trickleTimeMs equals zero, server has enough tokens and does not need to
 	// limit client consume token. So all token is granted to client right now.
 	if trickleTimeMs == 0 {
 		cfg.NewTokens = granted
 		cfg.NewRate = float64(bucket.GetSettings().FillRate)
 		counter.lastDeadline = time.Time{}
-		cfg.NotifyThreshold = math.Min((granted+counter.limiter.AvailableTokens(gc.run.now)), counter.avgRUPerSec*float64(defaultTargetPeriod)) * notifyFraction
+		cfg.NotifyThreshold = math.Min(granted+counter.limiter.AvailableTokens(gc.run.now), counter.avgRUPerSec*float64(defaultTargetPeriod)) * notifyFraction
 		// In the non-trickle case, clients can be allowed to accumulate more tokens.
 		if cfg.NewBurst >= 0 {
 			cfg.NewBurst = 0

client/resource_group/controller/limiter.go

@@ -64,7 +64,7 @@ func Every(interval time.Duration) Limit {
 //
 // Some changes about burst(b):
 //   - If b == 0, that means the limiter is unlimited capacity. default use in resource controller (burst with a rate within an unlimited capacity).
-//   - If b < 0, that means the limiter is unlimited capacity and r is ignored, can be seen as r == Inf (burst within a unlimited capacity).
+//   - If b < 0, that means the limiter is unlimited capacity and r is ignored, can be seen as r == Inf (burst within an unlimited capacity).
 //   - If b > 0, that means the limiter is limited capacity.
 type Limiter struct {
 	mu     sync.Mutex

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

@@ -116,8 +116,7 @@ func (s *Service) postResourceGroup(c *gin.Context) {
 		c.String(http.StatusBadRequest, err.Error())
 		return
 	}
-	nGroup := rmserver.FromProtoResourceGroup(&group)
-	if err := s.manager.AddResourceGroup(nGroup); err != nil {
+	if err := s.manager.AddResourceGroup(&group); err != nil {
 		c.String(http.StatusInternalServerError, err.Error())
 		return
 	}

pkg/mcs/resource_manager/server/grpc_service.go

@@ -126,8 +126,7 @@ func (s *Service) AddResourceGroup(ctx context.Context, req *rmpb.PutResourceGro
 	if err := s.checkServing(); err != nil {
 		return nil, err
 	}
-	rg := FromProtoResourceGroup(req.GetGroup())
-	err := s.manager.AddResourceGroup(rg)
+	err := s.manager.AddResourceGroup(req.GetGroup())
 	if err != nil {
 		return nil, err
 	}
@@ -180,6 +179,7 @@ func (s *Service) AcquireTokenBuckets(stream rmpb.ResourceManager_AcquireTokenBu
 			return err
 		}
 		targetPeriodMs := request.GetTargetRequestPeriodMs()
+		clientUniqueID := request.GetClientUniqueId()
 		resps := &rmpb.TokenBucketsResponse{}
 		for _, req := range request.Requests {
 			resourceGroupName := req.GetResourceGroupName()
@@ -203,7 +203,7 @@ 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)
+						tokens = rg.RequestRU(now, re.Value, targetPeriodMs, clientUniqueID)
 					}
 					if tokens == nil {
 						continue

pkg/mcs/resource_manager/server/manager.go

@@ -131,17 +131,22 @@ func (m *Manager) Init(ctx context.Context) {
 }
 
 // AddResourceGroup puts a resource group.
-func (m *Manager) AddResourceGroup(group *ResourceGroup) error {
+func (m *Manager) AddResourceGroup(grouppb *rmpb.ResourceGroup) error {
+	// Check the name.
+	if len(grouppb.Name) == 0 || len(grouppb.Name) > 32 {
+		return errors.New("invalid resource group name, the length should be in [1,32]")
+	}
+	// Check the Priority.
+	if grouppb.GetPriority() > 16 {
+		return errors.New("invalid resource group priority, the value should be in [0,16]")
+	}
 	m.RLock()
-	_, ok := m.groups[group.Name]
+	_, ok := m.groups[grouppb.Name]
 	m.RUnlock()
 	if ok {
 		return errors.New("this group already exists")
 	}
-	err := group.CheckAndInit()
-	if err != nil {
-		return err
-	}
+	group := FromProtoResourceGroup(grouppb)
 	m.Lock()
 	defer m.Unlock()
 	if err := group.persistSettings(m.storage); err != nil {
@@ -247,7 +252,9 @@ func (m *Manager) persistResourceGroupRunningState() {
 		group, ok := m.groups[keys[idx]]
 		m.RUnlock()
 		if ok {
+			m.Lock()
 			group.persistStates(m.storage)
+			m.Unlock()
 		}
 	}
 }

pkg/mcs/resource_manager/server/resource_group.go

@@ -34,14 +34,12 @@ type ResourceGroup struct {
 	Mode rmpb.GroupMode `json:"mode"`
 	// RU settings
 	RUSettings *RequestUnitSettings `json:"r_u_settings,omitempty"`
-	// raw resource settings
-	RawResourceSettings *RawResourceSettings `json:"raw_resource_settings,omitempty"`
-	Priority            uint32               `json:"priority"`
+	Priority   uint32               `json:"priority"`
 }
 
 // RequestUnitSettings is the definition of the RU settings.
 type RequestUnitSettings struct {
-	RU GroupTokenBucket `json:"r_u,omitempty"`
+	RU *GroupTokenBucket `json:"r_u,omitempty"`
 }
 
 // NewRequestUnitSettings creates a new RequestUnitSettings with the given token bucket.
@@ -51,22 +49,6 @@ func NewRequestUnitSettings(tokenBucket *rmpb.TokenBucket) *RequestUnitSettings
 	}
 }
 
-// RawResourceSettings is the definition of the native resource settings.
-type RawResourceSettings struct {
-	CPU              GroupTokenBucket `json:"cpu,omitempty"`
-	IOReadBandwidth  GroupTokenBucket `json:"io_read_bandwidth,omitempty"`
-	IOWriteBandwidth GroupTokenBucket `json:"io_write_bandwidth,omitempty"`
-}
-
-// NewRawResourceSettings creates a new RawResourceSettings with the given token buckets.
-func NewRawResourceSettings(cpu, ioRead, ioWrite *rmpb.TokenBucket) *RawResourceSettings {
-	return &RawResourceSettings{
-		CPU:              NewGroupTokenBucket(cpu),
-		IOReadBandwidth:  NewGroupTokenBucket(ioRead),
-		IOWriteBandwidth: NewGroupTokenBucket(ioWrite),
-	}
-}
-
 func (rg *ResourceGroup) String() string {
 	res, err := json.Marshal(rg)
 	if err != nil {
@@ -93,42 +75,13 @@ func (rg *ResourceGroup) Copy() *ResourceGroup {
 	return &newRG
 }
 
-// CheckAndInit checks the validity of the resource group and initializes the default values if not setting.
-// Only used to initialize the resource group when creating.
-func (rg *ResourceGroup) CheckAndInit() error {
-	if len(rg.Name) == 0 || len(rg.Name) > 32 {
-		return errors.New("invalid resource group name, the length should be in [1,32]")
-	}
-	if rg.Priority > 16 {
-		return errors.New("invalid resource group priority, the value should be in [0,16]")
-	}
-	switch rg.Mode {
-	case rmpb.GroupMode_RUMode:
-		if rg.RUSettings == nil {
-			rg.RUSettings = NewRequestUnitSettings(nil)
-		}
-		if rg.RawResourceSettings != nil {
-			return errors.New("invalid resource group settings, RU mode should not set raw resource settings")
-		}
-	case rmpb.GroupMode_RawMode:
-		if rg.RawResourceSettings == nil {
-			rg.RawResourceSettings = NewRawResourceSettings(nil, nil, nil)
-		}
-		if rg.RUSettings != nil {
-			return errors.New("invalid resource group settings, raw mode should not set RU settings")
-		}
-	default:
-		return errors.New("invalid resource group mode")
-	}
-	return nil
-}
-
 // PatchSettings patches the resource group settings.
 // Only used to patch the resource group when updating.
 // Note: the tokens is the delta value to patch.
 func (rg *ResourceGroup) PatchSettings(metaGroup *rmpb.ResourceGroup) error {
 	rg.Lock()
 	defer rg.Unlock()
+
 	if metaGroup.GetMode() != rg.Mode {
 		return errors.New("only support reconfigure in same mode, maybe you should delete and create a new one")
 	}
@@ -138,17 +91,14 @@ func (rg *ResourceGroup) PatchSettings(metaGroup *rmpb.ResourceGroup) error {
 	rg.Priority = metaGroup.Priority
 	switch rg.Mode {
 	case rmpb.GroupMode_RUMode:
-		if metaGroup.GetRUSettings() == nil {
+		settings := metaGroup.GetRUSettings()
+		if settings == nil {
 			return errors.New("invalid resource group settings, RU mode should set RU settings")
 		}
-		rg.RUSettings.RU.patch(metaGroup.GetRUSettings().GetRU())
+		rg.RUSettings.RU.patch(settings.GetRU())
+		log.Info("patch resource group ru settings", zap.String("name", rg.Name), zap.Any("settings", settings))
 	case rmpb.GroupMode_RawMode:
-		if metaGroup.GetRawResourceSettings() == nil {
-			return errors.New("invalid resource group settings, raw mode should set resource settings")
-		}
-		rg.RawResourceSettings.CPU.patch(metaGroup.GetRawResourceSettings().GetCpu())
-		rg.RawResourceSettings.IOReadBandwidth.patch(metaGroup.GetRawResourceSettings().GetIoRead())
-		rg.RawResourceSettings.IOWriteBandwidth.patch(metaGroup.GetRawResourceSettings().GetIoWrite())
+		panic("no implementation")
 	}
 	log.Info("patch resource group settings", zap.String("name", rg.Name), zap.String("settings", rg.String()))
 	return nil
@@ -163,17 +113,13 @@ func FromProtoResourceGroup(group *rmpb.ResourceGroup) *ResourceGroup {
 	}
 	switch group.GetMode() {
 	case rmpb.GroupMode_RUMode:
-		if settings := group.GetRUSettings(); settings != nil {
-			rg.RUSettings = NewRequestUnitSettings(settings.GetRU())
+		if group.GetRUSettings() == nil {
+			rg.RUSettings = NewRequestUnitSettings(nil)
+		} else {
+			rg.RUSettings = NewRequestUnitSettings(group.GetRUSettings().GetRU())
 		}
 	case rmpb.GroupMode_RawMode:
-		if settings := group.GetRawResourceSettings(); settings != nil {
-			rg.RawResourceSettings = NewRawResourceSettings(
-				settings.GetCpu(),
-				settings.GetIoRead(),
-				settings.GetIoWrite(),
-			)
-		}
+		panic("no implementation")
 	}
 	return rg
 }
@@ -182,21 +128,23 @@ func FromProtoResourceGroup(group *rmpb.ResourceGroup) *ResourceGroup {
 func (rg *ResourceGroup) RequestRU(
 	now time.Time,
 	neededTokens float64,
-	targetPeriodMs uint64,
+	targetPeriodMs, clientUniqueID uint64,
 ) *rmpb.GrantedRUTokenBucket {
 	rg.Lock()
 	defer rg.Unlock()
+
 	if rg.RUSettings == nil || rg.RUSettings.RU.Settings == nil {
 		return nil
 	}
-	tb, trickleTimeMs := rg.RUSettings.RU.request(now, neededTokens, targetPeriodMs)
+	tb, trickleTimeMs := rg.RUSettings.RU.request(now, neededTokens, targetPeriodMs, clientUniqueID)
 	return &rmpb.GrantedRUTokenBucket{GrantedTokens: tb, TrickleTimeMs: trickleTimeMs}
 }
 
 // IntoProtoResourceGroup converts a ResourceGroup to a rmpb.ResourceGroup.
 func (rg *ResourceGroup) IntoProtoResourceGroup() *rmpb.ResourceGroup {
 	rg.RLock()
 	defer rg.RUnlock()
+
 	switch rg.Mode {
 	case rmpb.GroupMode_RUMode: // RU mode
 		group := &rmpb.ResourceGroup{
@@ -209,17 +157,7 @@ func (rg *ResourceGroup) IntoProtoResourceGroup() *rmpb.ResourceGroup {
 		}
 		return group
 	case rmpb.GroupMode_RawMode: // Raw mode
-		group := &rmpb.ResourceGroup{
-			Name:     rg.Name,
-			Mode:     rmpb.GroupMode_RawMode,
-			Priority: rg.Priority,
-			RawResourceSettings: &rmpb.GroupRawResourceSettings{
-				Cpu:     rg.RawResourceSettings.CPU.GetTokenBucket(),
-				IoRead:  rg.RawResourceSettings.IOReadBandwidth.GetTokenBucket(),
-				IoWrite: rg.RawResourceSettings.IOWriteBandwidth.GetTokenBucket(),
-			},
-		}
-		return group
+		panic("no implementation")
 	}
 	return nil
 }
@@ -245,19 +183,15 @@ type GroupStates struct {
 func (rg *ResourceGroup) GetGroupStates() *GroupStates {
 	rg.RLock()
 	defer rg.RUnlock()
+
 	switch rg.Mode {
 	case rmpb.GroupMode_RUMode: // RU mode
 		tokens := &GroupStates{
 			RU: rg.RUSettings.RU.GroupTokenBucketState.Clone(),
 		}
 		return tokens
 	case rmpb.GroupMode_RawMode: // Raw mode
-		tokens := &GroupStates{
-			CPU:     rg.RawResourceSettings.CPU.GroupTokenBucketState.Clone(),
-			IORead:  rg.RawResourceSettings.IOReadBandwidth.GroupTokenBucketState.Clone(),
-			IOWrite: rg.RawResourceSettings.IOWriteBandwidth.GroupTokenBucketState.Clone(),
-		}
-		return tokens
+		panic("no implementation")
 	}
 	return nil
 }
@@ -267,18 +201,11 @@ func (rg *ResourceGroup) SetStatesIntoResourceGroup(states *GroupStates) {
 	switch rg.Mode {
 	case rmpb.GroupMode_RUMode:
 		if state := states.RU; state != nil {
-			rg.RUSettings.RU.GroupTokenBucketState = *state
+			rg.RUSettings.RU.setState(state)
+			log.Debug("update group token bucket state", zap.String("name", rg.Name), zap.Any("state", state))
 		}
 	case rmpb.GroupMode_RawMode:
-		if state := states.CPU; state != nil {
-			rg.RawResourceSettings.CPU.GroupTokenBucketState = *state
-		}
-		if state := states.IORead; state != nil {
-			rg.RawResourceSettings.IOReadBandwidth.GroupTokenBucketState = *state
-		}
-		if state := states.IOWrite; state != nil {
-			rg.RawResourceSettings.IOWriteBandwidth.GroupTokenBucketState = *state
-		}
+		panic("no implementation")
 	}
 }
 

pkg/mcs/resource_manager/server/resource_group_test.go

@@ -10,9 +10,7 @@ import (
 
 func TestPatchResourceGroup(t *testing.T) {
 	re := require.New(t)
-	rg1 := &ResourceGroup{Name: "test", Mode: rmpb.GroupMode_RUMode, RUSettings: &RequestUnitSettings{}}
-	err := rg1.CheckAndInit()
-	re.NoError(err)
+	rg := &ResourceGroup{Name: "test", Mode: rmpb.GroupMode_RUMode, RUSettings: NewRequestUnitSettings(nil)}
 	testCaseRU := []struct {
 		patchJSONString  string
 		expectJSONString string
@@ -26,40 +24,12 @@ func TestPatchResourceGroup(t *testing.T) {
 	}
 
 	for _, ca := range testCaseRU {
-		rg := rg1.Copy()
 		patch := &rmpb.ResourceGroup{}
 		err := json.Unmarshal([]byte(ca.patchJSONString), patch)
 		re.NoError(err)
 		err = rg.PatchSettings(patch)
 		re.NoError(err)
-		res, err := json.Marshal(rg)
-		re.NoError(err)
-		re.Equal(ca.expectJSONString, string(res))
-	}
-
-	rg2 := &ResourceGroup{Name: "test", Mode: rmpb.GroupMode_RawMode, RawResourceSettings: &RawResourceSettings{}}
-	err = rg2.CheckAndInit()
-	re.NoError(err)
-	testCaseResource := []struct {
-		patchJSONString  string
-		expectJSONString string
-	}{
-		{`{"name":"test", "mode":2, "raw_resource_settings": {"cpu":{"settings":{"fill_rate": 200000}}}}`,
-			`{"name":"test","mode":2,"raw_resource_settings":{"cpu":{"settings":{"fill_rate":200000},"state":{"initialized":false}},"io_read_bandwidth":{"state":{"initialized":false}},"io_write_bandwidth":{"state":{"initialized":false}}},"priority":0}`},
-		{`{"name":"test", "mode":2, "raw_resource_settings": {"io_read":{"settings":{"fill_rate": 200000,"burst_limit":1000000}}}}`,
-			`{"name":"test","mode":2,"raw_resource_settings":{"cpu":{"state":{"initialized":false}},"io_read_bandwidth":{"settings":{"fill_rate":200000,"burst_limit":1000000},"state":{"initialized":false}},"io_write_bandwidth":{"state":{"initialized":false}}},"priority":0}`},
-		{`{"name":"test", "mode":2, "raw_resource_settings": {"io_write":{"settings":{"fill_rate": 200000}}}, "priority": 16 }`,
-			`{"name":"test","mode":2,"raw_resource_settings":{"cpu":{"state":{"initialized":false}},"io_read_bandwidth":{"state":{"initialized":false}},"io_write_bandwidth":{"settings":{"fill_rate":200000},"state":{"initialized":false}}},"priority":16}`},
-	}
-
-	for _, ca := range testCaseResource {
-		rg := rg2.Copy()
-		patch := &rmpb.ResourceGroup{}
-		err := json.Unmarshal([]byte(ca.patchJSONString), patch)
-		re.NoError(err)
-		err = rg.PatchSettings(patch)
-		re.NoError(err)
-		res, err := json.Marshal(rg)
+		res, err := json.Marshal(rg.Copy())
 		re.NoError(err)
 		re.Equal(ca.expectJSONString, string(res))
 	}