From d12c3004bb995f2483a61732bb220f3858bf7e26 Mon Sep 17 00:00:00 2001 From: nolouch Date: Wed, 27 Dec 2023 15:13:02 +0800 Subject: [PATCH] resource_mananger: deep clone resource group Signed-off-by: nolouch --- go.mod | 1 + go.sum | 2 + pkg/mcs/resourcemanager/server/manager.go | 4 +- .../resourcemanager/server/resource_group.go | 45 +++++++++++++------ .../server/resource_group_test.go | 40 ++++++++++++++++- .../resourcemanager/server/token_buckets.go | 25 +++++++++-- tests/integrations/client/go.sum | 2 + tests/integrations/mcs/go.sum | 2 + tests/integrations/tso/go.sum | 2 + 9 files changed, 104 insertions(+), 19 deletions(-) diff --git a/go.mod b/go.mod index 0def4b4b6eac..31ffccd80263 100644 --- 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 index 37354b38f157..02de15a2c2af 100644 --- 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 index a4b49b38062b..bf0ab7e9b245 100644 --- 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 index 09f8a33de9f1..122d4de2e690 100644 --- 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,19 @@ type RequestUnitSettings struct { RU *GroupTokenBucket `json:"r_u,omitempty"` } +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 +76,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 index 8df1be6dccc5..da5f5c4f0e41 100644 --- 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 index 05a93c326733..0455fc99b00d 100644 --- a/pkg/mcs/resourcemanager/server/token_buckets.go +++ b/pkg/mcs/resourcemanager/server/token_buckets.go @@ -49,6 +49,21 @@ type GroupTokenBucket struct { GroupTokenBucketState `json:"state,omitempty"` } +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 +100,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 index 9e42dfa900c9..36c0d6758ab2 100644 --- 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 index 6397ab5b017c..d7cf59ea8c8b 100644 --- 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 index a00d29415823..7070f0699606 100644 --- 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=