resource_mananger: deep clone resource group (#7623) (#7624)

close #7206 resource_mananger: deep clone resource group Signed-off-by: ti-chi-bot <[email protected]> Signed-off-by: nolouch <[email protected]> Co-authored-by: ShuNing <[email protected]> Co-authored-by: nolouch <[email protected]>
tikv · Jan 3, 2024 · f3a948a · f3a948a
1 parent f36f49b
commit f3a948a
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 18 deletions.
diff --git a/pkg/mcs/resource_manager/server/manager.go b/pkg/mcs/resource_manager/server/manager.go
@@ -278,7 +278,7 @@ func (m *Manager) GetResourceGroup(name string) *ResourceGroup {
 	m.RLock()
 	defer m.RUnlock()
 	if group, ok := m.groups[name]; ok {
-		return group.Copy()
+		return group.Clone()
 	}
 	return nil
 }
@@ -298,7 +298,7 @@ func (m *Manager) GetResourceGroupList() []*ResourceGroup {
 	m.RLock()
 	res := make([]*ResourceGroup, 0, len(m.groups))
 	for _, group := range m.groups {
-		res = append(res, group.Copy())
+		res = append(res, group.Clone())
 	}
 	m.RUnlock()
 	sort.Slice(res, func(i, j int) bool {

diff --git a/pkg/mcs/resource_manager/server/resource_group.go b/pkg/mcs/resource_manager/server/resource_group.go
@@ -42,6 +42,20 @@ type RequestUnitSettings struct {
 	RU *GroupTokenBucket `json:"r_u,omitempty"`
 }
 
+// Clone returns a deep copy of the RequestUnitSettings.
+func (rus *RequestUnitSettings) Clone() *RequestUnitSettings {
+	if rus == nil {
+		return nil
+	}
+	var ru *GroupTokenBucket
+	if rus.RU != nil {
+		ru = rus.RU.Clone()
+	}
+	return &RequestUnitSettings{
+		RU: ru,
+	}
+}
+
 // NewRequestUnitSettings creates a new RequestUnitSettings with the given token bucket.
 func NewRequestUnitSettings(tokenBucket *rmpb.TokenBucket) *RequestUnitSettings {
 	return &RequestUnitSettings{
@@ -58,21 +72,17 @@ func (rg *ResourceGroup) String() string {
 	return string(res)
 }
 
-// Copy copies the resource group.
-func (rg *ResourceGroup) Copy() *ResourceGroup {
-	// TODO: use a better way to copy
+// Clone copies the resource group.
+func (rg *ResourceGroup) Clone() *ResourceGroup {
 	rg.RLock()
 	defer rg.RUnlock()
-	res, err := json.Marshal(rg)
-	if err != nil {
-		panic(err)
-	}
-	var newRG ResourceGroup
-	err = json.Unmarshal(res, &newRG)
-	if err != nil {
-		panic(err)
+	newRG := &ResourceGroup{
+		Name:       rg.Name,
+		Mode:       rg.Mode,
+		Priority:   rg.Priority,
+		RUSettings: rg.RUSettings.Clone(),
 	}
-	return &newRG
+	return newRG
 }
 
 func (rg *ResourceGroup) getRUToken() float64 {

diff --git a/pkg/mcs/resource_manager/server/resource_group_test.go b/pkg/mcs/resource_manager/server/resource_group_test.go
@@ -29,7 +29,7 @@ func TestPatchResourceGroup(t *testing.T) {
 		re.NoError(err)
 		err = rg.PatchSettings(patch)
 		re.NoError(err)
-		res, err := json.Marshal(rg.Copy())
+		res, err := json.Marshal(rg.Clone())
 		re.NoError(err)
 		re.Equal(ca.expectJSONString, string(res))
 	}

diff --git a/pkg/mcs/resource_manager/server/token_bukets.go b/pkg/mcs/resource_manager/server/token_bukets.go
@@ -49,6 +49,22 @@ type GroupTokenBucket struct {
 	GroupTokenBucketState `json:"state,omitempty"`
 }
 
+// Clone returns the deep copy of GroupTokenBucket
+func (gtb *GroupTokenBucket) Clone() *GroupTokenBucket {
+	if gtb == nil {
+		return nil
+	}
+	var settings *rmpb.TokenLimitSettings
+	if gtb.Settings != nil {
+		settings = proto.Clone(gtb.Settings).(*rmpb.TokenLimitSettings)
+	}
+	stateClone := *gtb.GroupTokenBucketState.Clone()
+	return &GroupTokenBucket{
+		Settings:              settings,
+		GroupTokenBucketState: stateClone,
+	}
+}
+
 func (gtb *GroupTokenBucket) setState(state *GroupTokenBucketState) {
 	gtb.Tokens = state.Tokens
 	gtb.LastUpdate = state.LastUpdate
@@ -85,10 +101,14 @@ type GroupTokenBucketState struct {
 
 // Clone returns the copy of GroupTokenBucketState
 func (gts *GroupTokenBucketState) Clone() *GroupTokenBucketState {
-	tokenSlots := make(map[uint64]*TokenSlot)
-	for id, tokens := range gts.tokenSlots {
-		tokenSlots[id] = tokens
+	var tokenSlots map[uint64]*TokenSlot
+	if gts.tokenSlots != nil {
+		tokenSlots = make(map[uint64]*TokenSlot)
+		for id, tokens := range gts.tokenSlots {
+			tokenSlots[id] = tokens
+		}
 	}
+
 	var lastUpdate *time.Time
 	if gts.LastUpdate != nil {
 		newLastUpdate := *gts.LastUpdate