From ebe73aa1864a5a04a477588933ed84a7e113406d Mon Sep 17 00:00:00 2001 From: Yuhao Zhang Date: Wed, 8 Jul 2026 13:05:00 +0800 Subject: [PATCH 1/3] client/resource_group: delay reporting write reservation Signed-off-by: Yuhao Zhang --- .../controller/group_controller.go | 11 ++-- .../controller/group_controller_test.go | 35 ++++++++++++ client/resource_group/controller/model.go | 56 +++++++++++++++++++ 3 files changed, 98 insertions(+), 4 deletions(-) diff --git a/client/resource_group/controller/group_controller.go b/client/resource_group/controller/group_controller.go index 04ac7e8b771..7a8b009ae44 100644 --- a/client/resource_group/controller/group_controller.go +++ b/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,6 +623,7 @@ 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 { @@ -631,7 +633,7 @@ func (gc *groupCostController) onResponseImpl( 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,6 +655,7 @@ 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() @@ -672,7 +675,7 @@ func (gc *groupCostController) onResponseWaitImpl( 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 diff --git a/client/resource_group/controller/group_controller_test.go b/client/resource_group/controller/group_controller_test.go index c9e912d6ffc..e8f80406ccb 100644 --- a/client/resource_group/controller/group_controller_test.go +++ b/client/resource_group/controller/group_controller_test.go @@ -209,6 +209,41 @@ 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 TestHandleTokenBucketUpdateEventCanceledByInitCounterNotify(t *testing.T) { re := require.New(t) gc := createTestGroupCostController(re) diff --git a/client/resource_group/controller/model.go b/client/resource_group/controller/model.go index 54721969811..45792de106d 100644 --- a/client/resource_group/controller/model.go +++ b/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 { + reservation := &rmpb.Consumption{} + if !req.IsWrite() { + return reservation + } + 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 { + refund := &rmpb.Consumption{} + if !req.IsWrite() { + return refund + } + 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 From 93e00420ba81d496fb61d09f75322b995c21f6c2 Mon Sep 17 00:00:00 2001 From: Yuhao Zhang Date: Wed, 8 Jul 2026 13:30:06 +0800 Subject: [PATCH 2/3] client/resource_group: refund failed write reservations Signed-off-by: Yuhao Zhang --- .../controller/group_controller.go | 34 +++++++--- .../controller/group_controller_test.go | 68 +++++++++++++++++++ client/resource_group/controller/limiter.go | 18 +++++ .../resource_group/controller/limiter_test.go | 34 ++++++++++ client/resource_group/controller/model.go | 8 +-- .../resource_group/controller/model_test.go | 15 ++++ 6 files changed, 162 insertions(+), 15 deletions(-) diff --git a/client/resource_group/controller/group_controller.go b/client/resource_group/controller/group_controller.go index 7a8b009ae44..7f4707755da 100644 --- a/client/resource_group/controller/group_controller.go +++ b/client/resource_group/controller/group_controller.go @@ -628,6 +628,9 @@ func (gc *groupCostController) onResponseImpl( counter := gc.run.requestUnitTokens if v := getRUValueFromConsumption(delta); v > 0 { counter.limiter.RemoveTokens(time.Now(), v) + } else if v < 0 { + // Failed writes pay back the local write reservation. + counter.limiter.RefundTokens(time.Now(), -v) } } @@ -658,19 +661,28 @@ func (gc *groupCostController) onResponseWaitImpl( 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 pay back the local write reservation. + 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() diff --git a/client/resource_group/controller/group_controller_test.go b/client/resource_group/controller/group_controller_test.go index e8f80406ccb..545b0e61002 100644 --- a/client/resource_group/controller/group_controller_test.go +++ b/client/resource_group/controller/group_controller_test.go @@ -244,6 +244,74 @@ func TestFailedWriteReservationDoesNotEnterReportedConsumption(t *testing.T) { "failed write refund must not make actual usage negative") } +func TestFailedWriteReservationRefundsLimiterOnResponseWait(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 write reservation") +} + +func TestFailedWriteReservationRefundsLimiterOnResponse(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 write reservation") +} + func TestHandleTokenBucketUpdateEventCanceledByInitCounterNotify(t *testing.T) { re := require.New(t) gc := createTestGroupCostController(re) diff --git a/client/resource_group/controller/limiter.go b/client/resource_group/controller/limiter.go index 9086a8026f8..50e8b0f09a6 100644 --- a/client/resource_group/controller/limiter.go +++ b/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 diff --git a/client/resource_group/controller/limiter_test.go b/client/resource_group/controller/limiter_test.go index 01cc7fe1977..d46c30f6587 100644 --- a/client/resource_group/controller/limiter_test.go +++ b/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) diff --git a/client/resource_group/controller/model.go b/client/resource_group/controller/model.go index 45792de106d..b6ba7551975 100644 --- a/client/resource_group/controller/model.go +++ b/client/resource_group/controller/model.go @@ -194,10 +194,10 @@ func cloneConsumption(consumption *rmpb.Consumption) *rmpb.Consumption { } func writeReservationConsumption(calculators []ResourceCalculator, req RequestInfo) *rmpb.Consumption { - reservation := &rmpb.Consumption{} if !req.IsWrite() { - return reservation + return nil } + reservation := &rmpb.Consumption{} for _, calc := range calculators { kvCalc, ok := calc.(*KVCalculator) if !ok { @@ -209,10 +209,10 @@ func writeReservationConsumption(calculators []ResourceCalculator, req RequestIn } func writeRefundConsumption(calculators []ResourceCalculator, req RequestInfo) *rmpb.Consumption { - refund := &rmpb.Consumption{} if !req.IsWrite() { - return refund + return nil } + refund := &rmpb.Consumption{} for _, calc := range calculators { kvCalc, ok := calc.(*KVCalculator) if !ok { diff --git a/client/resource_group/controller/model_test.go b/client/resource_group/controller/model_test.go index 3e27534fd97..3051ce75660 100644 --- a/client/resource_group/controller/model_test.go +++ b/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) From fda5879c8459c8c1bfd4b1b85eca7854b2b2209e Mon Sep 17 00:00:00 2001 From: Yuhao Zhang Date: Wed, 8 Jul 2026 13:39:49 +0800 Subject: [PATCH 3/3] client/resource_group: clarify failed write payback refund Signed-off-by: Yuhao Zhang --- client/resource_group/controller/group_controller.go | 4 ++-- client/resource_group/controller/group_controller_test.go | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/client/resource_group/controller/group_controller.go b/client/resource_group/controller/group_controller.go index 7f4707755da..eb3e9697766 100644 --- a/client/resource_group/controller/group_controller.go +++ b/client/resource_group/controller/group_controller.go @@ -629,7 +629,7 @@ func (gc *groupCostController) onResponseImpl( if v := getRUValueFromConsumption(delta); v > 0 { counter.limiter.RemoveTokens(time.Now(), v) } else if v < 0 { - // Failed writes pay back the local write reservation. + // Failed writes refund the response-side write payback delta. counter.limiter.RefundTokens(time.Now(), -v) } } @@ -677,7 +677,7 @@ func (gc *groupCostController) onResponseWaitImpl( gc.metrics.successfulRequestDuration.Observe(d.Seconds()) waitDuration += d } else if v < 0 { - // Failed writes pay back the local write reservation. + // Failed writes refund the response-side write payback delta. gc.run.requestUnitTokens.limiter.RefundTokens(time.Now(), -v) gc.metrics.successfulRequestDuration.Observe(0) } else { diff --git a/client/resource_group/controller/group_controller_test.go b/client/resource_group/controller/group_controller_test.go index 545b0e61002..6886f61b19b 100644 --- a/client/resource_group/controller/group_controller_test.go +++ b/client/resource_group/controller/group_controller_test.go @@ -244,7 +244,7 @@ func TestFailedWriteReservationDoesNotEnterReportedConsumption(t *testing.T) { "failed write refund must not make actual usage negative") } -func TestFailedWriteReservationRefundsLimiterOnResponseWait(t *testing.T) { +func TestFailedWritePaybackRefundsLimiterOnResponseWait(t *testing.T) { re := require.New(t) gc := createTestGroupCostController(re) @@ -275,10 +275,10 @@ func TestFailedWriteReservationRefundsLimiterOnResponseWait(t *testing.T) { tokensAfterResponse := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) re.InDelta(tokensAfterReservation+expectedRefund, tokensAfterResponse, 1.0, - "failed write payback should refund the local write reservation") + "failed write payback should refund the local limiter") } -func TestFailedWriteReservationRefundsLimiterOnResponse(t *testing.T) { +func TestFailedWritePaybackRefundsLimiterOnResponse(t *testing.T) { re := require.New(t) gc := createTestGroupCostController(re) @@ -309,7 +309,7 @@ func TestFailedWriteReservationRefundsLimiterOnResponse(t *testing.T) { tokensAfterResponse := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) re.InDelta(tokensAfterReservation+expectedRefund, tokensAfterResponse, 1.0, - "failed write payback should refund the local write reservation") + "failed write payback should refund the local limiter") } func TestHandleTokenBucketUpdateEventCanceledByInitCounterNotify(t *testing.T) {