From 78082cf4ce6e9522287f3e0881b89f846ab585d5 Mon Sep 17 00:00:00 2001
From: JmPotato <ghzpotato@gmail.com>
Date: Wed, 8 Feb 2023 17:13:26 +0800
Subject: [PATCH] Pay back WRU when the request fails

Signed-off-by: JmPotato <ghzpotato@gmail.com>
---
 client/resource_manager/client/client.go |  4 +--
 client/resource_manager/client/model.go  | 39 +++++++++++++++++++-----
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/client/resource_manager/client/client.go b/client/resource_manager/client/client.go
index 61c2a1048dd..e38a0e111ec 100644
--- a/client/resource_manager/client/client.go
+++ b/client/resource_manager/client/client.go
@@ -32,11 +32,11 @@ const (
 	maxNotificationChanLen = 200
 )
 
-// ResourceGroupKVInterceptor is used as quato limit controller for resource group using kv store.
+// ResourceGroupKVInterceptor is used as quota limit controller for resource group using kv store.
 type ResourceGroupKVInterceptor interface {
 	// OnRequestWait is used to check whether resource group has enough tokens. It maybe needs wait some time.
 	OnRequestWait(ctx context.Context, resourceGroupName string, info RequestInfo) error
-	// OnResponse is used to consume tokens atfer receiving response
+	// OnResponse is used to consume tokens after receiving response
 	OnResponse(ctx context.Context, resourceGroupName string, req RequestInfo, resp ResponseInfo) error
 }
 
diff --git a/client/resource_manager/client/model.go b/client/resource_manager/client/model.go
index 71cbb59d707..7c2ffe7716a 100644
--- a/client/resource_manager/client/model.go
+++ b/client/resource_manager/client/model.go
@@ -37,6 +37,9 @@ type RequestInfo interface {
 type ResponseInfo interface {
 	ReadBytes() uint64
 	KVCPUMs() uint64
+	// Succeed is used to tell whether the request is successfully returned.
+	// If not, we need to pay back the WRU cost of the request.
+	Succeed() bool
 }
 
 // ResourceCalculator is used to calculate the resource consumption of a request.
@@ -72,9 +75,7 @@ func (kc *KVCalculator) BeforeKVRequest(consumption *rmpb.Consumption, req Reque
 	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
+		kc.calculateWriteCost(consumption, req)
 	} else {
 		consumption.KvReadRpcCount += 1
 		// Read bytes could not be known before the request is executed,
@@ -83,20 +84,44 @@ func (kc *KVCalculator) BeforeKVRequest(consumption *rmpb.Consumption, req Reque
 	}
 }
 
+func (kc *KVCalculator) calculateWriteCost(consumption *rmpb.Consumption, req RequestInfo) {
+	writeBytes := float64(req.WriteBytes())
+	consumption.WriteBytes += writeBytes
+	consumption.WRU += float64(kc.WriteBaseCost) + float64(kc.WriteBytesCost)*writeBytes
+}
+
 // AfterKVRequest ...
 func (kc *KVCalculator) AfterKVRequest(consumption *rmpb.Consumption, req RequestInfo, res ResponseInfo) {
-	// 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.RRU += float64(kc.CPUMsCost) * kvCPUMs
+		// For now, we can only collect the KV CPU cost for a read request.
+		kc.calculateCPUCost(consumption, res)
+	} else if !res.Succeed() {
+		// If the write request is not successfully returned, we need to pay back the WRU cost.
+		kc.payBackWriteCost(consumption, req)
 	}
 	// A write request may also read data, which should be counted into the RRU cost.
+	// This part should be counted even if the request does not succeed.
+	kc.calculateReadCost(consumption, res)
+}
+
+func (kc *KVCalculator) calculateReadCost(consumption *rmpb.Consumption, res ResponseInfo) {
 	readBytes := float64(res.ReadBytes())
 	consumption.ReadBytes += readBytes
 	consumption.RRU += float64(kc.ReadBytesCost) * readBytes
 }
 
+func (kc *KVCalculator) calculateCPUCost(consumption *rmpb.Consumption, res ResponseInfo) {
+	kvCPUMs := float64(res.KVCPUMs())
+	consumption.TotalCpuTimeMs += kvCPUMs
+	consumption.RRU += float64(kc.CPUMsCost) * kvCPUMs
+}
+
+func (kc *KVCalculator) payBackWriteCost(consumption *rmpb.Consumption, req RequestInfo) {
+	writeBytes := float64(req.WriteBytes())
+	consumption.WriteBytes -= writeBytes
+	consumption.WRU -= float64(kc.WriteBaseCost) + float64(kc.WriteBytesCost)*writeBytes
+}
+
 // SQLCalculator is used to calculate the SQL-side consumption.
 type SQLCalculator struct {
 	*Config