Introduce the resource unit model and config

Signed-off-by: JmPotato <[email protected]>

pkg/mcs/resource_manager/client/config.go

@@ -0,0 +1,83 @@
+// 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,g
+// 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 client
+
+const (
+	defaultReadBaseCost     = 1
+	defaultReadCostPerByte  = 1. / 1024 / 1024
+	defaultWriteBaseCost    = 5
+	defaultWriteCostPerByte = 10. / 1024 / 1024
+	defaultWriteCPUMsCost   = 1
+)
+
+// RequestUnitConfig is the configuration of the request units, which determines the coefficients of
+// the RRU and WRU cost. This configuration should be modified carefully.
+type RequestUnitConfig struct {
+	// ReadBaseCost is the base cost for a read request. No matter how many bytes read/written or
+	// the CPU times taken for a request, this cost is inevitable.
+	ReadBaseCost float64 `toml:"read-base-cost" json:"read-base-cost"`
+	// ReadCostPerByte is the cost for each byte read. It's 1 MiB = 1 RRU by default.
+	ReadCostPerByte float64 `toml:"read-cost-per-byte" json:"read-cost-per-byte"`
+	// WriteBaseCost is the base cost for a write request. No matter how many bytes read/written or
+	// the CPU times taken for a request, this cost is inevitable.
+	WriteBaseCost float64 `toml:"write-base-cost" json:"write-base-cost"`
+	// WriteCostPerByte is the cost for each byte written. It's 1 MiB = 10 WRU by default.
+	WriteCostPerByte float64 `toml:"write-cost-per-byte" json:"write-cost-per-byte"`
+	// WriteCPUMsCost is the cost for each millisecond of CPU time taken by a write request.
+	// It's 1 millisecond = 1 WRU by default.
+	WriteCPUMsCost float64 `toml:"write-cpu-ms-cost" json:"write-cpu-ms-cost"`
+}
+
+// DefaultRequestUnitConfig returns the default request unit configuration.
+func DefaultRequestUnitConfig() *RequestUnitConfig {
+	return &RequestUnitConfig{
+		ReadBaseCost:     defaultReadBaseCost,
+		ReadCostPerByte:  defaultReadCostPerByte,
+		WriteBaseCost:    defaultWriteBaseCost,
+		WriteCostPerByte: defaultWriteCostPerByte,
+		WriteCPUMsCost:   defaultWriteCPUMsCost,
+	}
+}
+
+// Config is the configuration of the resource units, which gives the read/write request
+// units or request resource cost standards. It should be calculated by a given `RequestUnitConfig`
+// or `RequestResourceConfig`.
+type Config struct {
+	ReadBaseCost   RequestUnit
+	ReadBytesCost  RequestUnit
+	WriteBaseCost  RequestUnit
+	WriteBytesCost RequestUnit
+	WriteCPUMsCost RequestUnit
+	// TODO: add SQL computing CPU cost.
+}
+
+// DefaultConfig returns the default configuration.
+func DefaultConfig() *Config {
+	cfg := generateConfig(
+		DefaultRequestUnitConfig(),
+	)
+	return cfg
+}
+
+func generateConfig(ruConfig *RequestUnitConfig) *Config {
+	cfg := &Config{
+		ReadBaseCost:   RequestUnit(ruConfig.ReadBaseCost),
+		ReadBytesCost:  RequestUnit(ruConfig.ReadCostPerByte),
+		WriteBaseCost:  RequestUnit(ruConfig.WriteBaseCost),
+		WriteBytesCost: RequestUnit(ruConfig.WriteCostPerByte),
+		WriteCPUMsCost: RequestUnit(ruConfig.WriteCPUMsCost),
+	}
+	return cfg
+}

pkg/mcs/resource_manager/client/model.go

@@ -0,0 +1,100 @@
+// 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,g
+// 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 client
+
+import (
+	"context"
+
+	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
+)
+
+// RequestUnit is the basic unit of the resource request management, which has two types:
+//   - RRU: read request unit
+//   - WRU: write request unit
+type RequestUnit float64
+
+// RequestInfo is the interface of the request information provider. A request should be
+// able tell whether it's a write request and if so, the written bytes would also be provided.
+type RequestInfo interface {
+	IsWrite() bool
+	WriteBytes() uint64
+}
+
+// ResponseInfo is the interface of the response information provider. A response should be
+// able tell how many bytes it read and KV CPU cost in milliseconds.
+type ResponseInfo interface {
+	ReadBytes() uint64
+	KVCPUMs() uint64
+}
+
+// ResourceCalculator is used to calculate the resource consumption of a request.
+type ResourceCalculator interface {
+	Trickle(map[rmpb.ResourceType]float64, map[rmpb.RequestUnitType]float64, context.Context)
+	BeforeKVRequest(map[rmpb.ResourceType]float64, map[rmpb.RequestUnitType]float64, RequestInfo)
+	AfterKVRequest(map[rmpb.ResourceType]float64, map[rmpb.RequestUnitType]float64, RequestInfo, ResponseInfo)
+}
+
+// KVCalculator is used to calculate the KV request consumption.
+type KVCalculator struct {
+	*Config
+}
+
+func newKVCalculator(cfg *Config) *KVCalculator {
+	return &KVCalculator{Config: cfg}
+}
+
+func (kc *KVCalculator) Trickle(consumption *rmpb.Consumption, ctx context.Context) {
+}
+
+func (kc *KVCalculator) BeforeKVRequest(consumption *rmpb.Consumption, req RequestInfo) {
+	if req.IsWrite() {
+		consumption.KvWriteRpcCount += 1
+		writeBytes := float64(req.WriteBytes())
+		consumption.WriteBytes += writeBytes
+		consumption.WRU += float64(kc.WriteBaseCost) + float64(kc.WriteBytesCost)*writeBytes
+	} else {
+		consumption.KvReadRpcCount += 1
+		consumption.RRU += float64(kc.ReadBaseCost)
+	}
+}
+func (kc *KVCalculator) AfterKVRequest(consumption *rmpb.Consumption, req RequestInfo, res ResponseInfo) {
+	if req.IsWrite() {
+		kvCPUMs := float64(res.KVCPUMs())
+		consumption.TotalCpuTimeMs += kvCPUMs
+		consumption.WRU += float64(kc.WriteCPUMsCost) * kvCPUMs
+	} else {
+		readBytes := float64(res.ReadBytes())
+		consumption.ReadBytes += readBytes
+		consumption.RRU += float64(kc.ReadBytesCost) * readBytes
+	}
+}
+
+type SQLCPUCalculator struct {
+	*Config
+}
+
+func newSQLCPUCalculator(cfg *Config) *SQLCPUCalculator {
+	return &SQLCPUCalculator{Config: cfg}
+}
+
+// TODO: calculate the SQL CPU cost and related resource consumption.
+func (dsc *SQLCPUCalculator) Trickle(resource map[rmpb.ResourceType]float64, ru map[rmpb.RequestUnitType]float64, ctx context.Context) {
+}
+
+func (dsc *SQLCPUCalculator) BeforeKVRequest(resource map[rmpb.ResourceType]float64, ru map[rmpb.RequestUnitType]float64, req RequestInfo) {
+}
+
+func (dsc *SQLCPUCalculator) AfterKVRequest(resource map[rmpb.ResourceType]float64, ru map[rmpb.RequestUnitType]float64, req RequestInfo, res ResponseInfo) {
+}