client/resource_group: delay reporting write reservation

client/resource_group/controller/group_controller.go

@@ -572,9 +572,10 @@ func (gc *groupCostController) onRequestWaitImpl(
 	for _, calc := range gc.calculators {
 		calc.BeforeKVRequest(delta, info)
 	}
+	reportedDelta := reportedRequestConsumption(gc.calculators, info, delta)
 
 	gc.mu.Lock()
-	add(gc.mu.consumption, delta)
+	add(gc.mu.consumption, reportedDelta)
 	gc.mu.Unlock()
 
 	if !gc.burstable.Load() {
@@ -587,7 +588,7 @@ func (gc *groupCostController) onRequestWaitImpl(
 				gc.metrics.failedRequestCounterWithOthers.Inc()
 			}
 			gc.mu.Lock()
-			sub(gc.mu.consumption, delta)
+			sub(gc.mu.consumption, reportedDelta)
 			gc.mu.Unlock()
 			failpoint.Inject("triggerUpdate", func() {
 				gc.lowRUNotifyChan <- notifyMsg{}
@@ -622,16 +623,20 @@ func (gc *groupCostController) onResponseImpl(
 	for _, calc := range gc.calculators {
 		calc.AfterKVRequest(delta, req, resp)
 	}
+	reportedDelta := reportedResponseConsumption(gc.calculators, req, resp, delta)
 	if !gc.burstable.Load() {
 		counter := gc.run.requestUnitTokens
 		if v := getRUValueFromConsumption(delta); v > 0 {
 			counter.limiter.RemoveTokens(time.Now(), v)
+		} else if v < 0 {
+			// Failed writes refund the response-side write payback delta.
+			counter.limiter.RefundTokens(time.Now(), -v)
 		}
 	}
 
 	gc.mu.Lock()
 	// Record the consumption of the request
-	add(gc.mu.consumption, delta)
+	add(gc.mu.consumption, reportedDelta)
 	// Record the consumption of the request by store
 	count := &rmpb.Consumption{}
 	*count = *delta
@@ -653,26 +658,36 @@ func (gc *groupCostController) onResponseWaitImpl(
 	for _, calc := range gc.calculators {
 		calc.AfterKVRequest(delta, req, resp)
 	}
+	reportedDelta := reportedResponseConsumption(gc.calculators, req, resp, delta)
 	var waitDuration time.Duration
 	if !gc.burstable.Load() {
-		allowDebt := delta.ReadBytes+delta.WriteBytes < bigRequestThreshold || !gc.isThrottled.Load()
-		d, err := gc.acquireTokens(ctx, delta, &waitDuration, allowDebt)
-		if err != nil {
-			if errs.ErrClientResourceGroupThrottled.Equal(err) {
-				gc.metrics.failedRequestCounterWithThrottled.Inc()
-				gc.metrics.failedLimitReserveDuration.Observe(d.Seconds())
-			} else {
-				gc.metrics.failedRequestCounterWithOthers.Inc()
+		v := getRUValueFromConsumption(delta)
+		if v > 0 {
+			allowDebt := delta.ReadBytes+delta.WriteBytes < bigRequestThreshold || !gc.isThrottled.Load()
+			d, err := gc.acquireTokens(ctx, delta, &waitDuration, allowDebt)
+			if err != nil {
+				if errs.ErrClientResourceGroupThrottled.Equal(err) {
+					gc.metrics.failedRequestCounterWithThrottled.Inc()
+					gc.metrics.failedLimitReserveDuration.Observe(d.Seconds())
+				} else {
+					gc.metrics.failedRequestCounterWithOthers.Inc()
+				}
+				return nil, waitDuration, err
 			}
-			return nil, waitDuration, err
+			gc.metrics.successfulRequestDuration.Observe(d.Seconds())
+			waitDuration += d
+		} else if v < 0 {
+			// Failed writes refund the response-side write payback delta.
+			gc.run.requestUnitTokens.limiter.RefundTokens(time.Now(), -v)
+			gc.metrics.successfulRequestDuration.Observe(0)
+		} else {
+			gc.metrics.successfulRequestDuration.Observe(0)
 		}
-		gc.metrics.successfulRequestDuration.Observe(d.Seconds())
-		waitDuration += d
 	}
 
 	gc.mu.Lock()
 	// Record the consumption of the request
-	add(gc.mu.consumption, delta)
+	add(gc.mu.consumption, reportedDelta)
 	// Record the consumption of the request by store
 	count := &rmpb.Consumption{}
 	*count = *delta

client/resource_group/controller/group_controller_test.go

@@ -209,6 +209,109 @@ func TestOnResponseWaitConsumption(t *testing.T) {
 	verify()
 }
 
+func TestFailedWriteReservationDoesNotEnterReportedConsumption(t *testing.T) {
+	re := require.New(t)
+	gc := createTestGroupCostController(re)
+
+	req := &TestRequestInfo{
+		isWrite:     true,
+		writeBytes:  1024,
+		numReplicas: 3,
+	}
+	resp := &TestResponseInfo{
+		readBytes: 0,
+		succeed:   false,
+	}
+
+	tokenDelta, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req)
+	re.NoError(err)
+	re.Positive(tokenDelta.WRU)
+	gc.updateRunState()
+	requestReport := gc.collectRequestAndConsumption(periodicReport)
+	re.NotNil(requestReport)
+	re.Zero(requestReport.GetConsumptionSinceLastRequest().GetWRU(),
+		"write reservation should not be reported as actual usage before the response succeeds")
+	gc.handleTokenBucketResponse(&rmpb.TokenBucketResponse{})
+
+	refundDelta, _, err := gc.onResponseWaitImpl(context.TODO(), req, resp)
+	re.NoError(err)
+	re.Negative(refundDelta.WRU)
+	gc.run.lastRequestTime = time.Now().Add(-extendedReportingPeriodFactor * defaultTargetPeriod)
+	gc.updateRunState()
+	refundReport := gc.collectRequestAndConsumption(periodicReport)
+	re.NotNil(refundReport)
+	re.Zero(refundReport.GetConsumptionSinceLastRequest().GetWRU(),
+		"failed write refund must not make actual usage negative")
+}
+
+func TestFailedWritePaybackRefundsLimiterOnResponseWait(t *testing.T) {
+	re := require.New(t)
+	gc := createTestGroupCostController(re)
+
+	gc.run.requestUnitTokens.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{
+		newTokens:   100000,
+		newFillRate: 0,
+		newBurst:    0,
+	})
+
+	req := &TestRequestInfo{
+		isWrite:     true,
+		writeBytes:  4 * 1024,
+		numReplicas: 3,
+	}
+	resp := &TestResponseInfo{
+		readBytes: 0,
+		succeed:   false,
+	}
+
+	_, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req)
+	re.NoError(err)
+	tokensAfterReservation := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now())
+
+	refundDelta, _, err := gc.onResponseWaitImpl(context.TODO(), req, resp)
+	re.NoError(err)
+	expectedRefund := -getRUValueFromConsumption(refundDelta)
+	re.Positive(expectedRefund)
+	tokensAfterResponse := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now())
+
+	re.InDelta(tokensAfterReservation+expectedRefund, tokensAfterResponse, 1.0,
+		"failed write payback should refund the local limiter")
+}
+
+func TestFailedWritePaybackRefundsLimiterOnResponse(t *testing.T) {
+	re := require.New(t)
+	gc := createTestGroupCostController(re)
+
+	gc.run.requestUnitTokens.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{
+		newTokens:   100000,
+		newFillRate: 0,
+		newBurst:    0,
+	})
+
+	req := &TestRequestInfo{
+		isWrite:     true,
+		writeBytes:  4 * 1024,
+		numReplicas: 3,
+	}
+	resp := &TestResponseInfo{
+		readBytes: 0,
+		succeed:   false,
+	}
+
+	_, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req)
+	re.NoError(err)
+	tokensAfterReservation := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now())
+
+	refundDelta, err := gc.onResponseImpl(req, resp)
+	re.NoError(err)
+	expectedRefund := -getRUValueFromConsumption(refundDelta)
+	re.Positive(expectedRefund)
+	tokensAfterResponse := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now())
+
+	re.InDelta(tokensAfterReservation+expectedRefund, tokensAfterResponse, 1.0,
+		"failed write payback should refund the local limiter")
+}
+
 func TestHandleTokenBucketUpdateEventCanceledByInitCounterNotify(t *testing.T) {
 	re := require.New(t)
 	gc := createTestGroupCostController(re)

client/resource_group/controller/limiter.go

@@ -326,6 +326,24 @@ func (lim *Limiter) RemoveTokens(now time.Time, amount float64) {
 	lim.maybeNotify()
 }
 
+// RefundTokens adds tokens back to the limiter.
+//
+// Token overshoot above burst is intentional here: getTokens clamps the
+// value on the next call, so RefundTokens never permanently leaks past
+// the burst cap. maybeNotify is intentionally not called either because
+// refunding can only raise the token count, never push the limiter into
+// the low-token state.
+func (lim *Limiter) RefundTokens(now time.Time, amount float64) {
+	lim.mu.Lock()
+	defer lim.mu.Unlock()
+	if lim.burst < 0 || lim.fillRate == Inf {
+		return
+	}
+	_, tokens := lim.getTokens(now)
+	lim.updateLast(now)
+	lim.tokens = tokens + amount
+}
+
 type tokenBucketReconfigureArgs struct {
 	newTokens          float64
 	newFillRate        float64

client/resource_group/controller/limiter_test.go

@@ -219,6 +219,40 @@ func TestLastStaysMonotonicOnStaleNow(t *testing.T) {
 	}
 }
 
+func TestRefundTokens(t *testing.T) {
+	re := require.New(t)
+	nc := make(chan notifyMsg, 1)
+	lim := NewLimiter(t0, 0, 0, 100, nc)
+
+	// Consume some tokens.
+	lim.RemoveTokens(t0, 30)
+	checkTokens(re, lim, t0, 70)
+
+	// Refund part of them.
+	lim.RefundTokens(t0, 20)
+	checkTokens(re, lim, t0, 90)
+
+	// Refund beyond the initial amount - allowed (no burst cap when burst==0).
+	lim.RefundTokens(t0, 50)
+	checkTokens(re, lim, t0, 140)
+
+	// With burst > 0, refund up to burst stays.
+	limBurst := NewLimiter(t0, 0, 100, 50, nc)
+	limBurst.RemoveTokens(t0, 30)
+	checkTokens(re, limBurst, t0, 20)
+	limBurst.RefundTokens(t0, 80) // tokens = 20 + 80 = 100, burst = 100
+	checkTokens(re, limBurst, t0, 100)
+
+	// Refund beyond burst - capped by getTokens.
+	limBurst.RefundTokens(t0, 50) // tokens = 100 + 50 = 150, capped to 100
+	checkTokens(re, limBurst, t0, 100)
+
+	// Burstable (burst < 0): RefundTokens is a no-op.
+	limUnlimited := NewLimiter(t0, 0, -1, 100, nc)
+	limUnlimited.RefundTokens(t0, 50)
+	checkTokens(re, limUnlimited, t0, 100)
+}
+
 func TestNotify(t *testing.T) {
 	nc := make(chan notifyMsg, 1)
 	lim := NewLimiter(t0, 1, 0, 0, nc)

client/resource_group/controller/model.go

@@ -185,6 +185,62 @@ func (kc *KVCalculator) payBackWriteCost(consumption *rmpb.Consumption, req Requ
 	consumption.WRU -= float64(kc.WriteBaseCost) + float64(kc.WriteBytesCost)*writeBytes
 }
 
+func cloneConsumption(consumption *rmpb.Consumption) *rmpb.Consumption {
+	if consumption == nil {
+		return &rmpb.Consumption{}
+	}
+	cloned := *consumption
+	return &cloned
+}
+
+func writeReservationConsumption(calculators []ResourceCalculator, req RequestInfo) *rmpb.Consumption {
+	if !req.IsWrite() {
+		return nil
+	}
+	reservation := &rmpb.Consumption{}
+	for _, calc := range calculators {
+		kvCalc, ok := calc.(*KVCalculator)
+		if !ok {
+			continue
+		}
+		kvCalc.calculateWriteCost(reservation, req)
+	}
+	return reservation
+}
+
+func writeRefundConsumption(calculators []ResourceCalculator, req RequestInfo) *rmpb.Consumption {
+	if !req.IsWrite() {
+		return nil
+	}
+	refund := &rmpb.Consumption{}
+	for _, calc := range calculators {
+		kvCalc, ok := calc.(*KVCalculator)
+		if !ok {
+			continue
+		}
+		kvCalc.payBackWriteCost(refund, req)
+	}
+	return refund
+}
+
+func reportedRequestConsumption(calculators []ResourceCalculator, req RequestInfo, tokenDelta *rmpb.Consumption) *rmpb.Consumption {
+	reported := cloneConsumption(tokenDelta)
+	sub(reported, writeReservationConsumption(calculators, req))
+	return reported
+}
+
+func reportedResponseConsumption(calculators []ResourceCalculator, req RequestInfo, resp ResponseInfo, tokenDelta *rmpb.Consumption) *rmpb.Consumption {
+	reported := cloneConsumption(tokenDelta)
+	if req.IsWrite() {
+		if resp.Succeed() {
+			add(reported, writeReservationConsumption(calculators, req))
+		} else {
+			sub(reported, writeRefundConsumption(calculators, req))
+		}
+	}
+	return reported
+}
+
 // SQLCalculator is used to calculate the SQL-side consumption.
 type SQLCalculator struct {
 	*RUConfig

client/resource_group/controller/model_test.go

@@ -39,6 +39,21 @@ func TestGetRUValueFromConsumption(t *testing.T) {
 	re.Equal(expected, result)
 }
 
+func TestWriteConsumptionHelpersSkipReadRequests(t *testing.T) {
+	re := require.New(t)
+	calculators := []ResourceCalculator{newKVCalculator(DefaultRUConfig())}
+	req := &TestRequestInfo{isWrite: false}
+
+	re.Nil(writeReservationConsumption(calculators, req))
+	re.Nil(writeRefundConsumption(calculators, req))
+	re.Zero(testing.AllocsPerRun(1000, func() {
+		_ = writeReservationConsumption(calculators, req)
+	}))
+	re.Zero(testing.AllocsPerRun(1000, func() {
+		_ = writeRefundConsumption(calculators, req)
+	}))
+}
+
 func TestAdd(t *testing.T) {
 	// Positive test case
 	re := require.New(t)