resourcecontrol, tikvrpc: add PredictedReadBytes hint for RC paging pre-charge#1947
resourcecontrol, tikvrpc: add PredictedReadBytes hint for RC paging pre-charge#1947YuhaoZhang00 wants to merge 7 commits into
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)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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
![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
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>
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>
| bypass: bypass, | ||
| requestSize: uint64(req.GetSize()), | ||
| accessType: toPDAccessLocationType(req.AccessLocation), | ||
| predictedReadBytes: req.PredictedReadBytes, |
There was a problem hiding this comment.
With PD pre-charge enabled, this hint needs a settlement path even when the RPC fails before returning a response. Today OnResponseWait is only called when resp != nil, so a transport error/timeout after pre-charge can leave the predicted read RU charged without refund. Can we add a no-response settlement/refund path or a test that documents this behavior?
There was a problem hiding this comment.
Addressed by adding calling resourceControlInterceptor.OnRequestCancel() which does refund when resp == nil
Commit 54a00f4
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>
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
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
| // Transport-level failure produced no response — the settlement | ||
| // path will not run. Cancel the pre-charge so the predicted RU | ||
| // debited in OnRequestWait is refunded to the limiter. | ||
| resourceControlInterceptor.OnRequestCancel(ctx, resourceGroupName, reqInfo) |
There was a problem hiding this comment.
On the no-response path we refund the pre-charge in the PD controller, but RUDetails has already been updated with the OnRequestWait consumption above. Because OnRequestCancel does not return a refund delta and this branch does not apply a negative update, TiDB runtime RU stats will still include the predicted/base RU from the failed attempt. If the request is retried and later succeeds, the failed attempt can be counted as phantom RU in runtime stats even though controller tokens/consumption were rolled back. The async cancel path below has the same issue. Could we either have OnRequestCancel return the refund delta and update RUDetails, or keep the request-side consumption here and subtract it on cancel?
2 participants
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 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: newpredictedReadBytesfield populated fromtikvrpc.Request.PredictedReadBytesinMakeRequestInfo, exposed via aPredictedReadBytes()getter.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 required
PredictedReadBytes()method on itsRequestInfointerface and pre-charges RU before the RPC goes out; see the matching PD PR. When the hint is zero (the default for any caller that doesn't set the field, including all writes), behavior is unchanged — the request is billed by actual read bytes at settlement.Related PRs
Check List
Summary by CodeRabbit
New Features
Bug Fixes
Tests