Skip to content

Commit

Permalink
controller: avoid the invalid nil pointer (#6486)
Browse files Browse the repository at this point in the history 
close #6485, ref #6485

controller: avoid the invalid nil pointer

Signed-off-by: nolouch <nolouch@gmail.com>
  • Loading branch information
@nolouch
nolouch authored May 19, 2023
1 parent 76cddc7 commit 1613e22
Show file tree
Hide file tree
Showing 2 changed files with 256 additions and 8 deletions.
46 changes: 38 additions & 8 deletions client/resource_group/controller/model.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,18 +112,27 @@ func (kc *KVCalculator) AfterKVRequest(consumption *rmpb.Consumption, req Reques
}

func (kc *KVCalculator) calculateReadCost(consumption *rmpb.Consumption, res ResponseInfo) {
if consumption == nil {
return
}
readBytes := float64(res.ReadBytes())
consumption.ReadBytes += readBytes
consumption.RRU += float64(kc.ReadBytesCost) * readBytes
}

func (kc *KVCalculator) calculateCPUCost(consumption *rmpb.Consumption, res ResponseInfo) {
if consumption == nil {
return
}
kvCPUMs := float64(res.KVCPU().Nanoseconds()) / 1000000.0
consumption.TotalCpuTimeMs += kvCPUMs
consumption.RRU += float64(kc.CPUMsCost) * kvCPUMs
}

func (kc *KVCalculator) payBackWriteCost(consumption *rmpb.Consumption, req RequestInfo) {
if consumption == nil {
return
}
writeBytes := float64(req.WriteBytes())
consumption.WriteBytes -= writeBytes
consumption.WRU -= float64(kc.WriteBaseCost) + float64(kc.WriteBytesCost)*writeBytes
Expand All @@ -142,6 +151,9 @@ func newSQLCalculator(cfg *Config) *SQLCalculator {

// Trickle update sql layer CPU consumption.
func (dsc *SQLCalculator) Trickle(consumption *rmpb.Consumption) {
if consumption == nil {
return
}
delta := getSQLProcessCPUTime(dsc.isSingleGroupByKeyspace) - consumption.SqlLayerCpuTimeMs
consumption.TotalCpuTimeMs += delta
consumption.SqlLayerCpuTimeMs += delta
Expand All @@ -156,44 +168,59 @@ func (dsc *SQLCalculator) AfterKVRequest(consumption *rmpb.Consumption, req Requ
}

func getRUValueFromConsumption(custom *rmpb.Consumption, typ rmpb.RequestUnitType) float64 {
if typ == 0 {
if custom == nil {
return 0
}
if typ == rmpb.RequestUnitType_RU {
return custom.RRU + custom.WRU
}
return 0
}

func getRUTokenBucketSetting(group *rmpb.ResourceGroup, typ rmpb.RequestUnitType) *rmpb.TokenBucket {
if typ == 0 {
if group == nil {
return nil
}
if typ == rmpb.RequestUnitType_RU {
return group.RUSettings.RU
}
return nil
}

func getRawResourceValueFromConsumption(custom *rmpb.Consumption, typ rmpb.RawResourceType) float64 {
if custom == nil {
return 0
}
switch typ {
case 0:
case rmpb.RawResourceType_CPU:
return custom.TotalCpuTimeMs
case 1:
case rmpb.RawResourceType_IOReadFlow:
return custom.ReadBytes
case 2:
case rmpb.RawResourceType_IOWriteFlow:
return custom.WriteBytes
}
return 0
}

func getRawResourceTokenBucketSetting(group *rmpb.ResourceGroup, typ rmpb.RawResourceType) *rmpb.TokenBucket {
if group == nil {
return nil
}
switch typ {
case 0:
case rmpb.RawResourceType_CPU:
return group.RawResourceSettings.Cpu
case 1:
case rmpb.RawResourceType_IOReadFlow:
return group.RawResourceSettings.IoRead
case 2:
case rmpb.RawResourceType_IOWriteFlow:
return group.RawResourceSettings.IoWrite
}
return nil
}

func add(custom1 *rmpb.Consumption, custom2 *rmpb.Consumption) {
if custom1 == nil || custom2 == nil {
return
}
custom1.RRU += custom2.RRU
custom1.WRU += custom2.WRU
custom1.ReadBytes += custom2.ReadBytes
Expand All @@ -205,6 +232,9 @@ func add(custom1 *rmpb.Consumption, custom2 *rmpb.Consumption) {
}

func sub(custom1 *rmpb.Consumption, custom2 *rmpb.Consumption) {
if custom1 == nil || custom2 == nil {
return
}
custom1.RRU -= custom2.RRU
custom1.WRU -= custom2.WRU
custom1.ReadBytes -= custom2.ReadBytes
Expand Down
218 changes: 218 additions & 0 deletions client/resource_group/controller/model_test.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,218 @@
// Copyright 2023 TiKV Project Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package controller

import (
"testing"

rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
"github.com/stretchr/testify/require"
)

func TestGetRUValueFromConsumption(t *testing.T) {
// Positive test case
re := require.New(t)
custom := &rmpb.Consumption{RRU: 2.5, WRU: 3.5}
typ := rmpb.RequestUnitType_RU
expected := float64(6)

result := getRUValueFromConsumption(custom, typ)
re.Equal(expected, result)

// When custom is nil
custom = nil
expected = float64(0)

result = getRUValueFromConsumption(custom, typ)
re.Equal(expected, result)

// When typ is not RU
custom = &rmpb.Consumption{RRU: 2.5, WRU: 3.5}
typ = rmpb.RequestUnitType_RU
expected = float64(0)

result = getRUValueFromConsumption(custom, typ)
re.Equal(expected, result)
}

func TestGetRUTokenBucketSetting(t *testing.T) {
// Positive test case
re := require.New(t)
group := &rmpb.ResourceGroup{
RUSettings: &rmpb.GroupRequestUnitSettings{
RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 100}},
},
}
typ := rmpb.RequestUnitType_RU
expected := &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 100}}

result := getRUTokenBucketSetting(group, typ)
re.Equal(expected.GetSettings().GetFillRate(), result.GetSettings().GetFillRate())

// When group is nil
group = nil
expected = nil

result = getRUTokenBucketSetting(group, typ)
if result != expected {
t.Errorf("Expected nil but got %v", result)
}

// When typ is not RU
group = &rmpb.ResourceGroup{
RUSettings: &rmpb.GroupRequestUnitSettings{
RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 100}},
},
}
typ = rmpb.RequestUnitType_RU
expected = nil

result = getRUTokenBucketSetting(group, typ)
re.Equal(expected, result)
}

func TestGetRawResourceValueFromConsumption(t *testing.T) {
// Positive test case
re := require.New(t)
custom := &rmpb.Consumption{TotalCpuTimeMs: 50}
typ := rmpb.RawResourceType_CPU
expected := float64(50)

result := getRawResourceValueFromConsumption(custom, typ)
re.Equal(expected, result)

// When custom is nil
custom = nil
expected = float64(0)

result = getRawResourceValueFromConsumption(custom, typ)
re.Equal(expected, result)

// When typ is IOReadFlow
custom = &rmpb.Consumption{ReadBytes: 200}
typ = rmpb.RawResourceType_IOReadFlow
expected = float64(200)

result = getRawResourceValueFromConsumption(custom, typ)
re.Equal(expected, result)
}

func TestGetRawResourceTokenBucketSetting(t *testing.T) {
// Positive test case
re := require.New(t)
group := &rmpb.ResourceGroup{
RawResourceSettings: &rmpb.GroupRawResourceSettings{
Cpu: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 100}},
},
}
typ := rmpb.RawResourceType_CPU
expected := &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 100}}

