resourcecontrol, tikvrpc: add PredictedReadBytes hint for RC paging pre-charge#1947
Merged
ti-chi-bot[bot] merged 15 commits intoJul 8, 2026
Merged
Conversation
Add a client-go-internal PredictedReadBytes field on tikvrpc.Request so the caller (e.g. TiDB, maintaining a per-logical-scan EMA across paging cop RPCs) can supply a learned estimate of how many bytes the request will read. MakeRequestInfo propagates this into RequestInfo, and the new RequestInfo.PredictedReadBytes() getter satisfies the optional predictedReadBytesProvider interface that PD's resource_group/controller checks via type assertion. When the hint is > 0, PD uses it as the byte basis for RC paging pre-charge. Zero means the caller has no prediction (e.g. cold start); the request is not pre-charged and is billed at settlement time by actual read bytes only. The field is kept out of the proto because it is purely a client-side estimate consumed before the RPC is sent - TiKV neither needs nor reads it. Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
Drop the verbose restatement of how the hint is produced and consumed; keep only what a reader of this struct needs to know to set the field. Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
ti-chi-bot
Bot
added
dco-signoff: yes
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds a caller-supplied ChangesRequest Type Extension
Resource Control
Client Interceptor
Module pin
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
ti-chi-bot
Bot
added
the
size/M
This was referenced Apr 21, 2026
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
internal/resourcecontrol/resource_control_test.go (1)
52-74: Test coverage is adequate for the propagation contract.Both the hint-present and hint-absent (zero-value default) paths are asserted, and
IsWrite() == falseis verified so the read-path branch inMakeRequestInfois the one under test.One optional addition worth considering: a case asserting that for a write request (e.g.,
CmdPrewrite) the hint is not propagated (i.e.,PredictedReadBytes()returns0even when set on thetikvrpc.Request). This pins down the current write-path behavior so any future accidental propagation would be caught.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/resourcecontrol/resource_control_test.go` around lines 52 - 74, Add a test that verifies PredictedReadBytes is not propagated for write requests: create a tikvrpc.Request with Req set to a write command (e.g., &kvrpcpb.CmdPrewrite{}), Context with a Peer, and PredictedReadBytes non-zero; call MakeRequestInfo(req) and assert that the returned RequestInfo.PredictedReadBytes() == 0 and IsWrite() == true to ensure write-path does not carry the read hint. Ensure the new case mirrors the existing read-case structure so it fails if MakeRequestInfo incorrectly propagates the hint for writes.internal/resourcecontrol/resource_control.go (1)
90-131: Hint is silently dropped on write path — confirm this is intentional.
predictedReadBytesis only assigned in the read branch (Line 104). In the write branch (Lines 123–130) the field is left zero even if the caller setreq.PredictedReadBytes. Given the field is documented as a read-bytes estimate, this is almost certainly intentional — but since there's no guard at the call site, a caller mistakenly setting it on a write request gets silent data loss rather than a signal.Two low-cost options if you want to make the contract more explicit:
- Add a short comment in the write branch noting the hint is intentionally ignored for writes.
- Or drop the field at the
Requeststruct level into a helper that is only read on read paths (current design already effectively does this via the getter being consulted only by PD's read-path pre-charge — so a comment is probably enough).No change strictly required; flagging for visibility.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/resourcecontrol/resource_control.go` around lines 90 - 131, MakeRequestInfo drops req.PredictedReadBytes on the write path (in the branch building RequestInfo for write requests), causing the hint to be silently lost; either preserve it by setting predictedReadBytes: req.PredictedReadBytes in the RequestInfo returned by the write-path branch (the struct built at the end of MakeRequestInfo) or, if the drop is intentional, add an explicit comment in MakeRequestInfo (near the write-path return) stating that predictedReadBytes is ignored for write requests to avoid silent surprises; reference symbols: MakeRequestInfo, RequestInfo, predictedReadBytes, req.PredictedReadBytes.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@internal/resourcecontrol/resource_control_test.go`:
- Around line 52-74: Add a test that verifies PredictedReadBytes is not
propagated for write requests: create a tikvrpc.Request with Req set to a write
command (e.g., &kvrpcpb.CmdPrewrite{}), Context with a Peer, and
PredictedReadBytes non-zero; call MakeRequestInfo(req) and assert that the
returned RequestInfo.PredictedReadBytes() == 0 and IsWrite() == true to ensure
write-path does not carry the read hint. Ensure the new case mirrors the
existing read-case structure so it fails if MakeRequestInfo incorrectly
propagates the hint for writes.
In `@internal/resourcecontrol/resource_control.go`:
- Around line 90-131: MakeRequestInfo drops req.PredictedReadBytes on the write
path (in the branch building RequestInfo for write requests), causing the hint
to be silently lost; either preserve it by setting predictedReadBytes:
req.PredictedReadBytes in the RequestInfo returned by the write-path branch (the
struct built at the end of MakeRequestInfo) or, if the drop is intentional, add
an explicit comment in MakeRequestInfo (near the write-path return) stating that
predictedReadBytes is ignored for write requests to avoid silent surprises;
reference symbols: MakeRequestInfo, RequestInfo, predictedReadBytes,
req.PredictedReadBytes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: de10caf0-abb3-4c5e-9151-3f6c02498470
📒 Files selected for processing (3)
internal/resourcecontrol/resource_control.gointernal/resourcecontrol/resource_control_test.gotikvrpc/tikvrpc.go
YuhaoZhang00
added a commit
to YuhaoZhang00/tidb
that referenced
this pull request
Apr 22, 2026
Temporarily replace github.com/tikv/client-go/v2 and github.com/tikv/pd/client with the corresponding forks carrying the PredictedReadBytes hint and controller-side pre-charge logic so CI can build this PR before the two upstream PRs land: - github.com/YuhaoZhang00/client-go/v2 (tikv/client-go#1947) - github.com/YuhaoZhang00/pd/client (tikv/pd#10611) Will be reverted and replaced with a normal require bump once both PRs merge and tag new releases. Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
YuhaoZhang00
added a commit
to YuhaoZhang00/tidb
that referenced
this pull request
Apr 22, 2026
Temporarily replace github.com/tikv/client-go/v2 and github.com/tikv/pd/client with the corresponding forks carrying the PredictedReadBytes hint and controller-side pre-charge logic so CI can build this PR before the two upstream PRs land: - github.com/YuhaoZhang00/client-go/v2 (tikv/client-go#1947) - github.com/YuhaoZhang00/pd/client (tikv/pd#10611) Will be reverted and replaced with a normal require bump once both PRs merge and tag new releases. Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
PD promoted PredictedReadBytes to a required method on RequestInfo, so this is no longer an optional duck-typed satisfaction. Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
rleungx
reviewed
May 18, 2026
PD's resource-control controller needs to gate paging_* accounting on coprocessor requests only; non-cop reads (CmdGet, CmdBatchGet, CmdScan, internal lookups) reach the same interceptor but never participate in the EMA pre-charge path and must not appear in paging_nonprecharge_*. Add an isCop field on RequestInfo derived from tikvrpc.Request.Type at construction time, and expose it via IsCop() to satisfy PD's new RequestInfo interface method. Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
ti-chi-bot
Bot
added
size/L
interceptedClient.SendRequest and SendRequestAsync today call OnResponseWait only when the underlying RPC returns a non-nil response. That made sense in the legacy billing model (only base RRU / WRU at risk on transport failure) but became a real RU leak after the PD-side paging pre-charge: the predicted RU debited in OnRequestWait stayed permanently charged whenever the RPC failed before producing a response (connection drop, timeout, context cancellation). When resp == nil, invoke the new ResourceGroupKVInterceptor.OnRequestCancel so PD refunds the predicted RU to the limiter and rolls back the consumption row it added in OnRequestWait. Cancel and settlement are mutually exclusive on a given request — the if/else structure makes that invariant explicit at every call site. Pin the PD client replace to the matching tikv/pd#10611 commit; the replace will be removed once that PR is merged and tagged. Adds TestSendRequest{Cancels,DoesNotCancel}{,Async}* covering the three combinations (sync failure / sync success / async failure) through a recordingInterceptor mock so the wiring is verified end-to-end without touching a real PD. Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
ti-chi-bot
Bot
added
size/XL
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/client/client_interceptor_test.go (1)
216-237: ⚡ Quick winAdd async success-path assertion for settle-vs-cancel behavior.
Line 216 currently covers async transport failure only. Add a companion async success test to assert
OnResponseWaitruns andOnRequestCancelstays zero.Proposed test
+func TestSendRequestAsyncDoesNotCancelOnSuccess(t *testing.T) { + rec := withRecordingInterceptor(t) + client := NewInterceptedClient(respondingClient{}) + + done := make(chan struct{}) + var gotResp *tikvrpc.Response + var gotErr error + cb := async.NewCallback(nil, func(resp *tikvrpc.Response, err error) { + gotResp = resp + gotErr = err + close(done) + }) + + client.SendRequestAsync(context.Background(), "", newRGRequest(), cb) + <-done + + assert.NotNil(t, gotResp) + assert.NoError(t, gotErr) + assert.Equal(t, 1, rec.waitCalls) + assert.Equal(t, 1, rec.respCalls) + assert.Equal(t, 0, rec.cancelCalls) +}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/client/client_interceptor_test.go` around lines 216 - 237, Add a new async-success test companion to TestSendRequestAsyncCancelsPreChargeOnTransportFailure that verifies the settle vs cancel behavior on a successful transport: create a recording interceptor with withRecordingInterceptor(t), construct an intercepted client using a passing client (e.g., successfulClient or a client that returns a valid tikvrpc.Response), invoke client.SendRequestAsync(context.Background(), "", newRGRequest(), cb) with an async callback that captures resp and err and closes a done channel, wait on done, then assert resp is non-nil and err is nil and assert rec.waitCalls == 1, rec.respCalls == 1 and rec.cancelCalls == 0 (with a message like "OnRequestCancel must not run on the async path when resp is non-nil"); place the new test alongside TestSendRequestAsyncCancelsPreChargeOnTransportFailure and name it e.g. TestSendRequestAsyncSettlesPreChargeOnTransportSuccess.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@internal/client/client_interceptor_test.go`:
- Around line 216-237: Add a new async-success test companion to
TestSendRequestAsyncCancelsPreChargeOnTransportFailure that verifies the settle
vs cancel behavior on a successful transport: create a recording interceptor
with withRecordingInterceptor(t), construct an intercepted client using a
passing client (e.g., successfulClient or a client that returns a valid
tikvrpc.Response), invoke client.SendRequestAsync(context.Background(), "",
newRGRequest(), cb) with an async callback that captures resp and err and closes
a done channel, wait on done, then assert resp is non-nil and err is nil and
assert rec.waitCalls == 1, rec.respCalls == 1 and rec.cancelCalls == 0 (with a
message like "OnRequestCancel must not run on the async path when resp is
non-nil"); place the new test alongside
TestSendRequestAsyncCancelsPreChargeOnTransportFailure and name it e.g.
TestSendRequestAsyncSettlesPreChargeOnTransportSuccess.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d7a03552-a257-480d-8f7b-28bf13e33018
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (3)
go.modinternal/client/client_interceptor.gointernal/client/client_interceptor_test.go
rleungx
reviewed
May 19, 2026
ti-chi-bot
Bot
added
dco-signoff: no
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
YuhaoZhang00
force-pushed
the
feat/rc-paging-precharge
branch
from
ea180e1 to
0760c69
Compare
ti-chi-bot
Bot
removed
the
dco-signoff: no
ti-chi-bot
Bot
added
the
dco-signoff: yes
…echarge Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com> # Conflicts: # go.sum
YuhaoZhang00
force-pushed
the
feat/rc-paging-precharge
branch
from
b7acdd3 to
4e7e340
Compare
rleungx
approved these changes
Jun 23, 2026
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
…echarge Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
JmPotato
approved these changes
Jul 8, 2026
ti-chi-bot
Bot
added
lgtm
and removed
needs-1-more-lgtm
cfzjywxk
approved these changes
Jul 8, 2026
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: cfzjywxk, JmPotato, rleungx The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
This was referenced Jul 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Issue Number: close #1953
What is changed and how it works?
Adds an optional caller-supplied read-bytes estimate that the resource-control layer forwards to PD's controller as a pre-charge basis for paging-style coprocessor requests.
tikvrpc.Request: newPredictedReadBytes uint64field. This is a client-go-internal hint, not carried in the wire proto; it is consumed before the RPC is sent and TiKV never observes it.internal/resourcecontrol.RequestInfo: newpredictedReadBytesandisCopfields populated inMakeRequestInfo, exposed throughPredictedReadBytes()andIsCop().Together these let an upstream caller (TiDB cop iterator, driven by a per-logical-scan EMA) attach a read-bytes prediction to a paging RPC. PD's resource-group controller consumes this hint via the matching
RequestInfomethods and pre-charges RU before the RPC goes out. When the hint is zero, behavior is unchanged and the request is billed from actual read bytes at settlement.This PR also adds targeted unit coverage for:
PredictedReadBytes()propagation fromtikvrpc.Requesttoresourcecontrol.RequestInfo.IsCop()classification forCmdCop/CmdCopStreamversus other read requests.PD client sync
This PR also bumps
github.com/tikv/pd/clientto a merged PD master version. That brings in PD client interface changes from other PD PRs, so client-go's mocks and integration-test helpers are updated in the same PR to keep the upgraded dependency buildable:8ed2b3fb8e3c,client: fix resource group config watch recovery): meta-storage/resource-managerWatchnow returnschan *metastorage.WatchResponseso compact-revision information can be propagated.17e9771e4b9e,keyspace: load keyspace by ID through gRPC): the PD client keyspace interface now includesLoadKeyspaceByID(ctx, id uint32).integration_tests/go.modtemporarily replacesgithub.com/pingcap/tidbandgithub.com/pingcap/tidb/pkg/parserwith an adapted fork because the current upstream TiDB dependency still uses the old PD client mock signatures. After the TiDB-side PR lands, a follow-up client-go PR should remove that replace and bump back to the official TiDB master pseudo-version.Related PRs
Check List
Validation: