diff --git a/client/resource_group/controller/group_controller.go b/client/resource_group/controller/group_controller.go index 04ac7e8b77..eb3e969776 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,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 diff --git a/client/resource_group/controller/group_controller_test.go b/client/resource_group/controller/group_controller_test.go index c9e912d6ff..6886f61b19 100644 --- a/client/resource_group/controller/group_controller_test.go +++ b/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) diff --git a/client/resource_group/controller/limiter.go b/client/resource_group/controller/limiter.go index 9086a8026f..50e8b0f09a 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 01cc7fe197..d46c30f658 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 5472196981..b6ba755197 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 { + 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 diff --git a/client/resource_group/controller/model_test.go b/client/resource_group/controller/model_test.go index 3e27534fd9..3051ce7566 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)