resource_control: add PredictedReadBytes hint for RC paging pre-charge

Introduce an optional predictedReadBytesProvider interface on RequestInfo.
When a caller (e.g. TiDB maintaining a per-logical-scan EMA across paging
RPCs) supplies a non-zero PredictedReadBytes, BeforeKVRequest/AfterKVRequest
use that value as the byte basis for the paging pre-charge instead of
PagingSizeBytes. PagingSizeBytes remains the fallback and worst-case cap.

This lets TiDB replace the current fixed 4 MiB pre-charge (which matches
the paging byte budget but typically overshoots actual scanned bytes and
stalls concurrent workers at Phase 1) with a learned estimate, without
changing kvproto or TiKV behavior.

The hint is added as an optional interface (not a method on RequestInfo)
so existing RequestInfo implementations compile unchanged; they continue
to fall back to PagingSizeBytes.

Ref: per-logical-scan EMA pre-deduction design
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>

Author: YuhaoZhang00

client/resource_group/controller/group_controller_test.go

@@ -257,6 +257,70 @@ func TestPagingSizeBytesPreCharge(t *testing.T) {
 		"Without paging, Phase 1 should only charge baseCost")
 }
 
+func TestPredictedReadBytesOverridesPagingSizeBytes(t *testing.T) {
+	re := require.New(t)
+	cfg := DefaultRUConfig()
+	kvCalc := newKVCalculator(cfg)
+
+	// When PredictedReadBytes (the EMA hint) is > 0, it should override
+	// PagingSizeBytes as the pre-charge basis. PagingSizeBytes is kept as a
+	// safety cap and worst-case fallback; the hint gives a tighter estimate.
+	pagingSizeBytes := uint64(4 * 1024 * 1024) // 4 MB worst-case cap
+	predictedReadBytes := uint64(256 * 1024)   // 256 KB learned estimate
+	req := &TestRequestInfo{
+		isWrite:            false,
+		pagingSizeBytes:    pagingSizeBytes,
+		predictedReadBytes: predictedReadBytes,
+	}
+
+	phase1 := &rmpb.Consumption{}
+	kvCalc.BeforeKVRequest(phase1, req)
+
+	baseCost := float64(cfg.ReadBaseCost) + float64(cfg.ReadPerBatchBaseCost)*defaultAvgBatchProportion
+	hintCost := float64(cfg.ReadBytesCost) * float64(predictedReadBytes)
+	re.InDelta(baseCost+hintCost, phase1.RRU, 1e-6,
+		"Phase 1 should pre-charge based on PredictedReadBytes, not PagingSizeBytes")
+
+	// Phase 2 should subtract the same hint-based basis, preserving the
+	// invariant that Phase 1 + Phase 2 == baseCost + actualCost.
+	actualReadBytes := uint64(300 * 1024) // close to the prediction
+	resp := &TestResponseInfo{
+		readBytes: actualReadBytes,
+		kvCPU:     0,
+		succeed:   true,
+	}
+	phase2 := &rmpb.Consumption{}
+	kvCalc.AfterKVRequest(phase2, req, resp)
+
+	actualReadCost := float64(cfg.ReadBytesCost) * float64(actualReadBytes)
+	expectedPhase2RRU := actualReadCost - hintCost
+	re.InDelta(expectedPhase2RRU, phase2.RRU, 1e-6,
+		"Phase 2 should settle using the same hint basis as Phase 1")
+
+	totalRRU := phase1.RRU + phase2.RRU
+	re.InDelta(baseCost+actualReadCost, totalRRU, 1e-6,
+		"Total RRU across Phase 1+2 should still equal baseCost + actualCost")
+
+	// When the hint is zero, fall back to PagingSizeBytes (old behavior).
+	reqFallback := &TestRequestInfo{
+		isWrite:         false,
+		pagingSizeBytes: pagingSizeBytes,
+		// predictedReadBytes left at 0
+	}
+	phase1Fallback := &rmpb.Consumption{}
+	kvCalc.BeforeKVRequest(phase1Fallback, reqFallback)
+	fallbackBytesCost := float64(cfg.ReadBytesCost) * float64(pagingSizeBytes)
+	re.InDelta(baseCost+fallbackBytesCost, phase1Fallback.RRU, 1e-6,
+		"With no hint, pre-charge should fall back to PagingSizeBytes")
+
+	// When both are zero, no byte-based pre-charge is applied.
+	reqNone := &TestRequestInfo{isWrite: false}
+	phase1None := &rmpb.Consumption{}
+	kvCalc.BeforeKVRequest(phase1None, reqNone)
+	re.InDelta(baseCost, phase1None.RRU, 1e-6,
+		"Without hint or paging, Phase 1 should only charge baseCost")
+}
+
 func TestPagingPreChargeTokenRefund(t *testing.T) {
 	re := require.New(t)
 	gc := createTestGroupCostController(re)

client/resource_group/controller/model.go

@@ -57,6 +57,32 @@ type RequestInfo interface {
 	PagingSizeBytes() uint64
 }
 
+// predictedReadBytesProvider is an optional interface that a RequestInfo
+// implementation may satisfy to supply a learned estimate (e.g., from a
+// per-logical-scan EMA maintained in TiDB) of how many bytes the request
+// will read. When present and > 0, it overrides PagingSizeBytes as the
+// byte basis for RC paging pre-charge.
+//
+// Defined as an optional interface (not a method on RequestInfo) so older
+// RequestInfo implementations that have not been updated still compile and
+// behave as before (falling back to PagingSizeBytes).
+type predictedReadBytesProvider interface {
+	PredictedReadBytes() uint64
+}
+
+// estimatedReadBytes returns the byte basis used for RC paging pre-charge.
+// It prefers a learned PredictedReadBytes hint when the RequestInfo
+// implements the optional provider and returns a non-zero value; otherwise
+// it falls back to PagingSizeBytes (the paging byte budget / worst-case cap).
+func estimatedReadBytes(req RequestInfo) uint64 {
+	if p, ok := req.(predictedReadBytesProvider); ok {
+		if hint := p.PredictedReadBytes(); hint > 0 {
+			return hint
+		}
+	}
+	return req.PagingSizeBytes()
+}
+
 // 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 {
@@ -107,9 +133,11 @@ func (kc *KVCalculator) BeforeKVRequest(consumption *rmpb.Consumption, req Reque
 		consumption.RRU += float64(kc.ReadBaseCost) + float64(kc.ReadPerBatchBaseCost)*defaultAvgBatchProportion
 		// RC Paging pre-charge: if the request has a byte budget, pre-charge
 		// the estimated read bytes RU so that concurrent workers are throttled
-		// at Phase 1 instead of all hitting Phase 2 at once.
-		if pagingSizeBytes := req.PagingSizeBytes(); pagingSizeBytes > 0 {
-			consumption.RRU += float64(kc.ReadBytesCost) * float64(pagingSizeBytes)
+		// at Phase 1 instead of all hitting Phase 2 at once. Prefer a learned
+		// PredictedReadBytes hint when the caller supplies one; otherwise fall
+		// back to PagingSizeBytes (the paging byte budget / worst-case cap).
+		if bytesForEst := estimatedReadBytes(req); bytesForEst > 0 {
+			consumption.RRU += float64(kc.ReadBytesCost) * float64(bytesForEst)
 		}
 	}
 	if req.AccessLocationType() == AccessCrossZone {
@@ -147,9 +175,10 @@ func (kc *KVCalculator) AfterKVRequest(consumption *rmpb.Consumption, req Reques
 		kc.calculateCPUCost(consumption, res)
 		// RC Paging settlement: subtract the pre-charged bytes RU added in
 		// BeforeKVRequest so the net total (Phase 1 + Phase 2) equals
-		// baseCost + actualCost.
-		if pagingSizeBytes := req.PagingSizeBytes(); pagingSizeBytes > 0 {
-			consumption.RRU -= float64(kc.ReadBytesCost) * float64(pagingSizeBytes)
+		// baseCost + actualCost. Use the same basis (hint or PagingSizeBytes)
+		// that BeforeKVRequest used.
+		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.

client/resource_group/controller/testutil.go

@@ -22,12 +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
-	pagingSizeBytes uint64
+	isWrite            bool
+	writeBytes         uint64
+	numReplicas        int64
+	storeID            uint64
+	accessType         AccessLocationType
+	pagingSizeBytes    uint64
+	predictedReadBytes uint64
 }
 
 // NewTestRequestInfo creates a new TestRequestInfo.
@@ -76,6 +77,12 @@ func (tri *TestRequestInfo) PagingSizeBytes() uint64 {
 	return tri.pagingSizeBytes
 }
 
+// PredictedReadBytes implements the optional predictedReadBytesProvider
+// interface so tests can exercise the EMA-based hint path.
+func (tri *TestRequestInfo) PredictedReadBytes() uint64 {
+	return tri.predictedReadBytes
+}
+
 // TestResponseInfo is used to test the response info interface.
 type TestResponseInfo struct {
 	readBytes uint64