resource_mananger: deep clone resource group (tikv#7623)

close tikv#7206 resource_mananger: deep clone resource group Signed-off-by: nolouch <[email protected]> Co-authored-by: tongjian <[email protected]>
HuSharp · Dec 27, 2023 · ed9685a · ed9685a
1 parent b35b0cd
commit ed9685a
Show file tree

Hide file tree

Showing 9 changed files with 106 additions and 19 deletions.
diff --git a/go.mod b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/kms v1.20.8
 	github.com/aws/aws-sdk-go-v2/service/sts v1.18.7
 	github.com/axw/gocov v1.0.0
+	github.com/brianvoe/gofakeit/v6 v6.26.3
 	github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5
 	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e
 	github.com/coreos/go-semver v0.3.0

diff --git a/go.sum b/go.sum
@@ -72,6 +72,8 @@ github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4Yn
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch h1:KLE/YeX+9FNaGVW5MtImRVPhjDpfpgJhvkuYWBmOYbo=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch/go.mod h1:KjBLriHXe7L6fGceqWzTod8HUB/TP1WWDtfuSYtYXaI=
+github.com/brianvoe/gofakeit/v6 v6.26.3 h1:3ljYrjPwsUNAUFdUIr2jVg5EhKdcke/ZLop7uVg1Er8=
+github.com/brianvoe/gofakeit/v6 v6.26.3/go.mod h1:Xj58BMSnFqcn/fAQeSK+/PLtC5kSb7FJIq4JyGa8vEs=
 github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
 github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
 github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=

diff --git a/pkg/mcs/resourcemanager/server/manager.go b/pkg/mcs/resourcemanager/server/manager.go
@@ -282,7 +282,7 @@ func (m *Manager) GetResourceGroup(name string, withStats bool) *ResourceGroup {
 	m.RLock()
 	defer m.RUnlock()
 	if group, ok := m.groups[name]; ok {
-		return group.Copy(withStats)
+		return group.Clone(withStats)
 	}
 	return nil
 }
@@ -302,7 +302,7 @@ func (m *Manager) GetResourceGroupList(withStats bool) []*ResourceGroup {
 	m.RLock()
 	res := make([]*ResourceGroup, 0, len(m.groups))
 	for _, group := range m.groups {
-		res = append(res, group.Copy(withStats))
+		res = append(res, group.Clone(withStats))
 	}
 	m.RUnlock()
 	sort.Slice(res, func(i, j int) bool {

diff --git a/pkg/mcs/resourcemanager/server/resource_group.go b/pkg/mcs/resourcemanager/server/resource_group.go
@@ -19,6 +19,7 @@ import (
 	"encoding/json"
 	"time"
 
+	"github.com/gogo/protobuf/proto"
 	"github.com/pingcap/errors"
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
 	"github.com/pingcap/log"
@@ -46,6 +47,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{
@@ -62,24 +77,29 @@ func (rg *ResourceGroup) String() string {
 	return string(res)
 }
 
-// Copy copies the resource group.
-func (rg *ResourceGroup) Copy(withStats bool) *ResourceGroup {
-	// TODO: use a better way to copy
+// Clone copies the resource group.
+func (rg *ResourceGroup) Clone(withStats bool) *ResourceGroup {
 	rg.RLock()
 	defer rg.RUnlock()
-	res, err := json.Marshal(rg)
-	if err != nil {
-		panic(err)
+	newRG := &ResourceGroup{
+		Name:       rg.Name,
+		Mode:       rg.Mode,
+		Priority:   rg.Priority,
+		RUSettings: rg.RUSettings.Clone(),
 	}
-	var newRG ResourceGroup
-	err = json.Unmarshal(res, &newRG)
-	if err != nil {
-		panic(err)
+	if rg.Runaway != nil {
+		newRG.Runaway = proto.Clone(rg.Runaway).(*rmpb.RunawaySettings)
 	}
-	if !withStats {
-		newRG.RUConsumption = nil
+
+	if rg.Background != nil {
+		newRG.Background = proto.Clone(rg.Background).(*rmpb.BackgroundSettings)
 	}
-	return &newRG
+
+	if withStats && rg.RUConsumption != nil {
+		newRG.RUConsumption = proto.Clone(rg.RUConsumption).(*rmpb.Consumption)
+	}
+
+	return newRG
 }
 
 func (rg *ResourceGroup) getRUToken() float64 {

diff --git a/pkg/mcs/resourcemanager/server/resource_group_test.go b/pkg/mcs/resourcemanager/server/resource_group_test.go
@@ -2,8 +2,10 @@ package server
 
 import (
 	"encoding/json"
+	"reflect"
 	"testing"
 
+	"github.com/brianvoe/gofakeit/v6"
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
 	"github.com/stretchr/testify/require"
 )
@@ -29,8 +31,44 @@ func TestPatchResourceGroup(t *testing.T) {
 		re.NoError(err)
 		err = rg.PatchSettings(patch)
 		re.NoError(err)
-		res, err := json.Marshal(rg.Copy(false))
+		res, err := json.Marshal(rg.Clone(false))
 		re.NoError(err)
 		re.Equal(ca.expectJSONString, string(res))
 	}
 }
+
+func resetSizeCache(obj interface{}) {
+	resetSizeCacheRecursive(reflect.ValueOf(obj))
+}
+
+func resetSizeCacheRecursive(value reflect.Value) {
+	if value.Kind() == reflect.Ptr {
+		value = value.Elem()
+	}
+
+	if value.Kind() != reflect.Struct {
+		return
+	}
+
+	for i := 0; i < value.NumField(); i++ {
+		fieldValue := value.Field(i)
+		fieldType := value.Type().Field(i)
+
+		if fieldType.Name == "XXX_sizecache" && fieldType.Type.Kind() == reflect.Int32 {
+			fieldValue.SetInt(0)
+		} else {
+			resetSizeCacheRecursive(fieldValue)
+		}
+	}
+}
+
+func TestClone(t *testing.T) {
+	for i := 0; i <= 10; i++ {
+		var rg ResourceGroup
+		gofakeit.Struct(&rg)
+		// hack to reset XXX_sizecache, gofakeit will random set this field but proto clone will not copy this field.
+		resetSizeCache(&rg)
+		rgClone := rg.Clone(true)
+		require.EqualValues(t, &rg, rgClone)
+	}
+}
diff --git a/pkg/mcs/resourcemanager/server/token_buckets.go b/pkg/mcs/resourcemanager/server/token_buckets.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

diff --git a/tests/integrations/client/go.sum b/tests/integrations/client/go.sum
@@ -68,6 +68,8 @@ github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4Yn
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch h1:KLE/YeX+9FNaGVW5MtImRVPhjDpfpgJhvkuYWBmOYbo=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch/go.mod h1:KjBLriHXe7L6fGceqWzTod8HUB/TP1WWDtfuSYtYXaI=
+github.com/brianvoe/gofakeit/v6 v6.26.3 h1:3ljYrjPwsUNAUFdUIr2jVg5EhKdcke/ZLop7uVg1Er8=
+github.com/brianvoe/gofakeit/v6 v6.26.3/go.mod h1:Xj58BMSnFqcn/fAQeSK+/PLtC5kSb7FJIq4JyGa8vEs=
 github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
 github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
 github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=

diff --git a/tests/integrations/mcs/go.sum b/tests/integrations/mcs/go.sum
@@ -68,6 +68,8 @@ github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4Yn
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch h1:KLE/YeX+9FNaGVW5MtImRVPhjDpfpgJhvkuYWBmOYbo=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch/go.mod h1:KjBLriHXe7L6fGceqWzTod8HUB/TP1WWDtfuSYtYXaI=
+github.com/brianvoe/gofakeit/v6 v6.26.3 h1:3ljYrjPwsUNAUFdUIr2jVg5EhKdcke/ZLop7uVg1Er8=
+github.com/brianvoe/gofakeit/v6 v6.26.3/go.mod h1:Xj58BMSnFqcn/fAQeSK+/PLtC5kSb7FJIq4JyGa8vEs=
 github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
 github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
 github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=

diff --git a/tests/integrations/tso/go.sum b/tests/integrations/tso/go.sum
@@ -68,6 +68,8 @@ github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4Yn
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch h1:KLE/YeX+9FNaGVW5MtImRVPhjDpfpgJhvkuYWBmOYbo=
 github.com/breeswish/gin-jwt/v2 v2.6.4-jwt-patch/go.mod h1:KjBLriHXe7L6fGceqWzTod8HUB/TP1WWDtfuSYtYXaI=
+github.com/brianvoe/gofakeit/v6 v6.26.3 h1:3ljYrjPwsUNAUFdUIr2jVg5EhKdcke/ZLop7uVg1Er8=
+github.com/brianvoe/gofakeit/v6 v6.26.3/go.mod h1:Xj58BMSnFqcn/fAQeSK+/PLtC5kSb7FJIq4JyGa8vEs=
 github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
 github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
 github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=