result := getRawResourceTokenBucketSetting(group, typ)

re.Equal(expected.GetSettings().GetFillRate(), result.GetSettings().GetFillRate())

// When group is nil
group = nil
expected = nil

result = getRawResourceTokenBucketSetting(group, typ)
if result != expected {
t.Errorf("Expected nil but got %v", result)
}

// When typ is IOReadFlow
group = &rmpb.ResourceGroup{
RawResourceSettings: &rmpb.GroupRawResourceSettings{
IoRead: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 200}},
},
}
typ = rmpb.RawResourceType_IOReadFlow
expected = &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 200}}

result = getRawResourceTokenBucketSetting(group, typ)
re.Equal(expected.GetSettings().GetFillRate(), result.GetSettings().GetFillRate())
}

func TestAdd(t *testing.T) {
// Positive test case
re := require.New(t)
custom1 := &rmpb.Consumption{RRU: 2.5, WRU: 3.5}
custom2 := &rmpb.Consumption{RRU: 1.5, WRU: 2.5}
expected := &rmpb.Consumption{
RRU: 4,
WRU: 6,
ReadBytes: 0,
WriteBytes: 0,
TotalCpuTimeMs: 0,
SqlLayerCpuTimeMs: 0,
KvReadRpcCount: 0,
KvWriteRpcCount: 0,
}

add(custom1, custom2)
re.Equal(expected, custom1)

// When custom1 is nil
custom1 = nil
custom2 = &rmpb.Consumption{RRU: 1.5, WRU: 2.5}
expected = nil

add(custom1, custom2)
re.Equal(expected, custom1)

// When custom2 is nil
custom1 = &rmpb.Consumption{RRU: 2.5, WRU: 3.5}
custom2 = nil
expected = &rmpb.Consumption{RRU: 2.5, WRU: 3.5}

add(custom1, custom2)
re.Equal(expected, custom1)
}

func TestSub(t *testing.T) {
// Positive test case
re := require.New(t)
custom1 := &rmpb.Consumption{RRU: 2.5, WRU: 3.5}
custom2 := &rmpb.Consumption{RRU: 1.5, WRU: 2.5}
expected := &rmpb.Consumption{
RRU: 1,
WRU: 1,
ReadBytes: 0,
WriteBytes: 0,
TotalCpuTimeMs: 0,
SqlLayerCpuTimeMs: 0,
KvReadRpcCount: 0,
KvWriteRpcCount: 0,
}

sub(custom1, custom2)
re.Equal(expected, custom1)
// When custom1 is nil
custom1 = nil
custom2 = &rmpb.Consumption{RRU: 1.5, WRU: 2.5}
expected = nil

sub(custom1, custom2)
re.Equal(expected, custom1)

// When custom2 is nil
custom1 = &rmpb.Consumption{RRU: 2.5, WRU: 3.5}
custom2 = nil
expected = &rmpb.Consumption{RRU: 2.5, WRU: 3.5}

sub(custom1, custom2)
re.Equal(expected, custom1)
}

0 comments on commit 1613e22

Please sign in to comment.