Address the comments

Signed-off-by: JmPotato <ghzpotato@gmail.com>

pkg/mcs/resource_manager/client/config.go

@@ -17,9 +17,9 @@ package client
 const (
 	defaultReadBaseCost     = 1
 	defaultReadCostPerByte  = 1. / 1024 / 1024
+	defaultReadCPUMsCost    = 1
 	defaultWriteBaseCost    = 3
 	defaultWriteCostPerByte = 5. / 1024 / 1024
-	defaultWriteCPUMsCost   = 1
 )
 
 // RequestUnitConfig is the configuration of the request units, which determines the coefficients of
@@ -30,14 +30,14 @@ type RequestUnitConfig struct {
 	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"`
+	// ReadCPUMsCost is the cost for each millisecond of CPU time taken by a read request.
+	// It's 1 millisecond = 1 RRU by default.
+	ReadCPUMsCost float64 `toml:"write-cpu-ms-cost" json:"write-cpu-ms-cost"`
 	// 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 = 5 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.
@@ -47,7 +47,7 @@ func DefaultRequestUnitConfig() *RequestUnitConfig {
 		ReadCostPerByte:  defaultReadCostPerByte,
 		WriteBaseCost:    defaultWriteBaseCost,
 		WriteCostPerByte: defaultWriteCostPerByte,
-		WriteCPUMsCost:   defaultWriteCPUMsCost,
+		ReadCPUMsCost:    defaultReadCPUMsCost,
 	}
 }
 
@@ -57,9 +57,9 @@ func DefaultRequestUnitConfig() *RequestUnitConfig {
 type Config struct {
 	ReadBaseCost   RequestUnit
 	ReadBytesCost  RequestUnit
+	ReadCPUMsCost  RequestUnit
 	WriteBaseCost  RequestUnit
 	WriteBytesCost RequestUnit
-	WriteCPUMsCost RequestUnit
 	// TODO: add SQL computing CPU cost.
 }
 
@@ -75,9 +75,9 @@ func generateConfig(ruConfig *RequestUnitConfig) *Config {
 	cfg := &Config{
 		ReadBaseCost:   RequestUnit(ruConfig.ReadBaseCost),
 		ReadBytesCost:  RequestUnit(ruConfig.ReadCostPerByte),
+		ReadCPUMsCost:  RequestUnit(ruConfig.ReadCPUMsCost),
 		WriteBaseCost:  RequestUnit(ruConfig.WriteBaseCost),
 		WriteBytesCost: RequestUnit(ruConfig.WriteCostPerByte),
-		WriteCPUMsCost: RequestUnit(ruConfig.WriteCPUMsCost),
 	}
 	return cfg
 }

pkg/mcs/resource_manager/client/model.go

@@ -43,20 +43,22 @@ type ResponseInfo interface {
 type ResourceCalculator interface {
 	// Trickle is used to calculate the resource consumption periodically rather than on the request path.
 	// It's mainly used to calculate like the SQL CPU cost.
-	Trickle(context.Context, map[rmpb.ResourceType]float64, map[rmpb.RequestUnitType]float64)
+	Trickle(context.Context, *rmpb.Consumption)
 	// BeforeKVRequest is used to calculate the resource consumption before the KV request.
 	// It's mainly used to calculate the base and write request cost.
-	BeforeKVRequest(map[rmpb.ResourceType]float64, map[rmpb.RequestUnitType]float64, RequestInfo)
+	BeforeKVRequest(*rmpb.Consumption, RequestInfo)
 	// AfterKVRequest is used to calculate the resource consumption after the KV request.
 	// It's mainly used to calculate the read request cost and KV CPU cost.
-	AfterKVRequest(map[rmpb.ResourceType]float64, map[rmpb.RequestUnitType]float64, RequestInfo, ResponseInfo)
+	AfterKVRequest(*rmpb.Consumption, RequestInfo, ResponseInfo)
 }
 
 // KVCalculator is used to calculate the KV-side consumption.
 type KVCalculator struct {
 	*Config
 }
 
+var _ ResourceCalculator = (*KVCalculator)(nil)
+
 // func newKVCalculator(cfg *Config) *KVCalculator {
 // 	return &KVCalculator{Config: cfg}
 // }
@@ -69,46 +71,52 @@ func (kc *KVCalculator) Trickle(ctx context.Context, consumption *rmpb.Consumpti
 func (kc *KVCalculator) BeforeKVRequest(consumption *rmpb.Consumption, req RequestInfo) {
 	if req.IsWrite() {
 		consumption.KvWriteRpcCount += 1
+		// Write bytes are knowable in advance, so we can calculate the WRU cost here.
 		writeBytes := float64(req.WriteBytes())
 		consumption.WriteBytes += writeBytes
 		consumption.WRU += float64(kc.WriteBaseCost) + float64(kc.WriteBytesCost)*writeBytes
 	} else {
 		consumption.KvReadRpcCount += 1
+		// Read bytes could not be known before the request is executed,
+		// so we only add the base cost here.
 		consumption.RRU += float64(kc.ReadBaseCost)
 	}
 }
 
 // AfterKVRequest ...
 func (kc *KVCalculator) AfterKVRequest(consumption *rmpb.Consumption, req RequestInfo, res ResponseInfo) {
-	if req.IsWrite() {
+	// For now, we can only collect the KV CPU cost for a read request.
+	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
+		consumption.RRU += float64(kc.ReadCPUMsCost) * kvCPUMs
 	}
+	// A write request may also read data, which should also be counted into the RRU cost.
+	readBytes := float64(res.ReadBytes())
+	consumption.ReadBytes += readBytes
+	consumption.RRU += float64(kc.ReadBytesCost) * readBytes
 }
 
 // SQLCalculator is used to calculate the SQL-side consumption.
 type SQLCalculator struct {
 	*Config
 }
 
+var _ ResourceCalculator = (*SQLCalculator)(nil)
+
 // func newSQLCalculator(cfg *Config) *SQLCalculator {
 // 	return &SQLCalculator{Config: cfg}
 // }
 
 // Trickle ...
 // TODO: calculate the SQL CPU cost and related resource consumption.
-func (dsc *SQLCalculator) Trickle(ctx context.Context, resource map[rmpb.ResourceType]float64, ru map[rmpb.RequestUnitType]float64) {
+func (dsc *SQLCalculator) Trickle(ctx context.Context, consumption *rmpb.Consumption) {
 }
 
 // BeforeKVRequest ...
-func (dsc *SQLCalculator) BeforeKVRequest(resource map[rmpb.ResourceType]float64, ru map[rmpb.RequestUnitType]float64, req RequestInfo) {
+func (dsc *SQLCalculator) BeforeKVRequest(consumption *rmpb.Consumption, req RequestInfo) {
 }
 
 // AfterKVRequest ...
-func (dsc *SQLCalculator) AfterKVRequest(resource map[rmpb.ResourceType]float64, ru map[rmpb.RequestUnitType]float64, req RequestInfo, res ResponseInfo) {
+func (dsc *SQLCalculator) AfterKVRequest(consumption *rmpb.Consumption, req RequestInfo, res ResponseInfo) {
 }