Skip to content

Commit

Permalink
resource_mananger: deep clone resource group
Browse files Browse the repository at this point in the history
Signed-off-by: nolouch <nolouch@gmail.com>
  • Loading branch information
nolouch committed Dec 27, 2023
1 parent d3551ea commit d12c300
Show file tree
Hide file tree
Showing 9 changed files with 104 additions and 19 deletions.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
4 changes: 2 additions & 2 deletions pkg/mcs/resourcemanager/server/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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 {
Expand Down
45 changes: 32 additions & 13 deletions pkg/mcs/resourcemanager/server/resource_group.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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{
Expand All @@ -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 {
Expand Down
40 changes: 39 additions & 1 deletion pkg/mcs/resourcemanager/server/resource_group_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand All @@ -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)
}
}
25 changes: 22 additions & 3 deletions pkg/mcs/resourcemanager/server/token_buckets.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions tests/integrations/client/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
2 changes: 2 additions & 0 deletions tests/integrations/mcs/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
2 changes: 2 additions & 0 deletions tests/integrations/tso/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down

0 comments on commit d12c300

Please sign in to comment.