From e2ca4b0d3f76a100a5c9390e580dbb4edfc486be Mon Sep 17 00:00:00 2001 From: Hitachi Date: Wed, 8 Jul 2026 15:54:07 +0800 Subject: [PATCH 1/2] client/resource_group: add RC paging pre-charge with PredictedReadBytes hint (#10611) close tikv/pd#10612, closes tikv/pd#10985\n\nSigned-off-by: Yuhao Zhang (cherry picked from commit 4e05b9d2c2d3c7814f519e1b43053ac0b66f467f) Signed-off-by: JmPotato --- client/go.mod | 2 +- .../controller/global_controller.go | 1 + .../controller/group_controller.go | 148 +++- .../controller/group_controller_test.go | 723 ++++++++++++++++++ client/resource_group/controller/limiter.go | 18 + .../resource_group/controller/limiter_test.go | 42 + .../controller/metrics/metrics.go | 72 +- client/resource_group/controller/model.go | 108 +++ .../resource_group/controller/model_test.go | 193 +++++ client/resource_group/controller/testutil.go | 22 +- .../resourcemanager/server/manager_test.go | 22 + .../resourcemanager/server/metering_test.go | 20 + pkg/mcs/resourcemanager/server/metrics.go | 6 + .../resourcemanager/server/metrics_test.go | 216 ++++++ pkg/metering/utils.go | 3 + pkg/metering/utils_test.go | 30 + 16 files changed, 1585 insertions(+), 41 deletions(-) create mode 100644 pkg/metering/utils_test.go diff --git a/client/go.mod b/client/go.mod index d5857dc01d7..6bbbda44888 100644 --- a/client/go.mod +++ b/client/go.mod @@ -13,6 +13,7 @@ require ( github.com/pingcap/kvproto v0.0.0-20260703015547-e800f9e5a427 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 github.com/prometheus/client_golang v1.20.5 + github.com/prometheus/client_model v0.6.1 github.com/stretchr/testify v1.9.0 go.uber.org/goleak v1.1.11 go.uber.org/zap v1.24.0 @@ -30,7 +31,6 @@ require ( github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/prometheus/client_model v0.6.1 // indirect github.com/prometheus/common v0.55.0 // indirect github.com/prometheus/procfs v0.15.1 // indirect github.com/stretchr/objx v0.5.2 // indirect diff --git a/client/resource_group/controller/global_controller.go b/client/resource_group/controller/global_controller.go index 17079a59774..70450e1dbbd 100644 --- a/client/resource_group/controller/global_controller.go +++ b/client/resource_group/controller/global_controller.go @@ -646,6 +646,7 @@ func (c *ResourceGroupsController) cleanUpResourceGroup() { metrics.WriteByteCost.DeleteLabelValues(resourceGroupName, refundDirection) metrics.KVCPUCost.DeleteLabelValues(resourceGroupName) metrics.SQLCPUCost.DeleteLabelValues(resourceGroupName) + gc.metrics.deletePagingLabels(resourceGroupName) return true } gc.inactive = true diff --git a/client/resource_group/controller/group_controller.go b/client/resource_group/controller/group_controller.go index 405b2e2ad36..2eba9f7b777 100644 --- a/client/resource_group/controller/group_controller.go +++ b/client/resource_group/controller/group_controller.go @@ -123,6 +123,14 @@ type groupMetricsCollection struct { refundWriteByteCostCounter prometheus.Counter kvCPUCostCounter prometheus.Counter sqlCPUCostCounter prometheus.Counter + + // Paging pre-charge observers, cached per-RG to avoid WithLabelValues + // on the hot path. + prechargeCounter prometheus.Counter + prechargeBytesCounter prometheus.Counter + actualBytesCounter prometheus.Counter + predictionResidualBytes prometheus.Observer + noPrechargeCounter prometheus.Counter } const ( @@ -160,7 +168,50 @@ func initMetrics(oldName, name string) *groupMetricsCollection { refundWriteByteCostCounter: metrics.WriteByteCost.WithLabelValues(name, refundDirection), kvCPUCostCounter: metrics.KVCPUCost.WithLabelValues(name), sqlCPUCostCounter: metrics.SQLCPUCost.WithLabelValues(name), + + prechargeCounter: metrics.CopReadPrechargeCounter.WithLabelValues(name), + prechargeBytesCounter: metrics.PagingPrechargeBytesCounter.WithLabelValues(name), + actualBytesCounter: metrics.PagingActualBytesCounter.WithLabelValues(name), + predictionResidualBytes: metrics.PagingPredictionResidualBytes.WithLabelValues(name), + + noPrechargeCounter: metrics.CopReadNoPrechargeCounter.WithLabelValues(name), + } +} + +// deletePagingLabels removes the per-resource-group paging_* metric series +// when the group is being deleted or tombstoned, so stale label series do +// not linger in Prometheus until the process restarts. Keep this list in +// sync with initMetrics — adding a paging metric there must be paired +// with a deletion here. +func (*groupMetricsCollection) deletePagingLabels(name string) { + metrics.CopReadPrechargeCounter.DeleteLabelValues(name) + metrics.CopReadNoPrechargeCounter.DeleteLabelValues(name) + metrics.PagingPrechargeBytesCounter.DeleteLabelValues(name) + metrics.PagingActualBytesCounter.DeleteLabelValues(name) + metrics.PagingPredictionResidualBytes.DeleteLabelValues(name) +} + +// observePagingRequest records request-boundary cop read pre-charge counters. +// Callers must gate the call on pagingReadEstimate(...).ok so the metric stays +// scoped to coprocessor reads and excludes point gets, batch gets, scans, and +// other bounded-size reads. +func (gmc *groupMetricsCollection) observePagingRequest(bytesForEst uint64) { + if bytesForEst == 0 { + gmc.noPrechargeCounter.Inc() + return + } + gmc.prechargeCounter.Inc() + gmc.prechargeBytesCounter.Add(float64(bytesForEst)) +} + +// observePagingResponse records response-boundary paging metrics for +// precharged coprocessor RPCs. +func (gmc *groupMetricsCollection) observePagingResponse(bytesForEst, actual uint64) { + if bytesForEst == 0 { + return } + gmc.actualBytesCounter.Add(float64(actual)) + gmc.predictionResidualBytes.Observe(float64(actual) - float64(bytesForEst)) } type tokenCounter struct { @@ -767,10 +818,11 @@ func (gc *groupCostController) onRequestWaitImpl( for _, calc := range gc.calculators { calc.BeforeKVRequest(delta, info) } + reportedDelta := reportedRequestConsumption(gc.calculators, info, delta) gc.observeDemand(delta) gc.mu.Lock() - add(gc.mu.consumption, delta) + add(gc.mu.consumption, reportedDelta) gc.mu.Unlock() if !gc.burstable.Load() { @@ -783,7 +835,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{} @@ -798,6 +850,9 @@ func (gc *groupCostController) onRequestWaitImpl( } gc.observeConsumption(delta) gc.observeComponentConsumption(delta) + if bytesForEst, ok := pagingReadEstimate(info); ok { + gc.metrics.observePagingRequest(bytesForEst) + } gc.mu.Lock() // Calculate the penalty of the store @@ -823,26 +878,32 @@ func (gc *groupCostController) onResponseImpl( for _, calc := range gc.calculators { calc.AfterKVRequest(delta, req, resp) } + reportedDelta := reportedResponseConsumption(gc.calculators, req, delta) + // `count` is the full per-request consumption (BeforeKVRequest + AfterKVRequest). + count := &rmpb.Consumption{} + *count = *delta + for _, calc := range gc.calculators { + calc.BeforeKVRequest(count, req) + } + bytesForEst, isPagingRead := pagingReadEstimate(req) + if isPagingRead { + gc.metrics.observePagingResponse(bytesForEst, resp.ReadBytes()) + } gc.observeDemand(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 && isPagingRead && bytesForEst > 0 { + // Paging over-estimate: refund the excess pre-charge. + counter.limiter.RefundTokens(time.Now(), -v) } } gc.observeConsumption(delta) gc.observeComponentConsumption(delta) gc.mu.Lock() - // Record the consumption of the request - add(gc.mu.consumption, delta) - // Record the consumption of the request by store - count := &rmpb.Consumption{} - *count = *delta - // As the penalty is only counted when the request is completed, so here needs to calculate the write cost which is added in `BeforeKVRequest` - for _, calc := range gc.calculators { - calc.BeforeKVRequest(count, req) - } + add(gc.mu.consumption, reportedDelta) add(gc.mu.storeCounter[req.StoreID()], count) add(gc.mu.globalCounter, count) gc.mu.Unlock() @@ -857,39 +918,58 @@ func (gc *groupCostController) onResponseWaitImpl( for _, calc := range gc.calculators { calc.AfterKVRequest(delta, req, resp) } + reportedDelta := reportedResponseConsumption(gc.calculators, req, delta) + // `count` is the full per-request consumption (BeforeKVRequest + AfterKVRequest). + count := &rmpb.Consumption{} + *count = *delta + for _, calc := range gc.calculators { + calc.BeforeKVRequest(count, req) + } + bytesForEst, isPagingRead := pagingReadEstimate(req) gc.observeDemand(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() - } - if waitDuration > slowNotifyFilterDuration { - log.Warn("[resource group controller] response wait for tokens failed", gc.logFields(waitDuration, err)...) + v := getRUValueFromConsumption(delta) + if v > 0 && isPagingRead && bytesForEst > 0 { + // The precharged request has already consumed TiKV resources. + // Debit any positive settlement delta immediately so future + // requests observe the debt instead of queueing this completed + // request for response-side admission. + gc.run.requestUnitTokens.limiter.RemoveTokens(time.Now(), v) + gc.metrics.successfulRequestDuration.Observe(0) + } else 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() + } + if waitDuration > slowNotifyFilterDuration { + log.Warn("[resource group controller] response wait for tokens failed", gc.logFields(waitDuration, err)...) + } + return nil, waitDuration, err } - return nil, waitDuration, err + gc.metrics.successfulRequestDuration.Observe(d.Seconds()) + waitDuration += d + } else if v < 0 && isPagingRead && bytesForEst > 0 { + // Paging over-estimate: refund the excess pre-charge. + 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.observeConsumption(delta) gc.observeComponentConsumption(delta) + if isPagingRead { + gc.metrics.observePagingResponse(bytesForEst, resp.ReadBytes()) + } gc.mu.Lock() - // Record the consumption of the request - add(gc.mu.consumption, delta) - // Record the consumption of the request by store - count := &rmpb.Consumption{} - *count = *delta - // As the penalty is only counted when the request is completed, so here needs to calculate the write cost which is added in `BeforeKVRequest` - for _, calc := range gc.calculators { - calc.BeforeKVRequest(count, req) - } + add(gc.mu.consumption, reportedDelta) add(gc.mu.storeCounter[req.StoreID()], count) add(gc.mu.globalCounter, count) gc.mu.Unlock() diff --git a/client/resource_group/controller/group_controller_test.go b/client/resource_group/controller/group_controller_test.go index 643192a7e7b..9c278cdfbef 100644 --- a/client/resource_group/controller/group_controller_test.go +++ b/client/resource_group/controller/group_controller_test.go @@ -25,6 +25,7 @@ import ( "github.com/prometheus/client_golang/prometheus" promtestutil "github.com/prometheus/client_golang/prometheus/testutil" + dto "github.com/prometheus/client_model/go" "github.com/stretchr/testify/require" "github.com/pingcap/failpoint" @@ -34,6 +35,28 @@ import ( "github.com/tikv/pd/client/resource_group/controller/metrics" ) +func counterValue(re *require.Assertions, c interface{ Write(*dto.Metric) error }) float64 { + var m dto.Metric + re.NoError(c.Write(&m)) + return m.GetCounter().GetValue() +} + +func histogramSampleCount(re *require.Assertions, h prometheus.Observer) uint64 { + metric, ok := h.(prometheus.Metric) + re.True(ok) + if !ok { + return 0 + } + var m dto.Metric + re.NoError(metric.Write(&m)) + histogram := m.GetHistogram() + re.NotNil(histogram) + if histogram == nil { + return 0 + } + return histogram.GetSampleCount() +} + func createTestGroupCostController(re *require.Assertions, names ...string) *groupCostController { name := "test" if len(names) > 0 { @@ -91,6 +114,38 @@ func TestGroupControlBurstable(t *testing.T) { re.True(gc.burstable.Load()) } +func TestSmallPositiveConsumptionDoesNotBypassThreshold(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + now := time.Now() + + gc.run.consumption.RRU = 1 + gc.run.initialRequestCompleted = true + gc.run.lastRequestTime = now.Add(-defaultTargetPeriod) + gc.run.now = now + + report := gc.collectRequestAndConsumption(periodicReport) + re.Nil(report) +} + +func TestDirectReportConsumptionReportsOnlyConsumption(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + gc.addRUConsumption(&rmpb.Consumption{ + RRU: 7, + WRU: 3, + WriteBytes: 1024, + }) + gc.updateRunState() + + report := gc.collectRequestAndConsumption(periodicReport) + re.NotNil(report) + re.Equal(float64(7), report.GetConsumptionSinceLastRequest().GetRRU()) + re.Equal(float64(3), report.GetConsumptionSinceLastRequest().GetWRU()) + re.Equal(float64(1024), report.GetConsumptionSinceLastRequest().GetWriteBytes()) +} + func TestRequestAndResponseConsumption(t *testing.T) { re := require.New(t) gc := createTestGroupCostController(re, t.Name()) @@ -237,6 +292,649 @@ func TestOnResponseWaitConsumption(t *testing.T) { verify() } +func TestOnResponseWaitZeroSettlementObservesSuccessfulDuration(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + re.False(gc.burstable.Load()) + + samplesBefore := histogramSampleCount(re, gc.metrics.successfulRequestDuration) + _, waitTime, err := gc.onResponseWaitImpl(context.TODO(), &TestRequestInfo{isWrite: false}, &TestResponseInfo{succeed: true}) + re.NoError(err) + re.Zero(waitTime) + re.Equal(samplesBefore+1, histogramSampleCount(re, gc.metrics.successfulRequestDuration)) +} + +func TestPredictedReadBytesPreCharge(t *testing.T) { + re := require.New(t) + cfg := DefaultRUConfig() + kvCalc := newKVCalculator(cfg) + + // BeforeKVRequest with a PredictedReadBytes hint should pre-charge + // baseCost + predictedReadBytes * ReadBytesCost. + predictedReadBytes := uint64(256 * 1024) // learned EMA estimate + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: predictedReadBytes, + } + precharge := &rmpb.Consumption{} + kvCalc.BeforeKVRequest(precharge, req) + + baseCost := float64(cfg.ReadBaseCost) + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion + hintCost := float64(cfg.ReadBytesCost) * float64(predictedReadBytes) + re.InDelta(baseCost+hintCost, precharge.RRU, 1e-6, + "BeforeKVRequest should pre-charge based on PredictedReadBytes") + + // AfterKVRequest should subtract the same hint basis, preserving + // precharge + settle == baseCost + actualCost. + actualReadBytes := uint64(300 * 1024) // close to the prediction + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + kvCPU: 10 * time.Millisecond, + succeed: true, + } + settle := &rmpb.Consumption{} + kvCalc.AfterKVRequest(settle, req, resp) + + actualReadCost := float64(cfg.ReadBytesCost) * float64(actualReadBytes) + cpuCost := float64(cfg.CPUMsCost) * 10.0 + re.InDelta(actualReadCost+cpuCost-hintCost, settle.RRU, 1e-6, + "AfterKVRequest should settle using the same hint basis as BeforeKVRequest") + + totalRRU := precharge.RRU + settle.RRU + re.InDelta(baseCost+actualReadCost+cpuCost, totalRRU, 1e-6, + "Total RRU across precharge+settle should equal baseCost + actualCost") +} + +func TestPagingPrechargeDoesNotEnterReportedRequestConsumption(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + cfg := DefaultRUConfig() + kvCalc := gc.getKVCalculator() + + predictedReadBytes := uint64(4 * 1024 * 1024) + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: predictedReadBytes, + } + + tokenDelta, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + + baseCost := float64(kvCalc.ReadBaseCost) + float64(kvCalc.ReadPerBatchBaseCost)*defaultAvgBatchProportion + prechargeCost := float64(cfg.ReadBytesCost) * float64(predictedReadBytes) + re.InDelta(baseCost+prechargeCost, tokenDelta.RRU, 1e-6, + "request token ledger should still include paging pre-charge") + + gc.updateRunState() + report := gc.collectRequestAndConsumption(periodicReport) + re.NotNil(report) + re.InDelta(baseCost, report.GetConsumptionSinceLastRequest().GetRRU(), 1e-6, + "reported request consumption must keep metering aligned with master and exclude predicted-read reservation") +} + +func TestPagingPrechargeRefundDoesNotEnterReportedResponseConsumption(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + cfg := DefaultRUConfig() + + predictedReadBytes := uint64(4 * 1024 * 1024) + actualReadBytes := uint64(512 * 1024) + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: predictedReadBytes, + } + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + succeed: true, + } + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + gc.updateRunState() + requestReport := gc.collectRequestAndConsumption(periodicReport) + re.NotNil(requestReport) + gc.handleTokenBucketResponse(&rmpb.TokenBucketResponse{}) + + tokenSettlement, _, err := gc.onResponseWaitImpl(context.TODO(), req, resp) + re.NoError(err) + expectedTokenSettlement := float64(cfg.ReadBytesCost) * (float64(actualReadBytes) - float64(predictedReadBytes)) + re.InDelta(expectedTokenSettlement, tokenSettlement.RRU, 1e-6, + "response token ledger should still settle against the request pre-charge") + re.Negative(tokenSettlement.RRU) + + gc.run.lastRequestTime = time.Now().Add(-extendedReportingPeriodFactor * defaultTargetPeriod) + gc.updateRunState() + responseReport := gc.collectRequestAndConsumption(periodicReport) + re.NotNil(responseReport) + expectedReportedResponseRU := float64(cfg.ReadBytesCost) * float64(actualReadBytes) + re.InDelta(expectedReportedResponseRU, responseReport.GetConsumptionSinceLastRequest().GetRRU(), 1e-6, + "reported response consumption must expose actual read usage, not the paging refund") + re.GreaterOrEqual(responseReport.GetConsumptionSinceLastRequest().GetRRU(), 0.0) +} + +func TestPredictedReadBytesRequiresCopRead(t *testing.T) { + re := require.New(t) + cfg := DefaultRUConfig() + kvCalc := newKVCalculator(cfg) + + predictedReadBytes := uint64(256 * 1024) + req := &TestRequestInfo{ + isWrite: false, + isCop: false, + predictedReadBytes: predictedReadBytes, + } + precharge := &rmpb.Consumption{} + kvCalc.BeforeKVRequest(precharge, req) + + baseCost := float64(cfg.ReadBaseCost) + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion + re.InDelta(baseCost, precharge.RRU, 1e-6, + "non-cop reads must not consume the paging predicted-read hint") + + actualReadBytes := uint64(300 * 1024) + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + kvCPU: 10 * time.Millisecond, + succeed: true, + } + settle := &rmpb.Consumption{} + kvCalc.AfterKVRequest(settle, req, resp) + + actualReadCost := float64(cfg.ReadBytesCost) * float64(actualReadBytes) + cpuCost := float64(cfg.CPUMsCost) * 10.0 + re.InDelta(actualReadCost+cpuCost, settle.RRU, 1e-6, + "non-cop reads must not settle against the paging predicted-read hint") +} + +func TestPagingPreChargeTokenRefund(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + initialTokens := float64(100000) + gc.run.requestUnitTokens.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{ + newTokens: initialTokens, + newFillRate: 0, + newBurst: 0, + }) + + predictedReadBytes := uint64(4 * 1024 * 1024) // 4 MB pre-charge + actualReadBytes := uint64(1 * 1024 * 1024) // 1 MB actual + + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: predictedReadBytes, + } + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + succeed: true, + } + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + tokensAfterPreCharge := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + samplesBefore := histogramSampleCount(re, gc.metrics.successfulRequestDuration) + _, _, err = gc.onResponseWaitImpl(context.TODO(), req, resp) + re.NoError(err) + tokensAfterSettlement := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + re.Equal(samplesBefore+1, histogramSampleCount(re, gc.metrics.successfulRequestDuration)) + + // Refund (pre-charge - actual) should exceed the actual read cost, + // so the limiter ends settlement with more tokens than after pre-charge. + cfg := DefaultRUConfig() + preChargeCost := float64(cfg.ReadBytesCost) * float64(predictedReadBytes) + actualCost := float64(cfg.ReadBytesCost) * float64(actualReadBytes) + expectedRefund := preChargeCost - actualCost + re.Positive(expectedRefund, "sanity: pre-charge should exceed actual cost") + re.InDelta(tokensAfterPreCharge+expectedRefund, tokensAfterSettlement, 1.0, + "limiter should be refunded the excess pre-charged tokens") + + gc.mu.Lock() + netRRU := gc.mu.consumption.RRU + gc.mu.Unlock() + baseCost := float64(cfg.ReadBaseCost) + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion + re.InDelta(baseCost+actualCost, netRRU, 1e-6, + "net consumption should equal baseCost + actualReadCost") +} + +func TestPagingPreChargeNoRefundWhenActualExceedsEstimate(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + initialTokens := float64(100000) + gc.run.requestUnitTokens.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{ + newTokens: initialTokens, + newFillRate: 0, + newBurst: 0, + }) + + predictedReadBytes := uint64(1 * 1024 * 1024) // 1 MB pre-charge + actualReadBytes := uint64(4 * 1024 * 1024) // 4 MB actual (exceeds estimate) + + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: predictedReadBytes, + } + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + succeed: true, + } + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + tokensAfterPreCharge := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + _, _, err = gc.onResponseWaitImpl(context.TODO(), req, resp) + re.NoError(err) + tokensAfterSettlement := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + re.Less(tokensAfterSettlement, tokensAfterPreCharge, + "when actual exceeds pre-charge, settlement should consume tokens") +} + +func TestOnResponseWaitPrechargedPositiveSettlementDebitsImmediately(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + gc.mainCfg.WaitRetryTimes = 1 + gc.mainCfg.WaitRetryInterval = time.Millisecond + gc.mainCfg.LTBMaxWaitDuration = time.Millisecond + + limiter := gc.run.requestUnitTokens.limiter + limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{ + newTokens: 0, + newFillRate: 0, + newBurst: 0, + }) + limiter.RemoveTokens(time.Now(), limiter.AvailableTokens(time.Now())) + gc.isThrottled.Store(true) + + req := &TestRequestInfo{ + isWrite: false, + predictedReadBytes: 16 * 1024 * 1024, + isCop: true, + } + resp := &TestResponseInfo{ + readBytes: 20 * 1024 * 1024, + succeed: true, + } + settlementDelta := &rmpb.Consumption{} + for _, calc := range gc.calculators { + calc.AfterKVRequest(settlementDelta, req, resp) + } + re.Greater(getRUValueFromConsumption(settlementDelta), float64(0)) + re.GreaterOrEqual(settlementDelta.ReadBytes+settlementDelta.WriteBytes, float64(bigRequestThreshold)) + re.True(gc.isThrottled.Load()) + + tokensBefore := limiter.AvailableTokens(time.Now()) + samplesBefore := histogramSampleCount(re, gc.metrics.successfulRequestDuration) + _, waitDuration, err := gc.onResponseWaitImpl(context.Background(), req, resp) + re.NoError(err) + re.Zero(waitDuration) + + tokensAfter := limiter.AvailableTokens(time.Now()) + re.Less(tokensAfter, tokensBefore, + "precharged positive settlement should be debited immediately") + re.Equal(samplesBefore+1, histogramSampleCount(re, gc.metrics.successfulRequestDuration)) +} + +func TestOnResponseImplPagingRefund(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + initialTokens := float64(100000) + gc.run.requestUnitTokens.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{ + newTokens: initialTokens, + newFillRate: 0, + newBurst: 0, + }) + + predictedReadBytes := uint64(4 * 1024 * 1024) // 4 MB pre-charge + actualReadBytes := uint64(512 * 1024) // 512 KB actual + + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: predictedReadBytes, + } + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + succeed: true, + } + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + tokensAfterPreCharge := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + _, err = gc.onResponseImpl(req, resp) + re.NoError(err) + tokensAfterSettlement := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + re.Greater(tokensAfterSettlement, tokensAfterPreCharge, + "onResponseImpl should refund excess pre-charged tokens") +} + +func TestNegativePagingSettlementDoesNotForceConsumptionReport(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + gc.run.requestUnitTokens.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{ + newTokens: 100000, + newFillRate: 0, + newBurst: 0, + }) + + predictedReadBytes := uint64(4 * 1024 * 1024) + actualReadBytes := uint64(512 * 1024) + req := &TestRequestInfo{ + isWrite: false, + predictedReadBytes: predictedReadBytes, + isCop: true, + } + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + succeed: true, + } + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + gc.updateRunState() + prechargeReport := gc.collectRequestAndConsumption(periodicReport) + re.NotNil(prechargeReport) + gc.handleTokenBucketResponse(&rmpb.TokenBucketResponse{}) + + settlementDelta, _, err := gc.onResponseWaitImpl(context.TODO(), req, resp) + re.NoError(err) + gc.run.lastRequestTime = time.Now().Add(-defaultTargetPeriod) + gc.updateRunState() + report := gc.collectRequestAndConsumption(periodicReport) + + cfg := DefaultRUConfig() + expectedRRU := float64(cfg.ReadBytesCost) * (float64(actualReadBytes) - float64(predictedReadBytes)) + re.InDelta(expectedRRU, settlementDelta.RRU, 1e-6) + re.Nil(report, "negative token settlement alone should not force a consumption report") +} + +func TestPagingPreChargeRefundOnFailedRead(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + initialTokens := float64(100000) + gc.run.requestUnitTokens.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{ + newTokens: initialTokens, + newFillRate: 0, + newBurst: 0, + }) + + predictedReadBytes := uint64(4 * 1024 * 1024) // 4 MB pre-charge + + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: predictedReadBytes, + } + // Failed response: no bytes read, no CPU consumed, succeed=false. + resp := &TestResponseInfo{ + readBytes: 0, + kvCPU: 0, + succeed: false, + } + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + tokensAfterPreCharge := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + _, _, err = gc.onResponseWaitImpl(context.TODO(), req, resp) + re.NoError(err) + tokensAfterSettlement := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + // On a failed read, AfterKVRequest still runs: paging settlement subtracts + // ReadBytesCost*predicted and calculateReadCost adds 0, yielding a negative + // delta that flows through RefundTokens. ReadBaseCost is not refunded, + // matching existing behavior for non-paging read failures. + cfg := DefaultRUConfig() + expectedRefund := float64(cfg.ReadBytesCost) * float64(predictedReadBytes) + re.InDelta(tokensAfterPreCharge+expectedRefund, tokensAfterSettlement, 1.0, + "failed read with paging hint should refund ReadBytesCost*predicted") +} + +func TestFailedWriteDoesNotUsePagingRefundPath(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) + re.Negative(refundDelta.WRU) + tokensAfterResponse := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + re.InDelta(tokensAfterReservation, tokensAfterResponse, 1.0, + "feature PR should only refund paging over-estimates, not failed-write reservations") +} + +func TestFailedWriteOnResponseDoesNotUsePagingRefundPath(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) + re.Negative(refundDelta.WRU) + tokensAfterResponse := gc.run.requestUnitTokens.limiter.AvailableTokens(time.Now()) + + re.InDelta(tokensAfterReservation, tokensAfterResponse, 1.0, + "feature PR should only refund paging over-estimates, not failed-write reservations") +} + +func TestNonCopPredictedReadBytesResponseIgnoresPagingAccounting(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + predictedReadBytes := uint64(4 * 1024 * 1024) + actualReadBytes := uint64(512 * 1024) + req := &TestRequestInfo{ + isWrite: false, + predictedReadBytes: predictedReadBytes, + isCop: false, + } + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + kvCPU: 10 * time.Millisecond, + succeed: true, + } + + prechargeBefore := counterValue(re, gc.metrics.prechargeCounter) + actualBefore := counterValue(re, gc.metrics.actualBytesCounter) + noPrechargeBefore := counterValue(re, gc.metrics.noPrechargeCounter) + delta, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + cfg := DefaultRUConfig() + baseCost := float64(cfg.ReadBaseCost) + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion + re.InDelta(baseCost, delta.RRU, 1e-6, + "non-cop read hints must not add predicted read bytes to pre-charge") + + settlement, err := gc.onResponseImpl(req, resp) + re.NoError(err) + actualReadCost := float64(cfg.ReadBytesCost) * float64(actualReadBytes) + cpuCost := float64(cfg.CPUMsCost) * 10.0 + re.InDelta(actualReadCost+cpuCost, settlement.RRU, 1e-6, + "non-cop response settlement must bill actual read bytes and CPU without subtracting the ignored hint") + re.InDelta(prechargeBefore, counterValue(re, gc.metrics.prechargeCounter), 1e-9) + re.InDelta(actualBefore, counterValue(re, gc.metrics.actualBytesCounter), 1e-9) + re.InDelta(noPrechargeBefore, counterValue(re, gc.metrics.noPrechargeCounter), 1e-9) +} + +func TestDeletePagingLabelsResetsSeries(t *testing.T) { + re := require.New(t) + // Use a name no other test reuses so we measure only our own series. + name := "test-paging-cleanup-rg" + + gmc := initMetrics(name, name) + // Push a sample through every paging counter / histogram so each label + // series actually exists in the underlying Vec. + gmc.observePagingRequest(100) + gmc.observePagingResponse(100, 80) + gmc.observePagingRequest(0) + gmc.observePagingResponse(0, 200) + + // Sanity: cached counters are non-zero before cleanup. + re.Positive(counterValue(re, gmc.prechargeCounter)) + re.Positive(counterValue(re, gmc.actualBytesCounter)) + re.Positive(counterValue(re, gmc.noPrechargeCounter)) + re.Positive(histogramSampleCount(re, gmc.predictionResidualBytes)) + + gmc.deletePagingLabels(name) + + // After cleanup, refetching each Vec with the same label must yield a + // fresh zero-valued series. Covers every counter declared in + // initMetrics so a forgotten DeleteLabelValues in deletePagingLabels + // surfaces here. + for _, vec := range []*prometheus.CounterVec{ + metrics.CopReadPrechargeCounter, + metrics.CopReadNoPrechargeCounter, + metrics.PagingPrechargeBytesCounter, + metrics.PagingActualBytesCounter, + } { + re.Zero(counterValue(re, vec.WithLabelValues(name)), + "paging counter series for %q should be cleared by deletePagingLabels", name) + } + for _, vec := range []*prometheus.HistogramVec{ + metrics.PagingPredictionResidualBytes, + } { + re.Zero(histogramSampleCount(re, vec.WithLabelValues(name)), + "paging histogram series for %q should be cleared by deletePagingLabels", name) + } +} + +func TestCopReadNoPrechargeGatedByIsCop(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + // Reads without a pre-charge hint reach onResponseImpl through the same RC + // interceptor regardless of cmd type (CmdGet, CmdBatchGet, CmdCop, ...). + // Only IsCop()==true requests should land in cop_read_no_precharge_total; + // the rest must be ignored so the metric keeps its cop-read coverage + // semantics. + resp := &TestResponseInfo{readBytes: 1024, succeed: true} + nonCop := &TestRequestInfo{isWrite: false, isCop: false} + cop := &TestRequestInfo{isWrite: false, isCop: true} + + counterBefore := counterValue(re, gc.metrics.noPrechargeCounter) + residualSamplesBefore := histogramSampleCount(re, gc.metrics.predictionResidualBytes) + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), nonCop) + re.NoError(err) + + _, err = gc.onResponseImpl(nonCop, resp) + re.NoError(err) + re.InDelta(counterBefore, counterValue(re, gc.metrics.noPrechargeCounter), 1e-9, + "non-cop reads must not increment cop_read_no_precharge_total") + re.Equal(residualSamplesBefore, histogramSampleCount(re, gc.metrics.predictionResidualBytes), + "non-cop reads must not be included in paging prediction residuals") + + _, _, _, _, err = gc.onRequestWaitImpl(context.TODO(), cop) + re.NoError(err) + re.InDelta(counterBefore+1, counterValue(re, gc.metrics.noPrechargeCounter), 1e-9, + "cop reads without a hint must increment cop_read_no_precharge_total on the request path") + re.Equal(residualSamplesBefore, histogramSampleCount(re, gc.metrics.predictionResidualBytes), + "cop reads without a hint must not record prediction residuals before response") + + _, err = gc.onResponseImpl(cop, resp) + re.NoError(err) + re.InDelta(counterBefore+1, counterValue(re, gc.metrics.noPrechargeCounter), 1e-9, + "cop read responses must not increment cop_read_no_precharge_total again") + re.Equal(residualSamplesBefore, histogramSampleCount(re, gc.metrics.predictionResidualBytes), + "cop read responses without a hint must not be included in precharge prediction residuals") +} + +func TestCopReadNoPrechargeCounterObservedOnRequest(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + req := &TestRequestInfo{isWrite: false, isCop: true} + + counterBefore := counterValue(re, gc.metrics.noPrechargeCounter) + residualSamplesBefore := histogramSampleCount(re, gc.metrics.predictionResidualBytes) + + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.NoError(err) + + re.InDelta(counterBefore+1, counterValue(re, gc.metrics.noPrechargeCounter), 1e-9, + "cop reads without a prediction should be counted at the same request boundary as precharged reads") + re.Equal(residualSamplesBefore, histogramSampleCount(re, gc.metrics.predictionResidualBytes), + "prediction residuals are response-only and must not be recorded on the request path") +} + +func TestNoPreChargeWithoutPredictedReadBytes(t *testing.T) { + re := require.New(t) + cfg := DefaultRUConfig() + kvCalc := newKVCalculator(cfg) + baseCost := float64(cfg.ReadBaseCost) + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion + + // Without a PredictedReadBytes hint, BeforeKVRequest must not + // pre-charge; AfterKVRequest bills actual read bytes only. + req := &TestRequestInfo{isWrite: false} + precharge := &rmpb.Consumption{} + kvCalc.BeforeKVRequest(precharge, req) + re.InDelta(baseCost, precharge.RRU, 1e-6, + "Without a hint, BeforeKVRequest should only charge baseCost") + + actualReadBytes := uint64(2 * 1024 * 1024) + resp := &TestResponseInfo{ + readBytes: actualReadBytes, + kvCPU: 10 * time.Millisecond, + succeed: true, + } + settle := &rmpb.Consumption{} + kvCalc.AfterKVRequest(settle, req, resp) + + actualReadCost := float64(cfg.ReadBytesCost) * float64(actualReadBytes) + cpuCost := float64(cfg.CPUMsCost) * 10.0 + re.InDelta(actualReadCost+cpuCost, settle.RRU, 1e-6, + "AfterKVRequest should bill actual read cost only when nothing was pre-charged") +} + func TestHandleTokenBucketUpdateEventCanceledByInitCounterNotify(t *testing.T) { re := require.New(t) gc := createTestGroupCostController(re) @@ -327,6 +1025,31 @@ func TestResourceGroupThrottledError(t *testing.T) { re.True(errs.ErrClientResourceGroupThrottled.Equal(err)) } +func TestPagingPrechargeNotObservedOnThrottle(t *testing.T) { + re := require.New(t) + gc := createTestGroupCostController(re) + + // 4 GiB predicted ~= 65536 RU at default ReadCostPerByte (1/64KiB), + // exceeding the LTBMaxWaitDuration budget at the default FillRate so + // acquireTokens returns ErrClientResourceGroupThrottled. + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: 4 * 1024 * 1024 * 1024, + } + + before := counterValue(re, gc.metrics.prechargeCounter) + _, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req) + re.Error(err) + re.True(errs.ErrClientResourceGroupThrottled.Equal(err)) + after := counterValue(re, gc.metrics.prechargeCounter) + + // Throttled requests never reach OnResponse for settlement, so the + // precharge counter must not be incremented either. + re.Equal(before, after, + "throttled paging request should not inflate CopReadPrechargeCounter") +} + func TestAcquireTokensSignalAwareWait(t *testing.T) { re := require.New(t) diff --git a/client/resource_group/controller/limiter.go b/client/resource_group/controller/limiter.go index 658085d97dc..1907cde7cf2 100644 --- a/client/resource_group/controller/limiter.go +++ b/client/resource_group/controller/limiter.go @@ -333,6 +333,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 — +// 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..9c774ebb0c8 100644 --- a/client/resource_group/controller/limiter_test.go +++ b/client/resource_group/controller/limiter_test.go @@ -163,6 +163,14 @@ func TestLastStaysMonotonicOnStaleNow(t *testing.T) { expectedPool: 1090, checkPool: true, }, + { + name: "refund tokens", + mutate: func(lim *Limiter, _ *Reservation, staleNow time.Time) { + lim.RefundTokens(staleNow, 0) + }, + expectedPool: 1090, + checkPool: true, + }, // CancelAt is covered by TestAcquireTokensCancelKeepsLastMonotonic. { name: "reconfigure", @@ -219,6 +227,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/metrics/metrics.go b/client/resource_group/controller/metrics/metrics.go index 5981c11e38b..8a9aec18731 100644 --- a/client/resource_group/controller/metrics/metrics.go +++ b/client/resource_group/controller/metrics/metrics.go @@ -80,6 +80,22 @@ var ( KVCPUCost *prometheus.CounterVec // SQLCPUCost tracks client-side SQL CPU time in milliseconds per resource group. SQLCPUCost *prometheus.CounterVec + + // CopReadPrechargeCounter counts read coprocessor RPCs that carried a positive + // PredictedReadBytes hint and triggered RC read-byte pre-charge. + CopReadPrechargeCounter *prometheus.CounterVec + // CopReadNoPrechargeCounter counts read coprocessor RPCs without a positive + // PredictedReadBytes hint, so request-side read-byte pre-charge was skipped. + CopReadNoPrechargeCounter *prometheus.CounterVec + + // PagingPrechargeBytesCounter accumulates bytes used as the RC paging pre-charge basis + // (sum of predicted hints from coprocessor RPCs). + PagingPrechargeBytesCounter *prometheus.CounterVec + // PagingActualBytesCounter accumulates actual bytes read by pre-charged coprocessor RPCs. + PagingActualBytesCounter *prometheus.CounterVec + // PagingPredictionResidualBytes records the distribution of (actual - predicted) read + // bytes for pre-charged coprocessor RPCs. + PagingPredictionResidualBytes *prometheus.HistogramVec ) func init() { @@ -183,7 +199,6 @@ func initMetrics(constLabels prometheus.Labels) { // WithLabelValues is a heavy operation, define variable to avoid call it every time. FailedTokenRequestDuration = TokenRequestDuration.WithLabelValues("fail") SuccessfulTokenRequestDuration = TokenRequestDuration.WithLabelValues("success") - TokenConsumedByTypeCounter = prometheus.NewCounterVec( prometheus.CounterOpts{ Namespace: namespace, @@ -291,6 +306,56 @@ func initMetrics(constLabels prometheus.Labels) { Help: "Counter of the SQL CPU time cost in milliseconds for all resource groups on the client side.", ConstLabels: constLabels, }, []string{newResourceGroupNameLabel}) + + CopReadPrechargeCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: requestSubsystem, + Name: "cop_read_precharge_total", + Help: "Counter of read coprocessor RPCs that carried a positive PredictedReadBytes hint and triggered request-side read-byte pre-charge.", + ConstLabels: constLabels, + }, []string{newResourceGroupNameLabel}) + + CopReadNoPrechargeCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: requestSubsystem, + Name: "cop_read_no_precharge_total", + Help: "Counter of read coprocessor RPCs that reached the RC interceptor without a positive PredictedReadBytes hint, so request-side read-byte pre-charge was skipped.", + ConstLabels: constLabels, + }, []string{newResourceGroupNameLabel}) + + PagingPrechargeBytesCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: requestSubsystem, + Name: "paging_precharge_bytes_total", + Help: "Sum of bytes used as the RC paging pre-charge basis (predicted hint from coprocessor RPCs).", + ConstLabels: constLabels, + }, []string{newResourceGroupNameLabel}) + + PagingActualBytesCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: requestSubsystem, + Name: "paging_actual_bytes_total", + Help: "Sum of actual bytes read by pre-charged coprocessor RPCs.", + ConstLabels: constLabels, + }, []string{newResourceGroupNameLabel}) + + PagingPredictionResidualBytes = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Subsystem: requestSubsystem, + Name: "paging_prediction_residual_bytes", + // Signed residual = actual - predicted for a single coprocessor + // RPC response. Buckets span ±64MB to cover per-RPC prediction + // residuals in both directions while factor-4 spacing keeps + // resolution near zero. + Buckets: []float64{-67108864, -16777216, -4194304, -1048576, -262144, -65536, -16384, -4096, 0, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864}, + Help: "Histogram of signed (actual_read_bytes - predicted_read_bytes) for pre-charged coprocessor RPCs. Use bucket series for distribution/quantiles; _sum is non-monotonic because observations can be negative.", + ConstLabels: constLabels, + }, []string{newResourceGroupNameLabel}) } // InitAndRegisterMetrics initializes and register metrics. @@ -318,4 +383,9 @@ func InitAndRegisterMetrics(constLabels prometheus.Labels) { prometheus.MustRegister(WriteByteCost) prometheus.MustRegister(KVCPUCost) prometheus.MustRegister(SQLCPUCost) + prometheus.MustRegister(CopReadPrechargeCounter) + prometheus.MustRegister(CopReadNoPrechargeCounter) + prometheus.MustRegister(PagingPrechargeBytesCounter) + prometheus.MustRegister(PagingActualBytesCounter) + prometheus.MustRegister(PagingPredictionResidualBytes) } diff --git a/client/resource_group/controller/model.go b/client/resource_group/controller/model.go index 54721969811..68edc1a1380 100644 --- a/client/resource_group/controller/model.go +++ b/client/resource_group/controller/model.go @@ -45,6 +45,8 @@ const ( // RequestInfo is the interface of the request information provider. A request should be // able to tell whether it's a write request and if so, the written bytes would also be provided. +// Implementations may additionally provide PredictedReadBytes() and IsCop() methods to opt +// into paging pre-charge accounting; missing methods are treated as no hint and non-cop. type RequestInfo interface { IsWrite() bool WriteBytes() uint64 @@ -54,6 +56,57 @@ type RequestInfo interface { AccessLocationType() AccessLocationType } +type predictedReadBytesProvider interface { + // PredictedReadBytes returns the caller-supplied read-bytes hint. The + // controller only uses it for coprocessor reads; non-cop hints are + // ignored by paging accounting. + PredictedReadBytes() uint64 +} + +type copRequestInfo interface { + // IsCop reports whether this request targets the coprocessor endpoint + // (CmdCop / CmdCopStream). Only coprocessor reads participate in paging + // pre-charge, settlement, and metrics; point gets, batch gets, scans and + // other bounded-size reads bypass paging even when they carry a + // PredictedReadBytes hint. + IsCop() bool +} + +func predictedReadBytes(req RequestInfo) uint64 { + provider, ok := req.(predictedReadBytesProvider) + if !ok { + return 0 + } + return provider.PredictedReadBytes() +} + +func isCopRequest(req RequestInfo) bool { + copReq, ok := req.(copRequestInfo) + return ok && copReq.IsCop() +} + +// pagingReadEstimate returns the predicted read-bytes hint for coprocessor +// read requests. The boolean reports whether the request participates in +// cop-read pre-charge accounting at all; a true result with 0 bytes means +// the cop read did not carry a usable pre-charge hint. +func pagingReadEstimate(req RequestInfo) (uint64, bool) { + if req.IsWrite() || !isCopRequest(req) { + return 0, false + } + return predictedReadBytes(req), true +} + +// estimatedReadBytes returns the predicted read-bytes hint for coprocessor +// read requests. Writes and non-coprocessor reads always return 0 so cop-read +// pre-charge and settlement stay gated to coprocessor RPCs. +func estimatedReadBytes(req RequestInfo) uint64 { + bytesForEst, ok := pagingReadEstimate(req) + if !ok { + return 0 + } + return bytesForEst +} + // ResponseInfo is the interface of the response information provider. A response should be // able to tell how many bytes it read and KV CPU cost in milliseconds. type ResponseInfo interface { @@ -104,6 +157,10 @@ func (kc *KVCalculator) BeforeKVRequest(consumption *rmpb.Consumption, req Reque // Read bytes could not be known before the request is executed, // so we only add the base cost here. consumption.RRU += float64(kc.ReadBaseCost) + float64(kc.ReadPerBatchBaseCost)*defaultAvgBatchProportion + // Paging pre-charge + if bytesForEst := estimatedReadBytes(req); bytesForEst > 0 { + consumption.RRU += float64(kc.ReadBytesCost) * float64(bytesForEst) + } } if req.AccessLocationType() == AccessCrossZone { if req.IsWrite() { @@ -138,6 +195,10 @@ func (kc *KVCalculator) AfterKVRequest(consumption *rmpb.Consumption, req Reques if !req.IsWrite() { // For now, we can only collect the KV CPU cost for a read request. kc.calculateCPUCost(consumption, res) + // Paging settlement + if bytesForEst := estimatedReadBytes(req); bytesForEst > 0 { + consumption.RRU -= float64(kc.ReadBytesCost) * float64(bytesForEst) + } } else if !res.Succeed() { // If the write request is not successfully returned, we need to pay back the WRU cost. kc.payBackWriteCost(consumption, req) @@ -185,6 +246,53 @@ 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 (kc *KVCalculator) pagingPrechargeRRU(req RequestInfo) float64 { + bytesForEst := estimatedReadBytes(req) + if bytesForEst == 0 { + return 0 + } + return float64(kc.ReadBytesCost) * float64(bytesForEst) +} + +func adjustPagingPrechargeRRU(calculators []ResourceCalculator, consumption *rmpb.Consumption, req RequestInfo, sign float64) { + if consumption == nil { + return + } + for _, calc := range calculators { + kvCalc, ok := calc.(*KVCalculator) + if !ok { + continue + } + consumption.RRU += sign * kvCalc.pagingPrechargeRRU(req) + } +} + +func reportedRequestConsumption(calculators []ResourceCalculator, req RequestInfo, tokenDelta *rmpb.Consumption) *rmpb.Consumption { + if estimatedReadBytes(req) == 0 { + return tokenDelta + } + reported := cloneConsumption(tokenDelta) + adjustPagingPrechargeRRU(calculators, reported, req, -1) + return reported +} + +func reportedResponseConsumption(calculators []ResourceCalculator, req RequestInfo, tokenDelta *rmpb.Consumption) *rmpb.Consumption { + if estimatedReadBytes(req) == 0 { + return tokenDelta + } + reported := cloneConsumption(tokenDelta) + adjustPagingPrechargeRRU(calculators, reported, req, 1) + 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 3e27534fd97..6f31ed5e64e 100644 --- a/client/resource_group/controller/model_test.go +++ b/client/resource_group/controller/model_test.go @@ -22,6 +22,42 @@ import ( rmpb "github.com/pingcap/kvproto/pkg/resource_manager" ) +type legacyRequestInfo struct { + isWrite bool +} + +func (req *legacyRequestInfo) IsWrite() bool { + return req.isWrite +} + +func (*legacyRequestInfo) WriteBytes() uint64 { + return 0 +} + +func (*legacyRequestInfo) ReplicaNumber() int64 { + return 0 +} + +func (*legacyRequestInfo) StoreID() uint64 { + return 0 +} + +func (*legacyRequestInfo) RequestSize() uint64 { + return 0 +} + +func (*legacyRequestInfo) AccessLocationType() AccessLocationType { + return AccessUnknown +} + +type copRequestInfoWithoutPrediction struct { + legacyRequestInfo +} + +func (*copRequestInfoWithoutPrediction) IsCop() bool { + return true +} + func TestGetRUValueFromConsumption(t *testing.T) { // Positive test case re := require.New(t) @@ -133,6 +169,163 @@ func TestUpdateDeltaConsumption(t *testing.T) { re.Equal(now.TiflashRUV2, last.TiflashRUV2) } +func TestUpdateDeltaConsumptionReportsEmptyConsumption(t *testing.T) { + re := require.New(t) + last := &rmpb.Consumption{RRU: 10, WRU: 20, ReadBytes: 30} + now := &rmpb.Consumption{RRU: 10, WRU: 20, ReadBytes: 30} + + delta := updateDeltaConsumption(last, now) + + re.Equal(&rmpb.Consumption{}, delta) + re.Equal(now, last) +} + +func TestUpdateDeltaConsumptionIgnoresDecreasedRequestUnits(t *testing.T) { + re := require.New(t) + last := &rmpb.Consumption{ + RRU: 10, + WRU: 20, + ReadBytes: 100, + WriteBytes: 200, + TotalCpuTimeMs: 300, + SqlLayerCpuTimeMs: 400, + KvReadRpcCount: 500, + KvWriteRpcCount: 600, + ReadCrossAzTrafficBytes: 700, + WriteCrossAzTrafficBytes: 800, + TikvRUV2: 900, + TidbRUV2: 1000, + TiflashRUV2: 1100, + } + now := &rmpb.Consumption{ + RRU: 7, + WRU: 15, + ReadBytes: 90, + WriteBytes: 190, + TotalCpuTimeMs: 290, + SqlLayerCpuTimeMs: 390, + KvReadRpcCount: 490, + KvWriteRpcCount: 590, + ReadCrossAzTrafficBytes: 690, + WriteCrossAzTrafficBytes: 790, + TikvRUV2: 890, + TidbRUV2: 990, + TiflashRUV2: 1090, + } + + delta := updateDeltaConsumption(last, now) + + re.Equal(&rmpb.Consumption{}, delta) + re.Equal(float64(10), last.RRU) + re.Equal(float64(20), last.WRU) + re.Equal(float64(100), last.ReadBytes) + re.Equal(float64(200), last.WriteBytes) + re.Equal(float64(300), last.TotalCpuTimeMs) + re.Equal(float64(400), last.SqlLayerCpuTimeMs) + re.Equal(float64(500), last.KvReadRpcCount) + re.Equal(float64(600), last.KvWriteRpcCount) + re.Equal(uint64(700), last.ReadCrossAzTrafficBytes) + re.Equal(uint64(800), last.WriteCrossAzTrafficBytes) + re.Equal(float64(900), last.TikvRUV2) + re.Equal(float64(1000), last.TidbRUV2) + re.Equal(float64(1100), last.TiflashRUV2) +} + +func TestRequestInfoPagingProvidersAreOptional(t *testing.T) { + re := require.New(t) + cfg := DefaultRUConfig() + kvCalc := newKVCalculator(cfg) + req := &legacyRequestInfo{} + + bytesForEst, ok := pagingReadEstimate(req) + re.False(ok) + re.Zero(bytesForEst) + + consumption := &rmpb.Consumption{} + kvCalc.BeforeKVRequest(consumption, req) + re.InDelta(float64(cfg.ReadBaseCost)+float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion, consumption.RRU, 1e-6) + + resp := &TestResponseInfo{readBytes: 1024, succeed: true} + settle := &rmpb.Consumption{} + kvCalc.AfterKVRequest(settle, req, resp) + re.InDelta(float64(cfg.ReadBytesCost)*1024, settle.RRU, 1e-6) +} + +func TestRequestInfoMissingPredictionProviderReturnsZeroHint(t *testing.T) { + re := require.New(t) + req := &copRequestInfoWithoutPrediction{} + + bytesForEst, ok := pagingReadEstimate(req) + + re.True(ok) + re.Zero(bytesForEst) +} + +func TestReportedConsumptionStripsPagingPrecharge(t *testing.T) { + re := require.New(t) + cfg := DefaultRUConfig() + kvCalc := newKVCalculator(cfg) + calculators := []ResourceCalculator{kvCalc} + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: 1024, + } + + tokenDelta := &rmpb.Consumption{ + RRU: float64(cfg.ReadBaseCost) + + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion + + float64(cfg.ReadBytesCost)*1024, + } + reported := reportedRequestConsumption(calculators, req, tokenDelta) + + re.InDelta(float64(cfg.ReadBaseCost)+ + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion, + reported.RRU, 1e-6) + re.InDelta(tokenDelta.RRU, reported.RRU+float64(cfg.ReadBytesCost)*1024, 1e-6) +} + +func TestReportedConsumptionRestoresPagingActualUsage(t *testing.T) { + re := require.New(t) + cfg := DefaultRUConfig() + kvCalc := newKVCalculator(cfg) + calculators := []ResourceCalculator{kvCalc} + req := &TestRequestInfo{ + isWrite: false, + isCop: true, + predictedReadBytes: 4096, + } + + actualReadBytes := uint64(1024) + tokenDelta := &rmpb.Consumption{ + RRU: float64(cfg.ReadBytesCost) * (float64(actualReadBytes) - float64(req.predictedReadBytes)), + ReadBytes: float64(actualReadBytes), + } + reported := reportedResponseConsumption(calculators, req, tokenDelta) + + re.InDelta(float64(cfg.ReadBytesCost)*float64(actualReadBytes), reported.RRU, 1e-6) + re.Equal(float64(actualReadBytes), reported.ReadBytes) + re.Negative(tokenDelta.RRU) + re.GreaterOrEqual(reported.RRU, 0.0) +} + +func TestReportedConsumptionReusesDeltaWithoutPagingPrecharge(t *testing.T) { + re := require.New(t) + calculators := []ResourceCalculator{newKVCalculator(DefaultRUConfig())} + req := &TestRequestInfo{ + isWrite: false, + isCop: false, + predictedReadBytes: 1024, + } + tokenDelta := &rmpb.Consumption{RRU: 1, ReadBytes: 64} + + reportedRequest := reportedRequestConsumption(calculators, req, tokenDelta) + reportedResponse := reportedResponseConsumption(calculators, req, tokenDelta) + + re.Same(tokenDelta, reportedRequest) + re.Same(tokenDelta, reportedResponse) +} + func TestEqualRU(t *testing.T) { re := require.New(t) diff --git a/client/resource_group/controller/testutil.go b/client/resource_group/controller/testutil.go index a055b32fa6f..2284273ec61 100644 --- a/client/resource_group/controller/testutil.go +++ b/client/resource_group/controller/testutil.go @@ -22,11 +22,13 @@ import "time" // TestRequestInfo is used to test the request info interface. type TestRequestInfo struct { - isWrite bool - writeBytes uint64 - numReplicas int64 - storeID uint64 - accessType AccessLocationType + isWrite bool + writeBytes uint64 + numReplicas int64 + storeID uint64 + accessType AccessLocationType + predictedReadBytes uint64 + isCop bool } // NewTestRequestInfo creates a new TestRequestInfo. @@ -70,6 +72,16 @@ func (tri *TestRequestInfo) AccessLocationType() AccessLocationType { return tri.accessType } +// PredictedReadBytes implements the optional predictedReadBytesProvider interface. +func (tri *TestRequestInfo) PredictedReadBytes() uint64 { + return tri.predictedReadBytes +} + +// IsCop implements the optional copRequestInfo interface. +func (tri *TestRequestInfo) IsCop() bool { + return tri.isCop +} + // TestResponseInfo is used to test the response info interface. type TestResponseInfo struct { readBytes uint64 diff --git a/pkg/mcs/resourcemanager/server/manager_test.go b/pkg/mcs/resourcemanager/server/manager_test.go index 72ca915105e..cba0aea08ef 100644 --- a/pkg/mcs/resourcemanager/server/manager_test.go +++ b/pkg/mcs/resourcemanager/server/manager_test.go @@ -431,6 +431,28 @@ func checkBackgroundMetricsFlush(ctx context.Context, re *require.Assertions, ma }) } +func TestDispatchConsumptionIncludesOnlyConsumption(t *testing.T) { + re := require.New(t) + m := newManagerBase(&ControllerConfig{}, ResourceGroupWriteRoleLegacyAll) + req := &rmpb.TokenBucketRequest{ + ResourceGroupName: "test_group", + ConsumptionSinceLastRequest: &rmpb.Consumption{ + RRU: 12, + WRU: 8, + WriteBytes: 1024, + }, + KeyspaceId: &rmpb.KeyspaceIDValue{Value: 42}, + } + + err := m.dispatchConsumption(req) + re.NoError(err) + + item := <-m.consumptionDispatcher + re.Equal(uint32(42), item.keyspaceID) + re.Equal(req.GetResourceGroupName(), item.resourceGroupName) + re.Equal(req.GetConsumptionSinceLastRequest(), item.Consumption) +} + // Put a keyspace meta into the storage. func prepareKeyspaceName(ctx context.Context, re *require.Assertions, manager *Manager, keyspaceIDValue *rmpb.KeyspaceIDValue, keyspaceName string) { keyspaceMeta := &keyspacepb.KeyspaceMeta{ diff --git a/pkg/mcs/resourcemanager/server/metering_test.go b/pkg/mcs/resourcemanager/server/metering_test.go index fdad1f0f7b3..81050b19e9b 100644 --- a/pkg/mcs/resourcemanager/server/metering_test.go +++ b/pkg/mcs/resourcemanager/server/metering_test.go @@ -81,6 +81,26 @@ func TestRUCollectorCollectSingleKeyspace(t *testing.T) { re.Equal(metering.NewRUValue(40.0), record[meteringDataTiFlashRUV2Field]) } +func TestRUCollectorUsesConsumption(t *testing.T) { + re := require.New(t) + collector := newRUCollector() + + collector.Collect(&consumptionItem{ + keyspaceName: testKeyspaceName, + Consumption: &rmpb.Consumption{ + RRU: 12, + WRU: 8, + WriteBytes: 1024, + }, + }) + + records := collector.Aggregate() + re.Len(records, 1) + record := records[0] + re.Equal(metering.NewRUValue(20.0), record[meteringDataOLTPRUField]) + re.Equal(metering.NewBytesValue(1024), record[meteringDataWriteBytesField]) +} + func TestRUCollectorCollectMultipleKeyspaces(t *testing.T) { re := require.New(t) collector := newRUCollector() diff --git a/pkg/mcs/resourcemanager/server/metrics.go b/pkg/mcs/resourcemanager/server/metrics.go index 2f6e28ee1c8..12ae61dafe1 100644 --- a/pkg/mcs/resourcemanager/server/metrics.go +++ b/pkg/mcs/resourcemanager/server/metrics.go @@ -549,6 +549,9 @@ func (m *metrics) recordConsumption( ruLabelType = tiflashTypeLabel } consumption := consumptionInfo.Consumption + if consumption == nil { + return + } m.getMaxPerSecTracker(keyspaceID, keyspaceName, groupName).collect(consumption) m.getCounterMetrics(keyspaceID, keyspaceName, groupName, ruLabelType).add(consumption, controllerConfig, keyspaceID) m.insertConsumptionRecord(keyspaceID, groupName, ruLabelType, now) @@ -692,6 +695,9 @@ func calculateActiveRU(consumption *rmpb.Consumption, controllerConfig *Controll } func (m *counterMetrics) add(consumption *rmpb.Consumption, controllerConfig *ControllerConfig, keyspaceID uint32) { + if consumption == nil { + return + } // RU info. if consumption.RRU > 0 { m.RRUMetrics.Add(consumption.RRU) diff --git a/pkg/mcs/resourcemanager/server/metrics_test.go b/pkg/mcs/resourcemanager/server/metrics_test.go index 08c07b52e21..ffef1f31ab7 100644 --- a/pkg/mcs/resourcemanager/server/metrics_test.go +++ b/pkg/mcs/resourcemanager/server/metrics_test.go @@ -91,6 +91,222 @@ func TestActiveRequestUnitMetricUsesV2ForOverriddenKeyspace(t *testing.T) { re.Equal(float64(31), testutil.ToFloat64(metrics.ActiveRUMetrics)) } +func TestCounterMetricsIgnoresNegativeRequestUnitConsumption(t *testing.T) { + re := require.New(t) + + const ( + keyspaceName = "negative-ru-keyspace" + groupName = "negative-ru-group" + keyspaceID = uint32(303) + ) + + t.Cleanup(func() { + deleteDefaultLabelValuesForTest(keyspaceName, groupName) + }) + + metrics := newCounterMetrics(keyspaceName, groupName, defaultTypeLabel) + metrics.add(&rmpb.Consumption{ + RRU: 10, + WRU: 3, + }, &ControllerConfig{}, keyspaceID) + metrics.add(&rmpb.Consumption{ + RRU: -4, + WRU: -1, + }, &ControllerConfig{}, keyspaceID) + + re.Equal(float64(10), testutil.ToFloat64(metrics.RRUMetrics)) + re.Equal(float64(3), testutil.ToFloat64(metrics.WRUMetrics)) + re.Equal(float64(13), testutil.ToFloat64(metrics.ActiveRUMetrics)) +} + +func TestRecordConsumptionUsesActualRequestUnitCounters(t *testing.T) { + re := require.New(t) + + const ( + keyspaceName = "actual-counter-keyspace" + groupName = "actual-counter-group" + keyspaceID = uint32(404) + ) + + t.Cleanup(func() { + deleteDefaultLabelValuesForTest(keyspaceName, groupName) + }) + + m := newMetrics() + m.recordConsumption(&consumptionItem{ + keyspaceID: keyspaceID, + keyspaceName: keyspaceName, + resourceGroupName: groupName, + Consumption: &rmpb.Consumption{ + RRU: 12, + WRU: 8, + WriteBytes: 1024, + }, + }, &ControllerConfig{}, time.Now()) + + counter := m.getCounterMetrics(keyspaceID, keyspaceName, groupName, defaultTypeLabel) + re.Equal(float64(12), testutil.ToFloat64(counter.RRUMetrics)) + re.Equal(float64(8), testutil.ToFloat64(counter.WRUMetrics)) + re.Equal(float64(1024), testutil.ToFloat64(counter.WriteByteMetrics)) + re.Contains(m.consumptionRecordMap, consumptionRecordKey{ + keyspaceID: keyspaceID, + groupName: groupName, + ruType: defaultTypeLabel, + }) +} + +func TestRecordConsumptionKeepsRecordForEmptyConsumption(t *testing.T) { + re := require.New(t) + + const ( + keyspaceName = "empty-consumption-keyspace" + groupName = "empty-consumption-group" + keyspaceID = uint32(407) + ) + + t.Cleanup(func() { + deleteDefaultLabelValuesForTest(keyspaceName, groupName) + }) + + m := newMetrics() + now := time.Now() + m.recordConsumption(&consumptionItem{ + keyspaceID: keyspaceID, + keyspaceName: keyspaceName, + resourceGroupName: groupName, + Consumption: &rmpb.Consumption{}, + }, &ControllerConfig{}, now) + + key := consumptionRecordKey{ + keyspaceID: keyspaceID, + groupName: groupName, + ruType: defaultTypeLabel, + } + re.Equal(now, m.consumptionRecordMap[key]) + + counter := m.counterMetricsMap[metricsKey{keyspaceID: keyspaceID, groupName: groupName, ruType: defaultTypeLabel}] + re.NotNil(counter) + re.Zero(testutil.ToFloat64(counter.RRUMetrics)) + re.Zero(testutil.ToFloat64(counter.WRUMetrics)) + re.Zero(testutil.ToFloat64(counter.ActiveRUMetrics)) +} + +func TestRecordConsumptionIgnoresNilConsumption(t *testing.T) { + re := require.New(t) + + m := newMetrics() + re.NotPanics(func() { + m.recordConsumption(&consumptionItem{ + keyspaceID: 408, + keyspaceName: "nil-consumption-keyspace", + resourceGroupName: "nil-consumption-group", + }, &ControllerConfig{}, time.Now()) + }) + re.Empty(m.counterMetricsMap) + re.Empty(m.consumptionRecordMap) +} + +func TestRecordConsumptionDoesNotRefundNegativeConsumption(t *testing.T) { + re := require.New(t) + + const ( + keyspaceName = "negative-consumption-keyspace" + groupName = "negative-consumption-group" + keyspaceID = uint32(405) + ) + + t.Cleanup(func() { + deleteDefaultLabelValuesForTest(keyspaceName, groupName) + }) + + m := newMetrics() + m.recordConsumption(&consumptionItem{ + keyspaceID: keyspaceID, + keyspaceName: keyspaceName, + resourceGroupName: groupName, + Consumption: &rmpb.Consumption{ + RRU: -12, + WRU: -8, + }, + }, &ControllerConfig{}, time.Now()) + + counter := m.getCounterMetrics(keyspaceID, keyspaceName, groupName, defaultTypeLabel) + re.Zero(testutil.ToFloat64(counter.RRUMetrics)) + re.Zero(testutil.ToFloat64(counter.WRUMetrics)) +} + +func TestRecordConsumptionUsesActualForActiveRequestUnit(t *testing.T) { + re := require.New(t) + + const ( + keyspaceName = "actual-active-ru-keyspace" + groupName = "actual-active-ru-group" + keyspaceID = uint32(505) + ) + + t.Cleanup(func() { + deleteDefaultLabelValuesForTest(keyspaceName, groupName) + }) + + m := newMetrics() + m.recordConsumption(&consumptionItem{ + keyspaceID: keyspaceID, + keyspaceName: keyspaceName, + resourceGroupName: groupName, + Consumption: &rmpb.Consumption{ + RRU: 12, + WRU: 8, + TotalCpuTimeMs: 4, + SqlLayerCpuTimeMs: 3, + }, + }, &ControllerConfig{ + RequestUnit: RequestUnitConfig{CPUMsCost: 2}, + }, time.Now()) + + counter := m.getCounterMetrics(keyspaceID, keyspaceName, groupName, defaultTypeLabel) + re.Equal(float64(12), testutil.ToFloat64(counter.RRUMetrics)) + re.Equal(float64(8), testutil.ToFloat64(counter.WRUMetrics)) + re.Equal(float64(26), testutil.ToFloat64(counter.ActiveRUMetrics)) +} + +func TestRecordConsumptionUsesActualRUV2ForActiveRequestUnit(t *testing.T) { + re := require.New(t) + + const ( + keyspaceName = "actual-active-ruv2-keyspace" + groupName = "actual-active-ruv2-group" + keyspaceID = uint32(606) + ) + + t.Cleanup(func() { + deleteDefaultLabelValuesForTest(keyspaceName, groupName) + }) + + m := newMetrics() + m.recordConsumption(&consumptionItem{ + keyspaceID: keyspaceID, + keyspaceName: keyspaceName, + resourceGroupName: groupName, + Consumption: &rmpb.Consumption{ + TikvRUV2: 7, + TidbRUV2: 11, + TiflashRUV2: 13, + }, + }, &ControllerConfig{ + RUVersionPolicy: &RUVersionPolicy{ + Default: RUVersionV1, + Overrides: map[uint32]RUVersion{ + keyspaceID: RUVersionV2, + }, + }, + }, time.Now()) + + counter := m.getCounterMetrics(keyspaceID, keyspaceName, groupName, defaultTypeLabel) + re.Zero(testutil.ToFloat64(counter.RRUMetrics)) + re.Zero(testutil.ToFloat64(counter.WRUMetrics)) + re.Equal(float64(31), testutil.ToFloat64(counter.ActiveRUMetrics)) +} + func TestMaxPerSecCostTracker(t *testing.T) { re := require.New(t) tracker := newMaxPerSecCostTracker("test", "test", defaultCollectIntervalSec) diff --git a/pkg/metering/utils.go b/pkg/metering/utils.go index 6147767e367..2cb66b38f96 100644 --- a/pkg/metering/utils.go +++ b/pkg/metering/utils.go @@ -37,6 +37,9 @@ const ( // NewRUValue creates a new metering RU value. func NewRUValue(value float64) common.MeteringValue { + if value < 0 { + value = 0 + } return common.MeteringValue{Value: uint64(value), Unit: UnitRU} } diff --git a/pkg/metering/utils_test.go b/pkg/metering/utils_test.go new file mode 100644 index 00000000000..f11fdcd74f8 --- /dev/null +++ b/pkg/metering/utils_test.go @@ -0,0 +1,30 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package metering + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/pingcap/metering_sdk/common" +) + +func TestNewRUValueClampsNegativeValues(t *testing.T) { + re := require.New(t) + + re.Equal(common.MeteringValue{Value: 0, Unit: UnitRU}, NewRUValue(-1)) + re.Equal(common.MeteringValue{Value: 1, Unit: UnitRU}, NewRUValue(1.5)) +} From e561e26ac7a474aa2c35302ed0c2be1313632398 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Mon, 6 Jul 2026 14:45:07 +0800 Subject: [PATCH 2/2] client/resource_group: expose runtime state (#10975) ref tikv/pd#10612\n\nExpose resource group runtime state from the PD client. The runtime state currently reports whether the local effective request-unit token bucket has a finite burst limit (`HasLimitedBurst`), based on token bucket responses. This lets callers observe Resource Control runtime signals without changing resource group metadata semantics.\n\nSigned-off-by: JmPotato (cherry picked from commit 14ffe2e809ee3405cbf379e1c68a4c3ddf7a9f74) Signed-off-by: JmPotato --- .../controller/global_controller.go | 21 +++++ .../controller/global_controller_test.go | 90 +++++++++++++++++++ .../controller/group_controller.go | 11 ++- .../controller/group_controller_test.go | 2 +- 4 files changed, 117 insertions(+), 7 deletions(-) diff --git a/client/resource_group/controller/global_controller.go b/client/resource_group/controller/global_controller.go index 70450e1dbbd..66675a2dfdb 100644 --- a/client/resource_group/controller/global_controller.go +++ b/client/resource_group/controller/global_controller.go @@ -824,6 +824,27 @@ func (c *ResourceGroupsController) GetResourceGroup(resourceGroupName string) (* return gc.getMeta(), nil } +// ResourceGroupRuntimeState describes generic local runtime signals derived +// from token bucket responses. +type ResourceGroupRuntimeState struct { + // HasLimitedBurst is true when the request-unit token bucket has a finite + // burst limit instead of unlimited burst. + HasLimitedBurst bool +} + +// GetResourceGroupRuntimeState returns the resource group's local runtime +// state. The second return value is false when the controller has no usable +// local state for the group. +func (c *ResourceGroupsController) GetResourceGroupRuntimeState(resourceGroupName string) (ResourceGroupRuntimeState, bool) { + gc, ok := c.loadGroupController(resourceGroupName) + if !ok || gc.tombstone.Load() || !gc.initialRequestCompleted.Load() { + return ResourceGroupRuntimeState{}, false + } + return ResourceGroupRuntimeState{ + HasLimitedBurst: !gc.burstable.Load(), + }, true +} + // ReportConsumption is used to report ru consumption directly. // // Currently, this interface is used to report the consumption for TiFlash MPP cost diff --git a/client/resource_group/controller/global_controller_test.go b/client/resource_group/controller/global_controller_test.go index b7018fc0953..359de785673 100644 --- a/client/resource_group/controller/global_controller_test.go +++ b/client/resource_group/controller/global_controller_test.go @@ -448,6 +448,96 @@ func TestGetResourceGroup(t *testing.T) { re.Nil(gc02) } +func TestGetResourceGroupRuntimeState(t *testing.T) { + testCases := []struct { + name string + responseBurstLimit int64 + sendResponse bool + resourceGroupName string + expectOK bool + expectResourceGroupRuntimeState ResourceGroupRuntimeState + }{ + { + name: "unknown before first token response", + resourceGroupName: "test-group", + }, + { + name: "unlimited response remains burstable", + responseBurstLimit: -1, + sendResponse: true, + resourceGroupName: "test-group", + expectOK: true, + }, + { + name: "override limited burst from token response", + responseBurstLimit: 100, + sendResponse: true, + resourceGroupName: "test-group", + expectOK: true, + expectResourceGroupRuntimeState: ResourceGroupRuntimeState{ + HasLimitedBurst: true, + }, + }, + { + name: "unknown resource group", + responseBurstLimit: 100, + sendResponse: true, + resourceGroupName: "unknown", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + re := require.New(t) + + group := &rmpb.ResourceGroup{ + Name: "test-group", + Mode: rmpb.GroupMode_RUMode, + RUSettings: &rmpb.GroupRequestUnitSettings{ + RU: &rmpb.TokenBucket{ + Settings: &rmpb.TokenLimitSettings{ + FillRate: 1000, + BurstLimit: -1, + }, + }, + }, + } + gc, err := newGroupCostController(group, 1001, DefaultRUConfig(), make(chan notifyMsg), make(chan *groupCostController)) + re.NoError(err) + + controller := &ResourceGroupsController{} + controller.groupsController.Store(group.Name, gc) + + if tc.sendResponse { + controller.handleTokenBucketResponse([]*rmpb.TokenBucketResponse{ + { + ResourceGroupName: group.Name, + GrantedRUTokens: []*rmpb.GrantedRUTokenBucket{ + { + GrantedTokens: &rmpb.TokenBucket{ + Settings: &rmpb.TokenLimitSettings{ + FillRate: 1000, + BurstLimit: tc.responseBurstLimit, + }, + Tokens: 1000, + }, + }, + }, + }, + }) + // Mirror the main loop, which refreshes the derived state + // (e.g. burstable) after handling token bucket responses. + gc.updateRunState() + gc.updateAvgRequestResourcePerSec() + } + + state, ok := controller.GetResourceGroupRuntimeState(tc.resourceGroupName) + re.Equal(tc.expectOK, ok) + re.Equal(tc.expectResourceGroupRuntimeState, state) + }) + } +} + func TestTokenBucketsRequestWithKeyspaceID(t *testing.T) { re := require.New(t) ctx, cancel := context.WithCancel(context.Background()) diff --git a/client/resource_group/controller/group_controller.go b/client/resource_group/controller/group_controller.go index 2eba9f7b777..626fadf2802 100644 --- a/client/resource_group/controller/group_controller.go +++ b/client/resource_group/controller/group_controller.go @@ -85,13 +85,12 @@ type groupCostController struct { // lastResourceConsumptions []*rmpb.ResourceItem lastRequestConsumption *rmpb.Consumption - // initialRequestCompleted is set to true when the first token bucket - // request completes successfully. - initialRequestCompleted bool - requestUnitTokens *tokenCounter } + // initialRequestCompleted is set to true when the first token bucket + // request completes successfully. + initialRequestCompleted atomic.Bool // tombstone is set to true when the resource group is deleted. tombstone atomic.Bool // inactive is set to true when the resource group has not been updated for a long time. @@ -486,7 +485,7 @@ func calcMovingAvgAt(now time.Time, avg *float64, lastRU *float64, lastTime *tim } func (gc *groupCostController) shouldReportConsumption() bool { - if !gc.run.initialRequestCompleted { + if !gc.initialRequestCompleted.Load() { return true } timeSinceLastRequest := gc.run.now.Sub(gc.run.lastRequestTime) @@ -510,7 +509,7 @@ func (gc *groupCostController) shouldReportConsumption() bool { func (gc *groupCostController) handleTokenBucketResponse(resp *rmpb.TokenBucketResponse) { gc.run.requestInProgress = false gc.handleRUTokenResponse(resp) - gc.run.initialRequestCompleted = true + gc.initialRequestCompleted.Store(true) } func (gc *groupCostController) handleRUTokenResponse(resp *rmpb.TokenBucketResponse) { diff --git a/client/resource_group/controller/group_controller_test.go b/client/resource_group/controller/group_controller_test.go index 9c278cdfbef..aa71acffa52 100644 --- a/client/resource_group/controller/group_controller_test.go +++ b/client/resource_group/controller/group_controller_test.go @@ -120,7 +120,7 @@ func TestSmallPositiveConsumptionDoesNotBypassThreshold(t *testing.T) { now := time.Now() gc.run.consumption.RRU = 1 - gc.run.initialRequestCompleted = true + gc.initialRequestCompleted.Store(true) gc.run.lastRequestTime = now.Add(-defaultTargetPeriod) gc.run.now = now