locate: implement request attempt admission tracking and related statistics

[solotzg]

Problem

Issue Number: close #2035

TiDB needs to limit concurrent coprocessor requests per query and per TiKV
store. The target store is selected inside client-go and may change during
region retries, so acquiring admission before calling client-go cannot
accurately limit the actual target store.

The admission token should cover exactly one RPC attempt and must be released
before another attempt is admitted.

What changed

• Add tikvrpc.RequestAttemptAdmissionFunc and
  Request.RequestAttemptAdmission.
• Invoke the callback after selecting the actual target store and before every
  RPC attempt, including retries.
• Pass the target store ID rather than the proxy store ID.
• Invoke the returned release function exactly once after the attempt finishes.
• Release admission when store-token acquisition fails or when admission
  returns both a release function and an error.
• Run potentially blocking admission asynchronously in SendReqAsync.
• Exclude admission wait time from async RPC runtime statistics.
• Define admission timeout semantics: admission follows the request context;
  the per-attempt RPC timeout starts after admission succeeds.

When the callback is nil, existing request behavior is unchanged.

Error and cancellation semantics

• An admission error terminates the request without sending an RPC.
• Admission implementations must return promptly when the request context is
  canceled.
• Admission wait is not bounded by the per-attempt RPC timeout.

Tests

Added coverage for:

• synchronous and asynchronous admission;
• admission rejection and context cancellation;
• release on store-token failure;
• release exactly once per attempt;
• selecting the new store after a retry;
• releasing the previous attempt before retry admission;
• excluding admission wait from async RPC runtime statistics.

Tests run:

• go test ./internal/locate -run 'TestRegionRequestToSingleStore/TestRequestAttemptAdmission|TestRegionRequestToThreeStores/TestRequestAttemptAdmissionUsesRetryStore' -count=10
• go test -race ./internal/locate -run 'TestRegionRequestToSingleStore/TestRequestAttemptAdmission|TestRegionRequestToThreeStores/TestRequestAttemptAdmissionUsesRetryStore' -count=1

Compatibility

This is an additive client API. Existing callers are unaffected unless they set
RequestAttemptAdmission.

The associated TiDB change installs the callback only for TiKV coprocessor
requests.

Related

• distsql, kv: add query-level per-store cop request limiter pingcap/tidb#69360

[ti-chi-bot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign bufferflies for approval. For more information see the Code Review Process.
Please ensure that each of them provides their approval before proceeding.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[ti-chi-bot]

Hi @solotzg. Thanks for your PR.

I'm waiting for a tikv member to verify that this patch is reasonable to test. If it is, they should reply with /ok-to-test on its own line. Until that is done, I will not automatically test new commits in this PR, but the usual testing commands by org members will still work. Regular contributors should join the org to skip this step.

Once the patch is verified, the new status will be reflected by the ok-to-test label.

I understand the commands that are listed here.

Details

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.

[ti-chi-bot]

Welcome @solotzg!

It looks like this is your first PR to tikv/client-go 🎉.

I'm the bot to help you request reviewers, add labels and more, See available commands.

We want to make sure your contribution gets all the attention it needs!



Thank you, and welcome to tikv/client-go. 😃

[coderabbitai]

📝 Walkthrough

Walkthrough

Request-attempt admission is added to the RPC request contract and integrated into asynchronous and retry execution. Admission release handling covers successful attempts, errors, cancellation, RPC statistics, and store-limit termination.

Changes

Request Attempt Admission

Layer / File(s)	Summary
Admission contract and execution flow tikvrpc/tikvrpc.go, internal/locate/region_request.go	Adds the admission callback contract and acquires admission for selected stores before asynchronous and retried RPC attempts, releasing capacity when attempts finish or fail.
Admission behavior and retry validation internal/locate/region_request_test.go, internal/locate/region_request3_test.go	Tests successful and rejected admission, cancellation, RPC-stat timing, release counts, store-limit handling, and retry-store selection for synchronous and asynchronous requests.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Request
  participant SendReqAsync
  participant sendReqState
  participant RequestAttemptAdmission
  participant RPC
  Request->>SendReqAsync: submit request
  SendReqAsync->>sendReqState: acquire admission for selected store
  sendReqState->>RequestAttemptAdmission: admit RPC attempt
  RequestAttemptAdmission-->>sendReqState: release function or error
  sendReqState->>RPC: send admitted attempt
  RPC-->>sendReqState: return attempt result
  sendReqState->>RequestAttemptAdmission: release admission

Loading

Suggested reviewers: ekexium, lcwangchao

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the PR’s main change: request-attempt admission tracking and related locate statistics.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@internal/locate/region_request.go`:
- Around line 953-971: Update acquireRequestAttemptAdmission to apply
s.args.timeout when invoking req.RequestAttemptAdmission, ensuring admission
cannot outlive the per-attempt RPC deadline. Use the existing request context as
the parent, preserve cancellation/error propagation, and retain the current
wait-time recording and release handling.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ba93397c-a8aa-4bcf-ba80-616225ab0c5a

📥 Commits

Reviewing files that changed from the base of the PR and between 3d4b3ea and b4f7076.

📒 Files selected for processing (4)

• internal/locate/region_request.go
• internal/locate/region_request3_test.go
• internal/locate/region_request_test.go
• tikvrpc/tikvrpc.go

[solotzg]

/cc @windtalker

[ti-chi-bot]

@solotzg: GitHub didn't allow me to request PR reviews from the following users: windtalker.

Note that only tikv members and repo collaborators can review this PR, and authors cannot review their own PRs.

Details

In response to this:

/cc @windtalker

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.

[solotzg]

/cc @windtalker @cfzjywxk

[ti-chi-bot]

@solotzg: GitHub didn't allow me to request PR reviews from the following users: windtalker.

Note that only tikv members and repo collaborators can review this PR, and authors cannot review their own PRs.

Details

In response to this:

/cc @windtalker @cfzjywxk

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.