fix(resourcecontrol): bypass RU accounting for NextGen background activities#1930
fix(resourcecontrol): bypass RU accounting for NextGen background activities#1930YuhaoZhang00 wants to merge 2 commits into
Conversation
ti-chi-bot
Bot
added
do-not-merge/work-in-progress
📝 WalkthroughWalkthroughThe change adds internal request-source constants, expands NextGen resource-control bypass matching for background operations and analyze requests, preserves the existing fallback behavior, and adds table-driven tests for the resulting bypass decisions. ChangesResource-control bypasses
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 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
contribution
ti-chi-bot
Bot
removed
the
do-not-merge/work-in-progress
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
util/request_source.go (1)
44-53: Consider renamingInternalImportIntotoInternalTxnImportIntofor naming consistency.The other constants in this block follow the
InternalTxn*naming pattern (e.g.,InternalTxnBR,InternalTxnAddIndex,InternalTxnWorkloadLearning), butInternalImportIntodoes not include theTxninfix. This inconsistency may cause confusion when developers search for or reference these constants.🔧 Suggested rename
- // InternalImportInto is the type of IMPORT INTO usage - InternalImportInto = "ImportInto" + // InternalTxnImportInto is the type of IMPORT INTO usage + InternalTxnImportInto = "ImportInto"Note: If renamed, update the reference in
internal/resourcecontrol/resource_control.go(bypassResourceSourceList) accordingly.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@util/request_source.go` around lines 44 - 53, The constant name InternalImportInto is inconsistent with the InternalTxn* pattern; rename InternalImportInto to InternalTxnImportInto across the codebase (replace occurrences of InternalImportInto with InternalTxnImportInto) and update any references such as bypassResourceSourceList in internal/resourcecontrol/resource_control.go to use InternalTxnImportInto; ensure you update imports/usage sites and run tests to verify no remaining references to InternalImportInto.internal/client/client_interceptor.go (1)
158-170: Consider defining constants for RU type labels.The string literals
"rru"and"wru"are used directly. For consistency with other label constants in the metrics package and to avoid typos in future uses, consider defining these as constants.🔧 Suggested constants (in metrics/metrics.go)
const ( // ... existing constants ... LblRRU = "rru" LblWRU = "wru" )Then use:
- metrics.TiKVResourceControlRUCounter.WithLabelValues(resourceGroupName, requestSource, "rru").Add(consumption.RRU) + metrics.TiKVResourceControlRUCounter.WithLabelValues(resourceGroupName, requestSource, metrics.LblRRU).Add(consumption.RRU)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/client/client_interceptor.go` around lines 158 - 170, The function recordResourceControlMetrics uses hardcoded RU label strings ("rru", "wru"); define constants in the metrics package (e.g., LblRRU = "rru" and LblWRU = "wru" in metrics/metrics.go) and replace the literals in recordResourceControlMetrics (and any other usages) to use metrics.LblRRU and metrics.LblWRU when calling metrics.TiKVResourceControlRUCounter.WithLabelValues to ensure consistency and avoid typos.internal/resourcecontrol/resource_control.go (1)
98-103: Consider edge cases with substring matching.The
strings.Containsapproach could theoretically cause false positives if a future request source contains a bypass source as a substring (e.g., a hypothetical "no_add_index" would match "add_index"). Given the current naming conventions in TiDB, this seems unlikely, but worth noting for future reference.🤖 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 98 - 103, The current substring check using strings.Contains(requestSource, source) can produce false positives; update the loop in resource_control.go (bypassResourceSourceList and requestSource) to use a safer comparison such as exact equality (requestSource == source) or, if requestSource can contain multiple tokens, split requestSource into tokens and check membership, or use boundaries (strings.HasPrefix/HasSuffix with delimiters) so only whole-source matches bypass; implement one of these approaches in the loop to replace strings.Contains.
🤖 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/client/client_interceptor.go`:
- Around line 158-170: The function recordResourceControlMetrics uses hardcoded
RU label strings ("rru", "wru"); define constants in the metrics package (e.g.,
LblRRU = "rru" and LblWRU = "wru" in metrics/metrics.go) and replace the
literals in recordResourceControlMetrics (and any other usages) to use
metrics.LblRRU and metrics.LblWRU when calling
metrics.TiKVResourceControlRUCounter.WithLabelValues to ensure consistency and
avoid typos.
In `@internal/resourcecontrol/resource_control.go`:
- Around line 98-103: The current substring check using
strings.Contains(requestSource, source) can produce false positives; update the
loop in resource_control.go (bypassResourceSourceList and requestSource) to use
a safer comparison such as exact equality (requestSource == source) or, if
requestSource can contain multiple tokens, split requestSource into tokens and
check membership, or use boundaries (strings.HasPrefix/HasSuffix with
delimiters) so only whole-source matches bypass; implement one of these
approaches in the loop to replace strings.Contains.
In `@util/request_source.go`:
- Around line 44-53: The constant name InternalImportInto is inconsistent with
the InternalTxn* pattern; rename InternalImportInto to InternalTxnImportInto
across the codebase (replace occurrences of InternalImportInto with
InternalTxnImportInto) and update any references such as
bypassResourceSourceList in internal/resourcecontrol/resource_control.go to use
InternalTxnImportInto; ensure you update imports/usage sites and run tests to
verify no remaining references to InternalImportInto.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9424dd89-6885-425f-95f8-0aec17877926
📒 Files selected for processing (4)
internal/client/client_interceptor.gointernal/resourcecontrol/resource_control.gometrics/metrics.goutil/request_source.go
| } | ||
| recordResourceControlMetrics(resourceGroupName, req.GetRequestSource(), consumption) | ||
| } else if reqInfo != nil && reqInfo.Bypass() { | ||
| metrics.TiKVResourceControlBypassedCounter.WithLabelValues(resourceGroupName, req.GetRequestSource()).Inc() |
There was a problem hiding this comment.
Avoid introducing WithLabelValues op on the hot path, should cache it in advance.
| return resp, err | ||
| }) | ||
| } else if reqInfo != nil && reqInfo.Bypass() { | ||
| metrics.TiKVResourceControlBypassedCounter.WithLabelValues(resourceGroupName, req.GetRequestSource()).Inc() |
| metrics.TiKVResourceControlRUCounter.WithLabelValues(resourceGroupName, requestSource, "rru").Add(consumption.RRU) | ||
| } | ||
| if consumption.WRU > 0 { | ||
| metrics.TiKVResourceControlRUCounter.WithLabelValues(resourceGroupName, requestSource, "wru").Add(consumption.WRU) |
ti-chi-bot
Bot
added
the
do-not-merge/work-in-progress
|
Remove the client-go local resource-control metrics introduced in earlier iterations |
ti-chi-bot
Bot
removed
the
do-not-merge/work-in-progress
Signed-off-by: JmPotato <github@ipotato.me>
Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>
993380a to
d7e0184
Compare
|
[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
the
do-not-merge/work-in-progress
There was a problem hiding this comment.
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/resourcecontrol/resource_control.go`:
- Around line 106-112: Prefix each NextGen resource-source match in
internal/resourcecontrol/resource_control.go:106-112 with
util.InternalRequestPrefix before checking requestSource, and apply the same
prefix to util.InternalTxnStats at
internal/resourcecontrol/resource_control.go:94-95. Update RequestSource values
in internal/resourcecontrol/resource_control_test.go:65-131 to use
util.InternalRequestPrefix with the relevant internal source constants.
🪄 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: 1d5f6a46-9046-4b4b-88af-6a1ded001bcd
📒 Files selected for processing (3)
internal/resourcecontrol/resource_control.gointernal/resourcecontrol/resource_control_test.goutil/request_source.go
🚧 Files skipped from review as they are similar to previous changes (1)
- util/request_source.go
| // Check other resource source types that should be bypassed. | ||
| for _, source := range bypassResourceSourceList { | ||
| if strings.Contains(requestSource, source) { | ||
| return true | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy lift
Resource-control bypass vulnerability via user-controlled explicit_request_source.
The NextGen strings.Contains checks for internal request sources omit the "internal_" prefix. Because the requestSource string is a composite of labels (e.g., <origin>_<source>_<explicit_source>), checking for bare substrings like "br" or "add_index" will incorrectly bypass resource accounting for external user requests if they happen to use those strings in their explicit source. For example, an external request with explicit_request_source set to "library" or "zebra" yields a requestSource like "external_unknown_library", which contains "br" and is erroneously bypassed.
The fallback logic on line 117 correctly prevents this by prepending util.InternalRequestPrefix. Apply this prefix to the NextGen checks and their corresponding tests:
internal/resourcecontrol/resource_control.go#L106-L112: Prependutil.InternalRequestPrefixtosourcewhen matchingrequestSource.internal/resourcecontrol/resource_control.go#L94-L95: Prependutil.InternalRequestPrefixtoutil.InternalTxnStatswhen matching.internal/resourcecontrol/resource_control_test.go#L65-L131: Update theRequestSourceassignments in the test cases to include the prefix (util.InternalRequestPrefix + util.InternalTxnAddIndex, etc.) so they pass the corrected matching logic.
🛡️ Proposed fixes for the matching logic
- if strings.Contains(requestSource, util.InternalTxnStats) {
+ if strings.Contains(requestSource, util.InternalRequestPrefix+util.InternalTxnStats) {
var tp int64
switch req.Type {
case tikvrpc.CmdBatchCop:
@@ -105,7 +105,7 @@
}
// Check other resource source types that should be bypassed.
for _, source := range bypassResourceSourceList {
- if strings.Contains(requestSource, source) {
+ if strings.Contains(requestSource, util.InternalRequestPrefix+source) {
return true
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Check other resource source types that should be bypassed. | |
| for _, source := range bypassResourceSourceList { | |
| if strings.Contains(requestSource, source) { | |
| return true | |
| } | |
| } | |
| } | |
| // Check other resource source types that should be bypassed. | |
| for _, source := range bypassResourceSourceList { | |
| if strings.Contains(requestSource, util.InternalRequestPrefix+source) { | |
| return true | |
| } | |
| } | |
| } |
📍 Affects 2 files
internal/resourcecontrol/resource_control.go#L106-L112(this comment)internal/resourcecontrol/resource_control.go#L94-L95internal/resourcecontrol/resource_control_test.go#L65-L131
🤖 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/resourcecontrol/resource_control.go` around lines 106 - 112, Prefix
each NextGen resource-source match in
internal/resourcecontrol/resource_control.go:106-112 with
util.InternalRequestPrefix before checking requestSource, and apply the same
prefix to util.InternalTxnStats at
internal/resourcecontrol/resource_control.go:94-95. Update RequestSource values
in internal/resourcecontrol/resource_control_test.go:65-131 to use
util.InternalRequestPrefix with the relevant internal source constants.
2 participants
Closes #1954
Tracking issue (TiDB side): pingcap/tidb#64339
Supersedes #1809.
What problem does this PR solve?
Under the
nextgenbuild tag, resource control should bypass selected heavy internal activities so they do not consume the user-facing RU quota.What is changed and how does it work?
Expand
shouldBypass()for the following NextGen request sources:add_indexandmerge_temp_indexbrImportIntoWorkloadLearningThe existing internal Analyze and
internal_othersbypass behavior is preserved. This PR changes bypass behavior only; request-source RU observability is tracked separately by #1929 and tikv/pd#10588.The original implementation commit from #1809 remains in this PR history with its original author attribution.
Tests
Earlier manual verification on a local NextGen cluster confirmed that
add_indexand internal Analyze requests were bypassed without accruing RU, while normal user queries and non-Analyze stats operations continued to consume RU.Release note