From 1613e229103d1665fe3b6358ca52390a8075ff36 Mon Sep 17 00:00:00 2001 From: ShuNing Date: Fri, 19 May 2023 18:15:37 +0800 Subject: [PATCH] controller: avoid the invalid nil pointer (#6486) close tikv/pd#6485, ref tikv/pd#6485 controller: avoid the invalid nil pointer Signed-off-by: nolouch --- client/resource_group/controller/model.go | 46 +++- .../resource_group/controller/model_test.go | 218 ++++++++++++++++++ 2 files changed, 256 insertions(+), 8 deletions(-) create mode 100644 client/resource_group/controller/model_test.go diff --git a/client/resource_group/controller/model.go b/client/resource_group/controller/model.go index 82e7fb2dd5f..e8b5eb1679d 100644 --- a/client/resource_group/controller/model.go +++ b/client/resource_group/controller/model.go @@ -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 @@ -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 @@ -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 @@ -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 diff --git a/client/resource_group/controller/model_test.go b/client/resource_group/controller/model_test.go new file mode 100644 index 00000000000..94007485520 --- /dev/null +++ b/client/resource_group/controller/model_test.go @@ -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) +}