txnkv: support shared lock upgrades - #2014
Conversation
|
Skipping CI for Draft Pull Request. |
ti-chi-bot
Bot
added
the
do-not-merge/work-in-progress
|
[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
dco-signoff: yes
📝 WalkthroughWalkthroughThe PR adds shared-lock upgrade support for eligible pessimistic transactions. It introduces typed lock errors, decodes and redacts their keys, tracks fatal and undetermined transaction states, and adds transaction and codec tests. ChangesShared-lock upgrade support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Transaction
participant lockKeys
participant twoPhaseCommitter
participant TiKV
Transaction->>lockKeys: submit lock requests
lockKeys->>twoPhaseCommitter: initialize transaction state
twoPhaseCommitter->>TiKV: acquire ordinary locks
twoPhaseCommitter->>TiKV: upgrade shared locks
TiKV-->>twoPhaseCommitter: return lock result or typed error
twoPhaseCommitter-->>Transaction: update state or return error
Possibly related PRs
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
the
size/XXL
wfxr
changed the title
f501e4b to
2741a6a
Compare
837a720 to
afb7539
Compare
ti-chi-bot
Bot
removed
the
do-not-merge/work-in-progress
|
/retest |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
go.mod (1)
18-38: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftDo not merge the fork
replacedirective.
replacedirectives apply only to the main module. Downstream modules that importgithub.com/tikv/client-go/v2resolvegithub.com/pingcap/kvprotoat the version in Line 18, which does not containLockUpgradeConflictorSharedLockLost. Those builds fail to compileerror/error.go,internal/apicodec/codec_v2.go, andutil/redact/redact.go.Block this PR until pingcap/kvproto#1495 merges, then pin the upstream pseudo-version and remove the
replaceline.🤖 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 `@go.mod` around lines 18 - 38, Remove the github.com/pingcap/kvproto replace directive and do not merge the fork. After pingcap/kvproto#1495 is merged, update the github.com/pingcap/kvproto dependency to the corresponding upstream pseudo-version containing LockUpgradeConflict and SharedLockLost, ensuring downstream modules compile without relying on the local replacement.
🧹 Nitpick comments (6)
error/error.go (1)
360-371: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the asymmetric ordering of the two new checks.
SharedLockLostis checked beforeConflict, butLockUpgradeConflictis checked after it. The precedence is deliberate:error/error_test.goLine 44 setsConflicttogether withSharedLockLostand requires the shared-lock-lost type. The reasonLockUpgradeConflictsits belowConflictis not stated. If TiKV can set bothConflictandLockUpgradeConflictin oneKeyError, the caller receivesErrWriteConflict, andtxn.goisLockUpgradeResultUndeterminedthen treats the outcome as deterministic through the write-conflict branch instead of the lock-upgrade branch. Add a short comment that records the intended precedence.🤖 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 `@error/error.go` around lines 360 - 371, Add a concise comment in the error-conversion logic around SharedLockLost, Conflict, and LockUpgradeConflict documenting that their asymmetric ordering is intentional: SharedLockLost takes precedence over Conflict, while Conflict takes precedence over LockUpgradeConflict so combined errors follow the write-conflict handling path. Preserve the existing check order and behavior.txnkv/transaction/txn.go (3)
1595-1604: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winDocument the per-key RPC cost and the partial-failure state.
The loop sends one
PessimisticLockRPC per upgrade key, sequentially, on the calling thread. For N upgrade keys the statement pays N round trips.txnkv/transaction/txn_test.goLine 411 asserts one key per upgrade request, so the batching is intentional.The loop also leaves a partial state on failure. If the first upgrade key succeeds and the second fails, the first key is exclusive, the second stays shared, and the error returns. No test covers two upgrade keys with a failure on the second.
Add a comment that states both properties, and add a test for the multi-upgrade-key partial failure.
🤖 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 `@txnkv/transaction/txn.go` around lines 1595 - 1604, Document the upgrade loop around txn.lockPessimisticKeyGroup to state that it performs one sequential PessimisticLock RPC per upgrade key and that failures can leave earlier keys upgraded while later keys remain shared. Add a transaction test covering multiple upgrade keys where the second upgrade fails, asserting the first is exclusive, the second remains shared, and the error is returned.
1481-1510: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffThe non-upgrade error path duplicates the legacy path.
Lines 1481-1510 repeat the deadlock detection, the
rollbackForUpdateTSselection, the async rollback, and the retry sleep from Lines 1876-1917. The two copies already differ: the legacy copy also adjuststxn.lockedCntand clearsaggressiveLockingContext.currentLockedKeys. Two copies of the same failure protocol will drift further.Extract the shared failure handling into one helper, then let both call sites pass the parts that differ (
allKeysversuskeys, and the aggressive-locking cleanup).🤖 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 `@txnkv/transaction/txn.go` around lines 1481 - 1510, Extract the duplicated pessimistic-lock failure handling from the current block and the legacy path into a shared helper, preserving deadlock detection, callback invocation, rollbackForUpdateTS selection, async rollback, and retry sleeps. Have both callers provide their respective key collection (`allKeys` or `keys`) and an optional aggressive-locking cleanup so the legacy path still updates `txn.lockedCnt` and clears `aggressiveLockingContext.currentLockedKeys` without duplicating the protocol.
802-807: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExplain why this gate runs before
defer txn.close().This block intentionally differs from
getTxnStateErrat Line 1396. It returns only the fatal error, and it runs beforedefer txn.close()on Line 808, so the transaction stays valid and the caller can still callRollback.txnkv/transaction/txn_test.goLine 626 asserts that behavior. The undetermined case falls through to Line 859, which runs afterdefer txn.close()and therefore invalidates the transaction, as asserted attxnkv/transaction/txn_test.goLine 750.Nothing in the code records this distinction. A later refactor that replaces this block with
getTxnStateErr()would remove rollback availability after a fatal error. Add a comment.♻️ Proposed comment
+ // Check the fatal error before `defer txn.close()` below, so a poisoned + // transaction stays valid and the caller can still roll back the locks it + // holds. The undetermined case is handled after `defer txn.close()`, because + // an undetermined result must invalidate the transaction. if txn.committer != nil { undeterminedErr, fatalTxnErr := txn.committer.getTxnStateErrs() if undeterminedErr == nil && fatalTxnErr != nil { return fatalTxnErr } }🤖 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 `@txnkv/transaction/txn.go` around lines 802 - 807, Add a concise explanatory comment immediately before the txn.committer fatal-error gate, documenting that it must run before defer txn.close() so fatal errors preserve transaction validity for caller Rollback, while undetermined errors continue to the later post-close path; keep the existing control flow unchanged.txnkv/transaction/txn_test.go (1)
323-336: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe recorder aliases one backing array.
Line 323 allocates
keyswithlen(lockReq.Mutations). Line 329 storeskeys[:0]in the summary. Line 335 then appends into that same backing array while the loop still writeskeys[i]at Line 325. Each append overwrites the slot that the same iteration just filled with the identical value, so the result is correct today. The aliasing is not obvious, and it breaks if the write order or the append order changes.Build the summary after the loop instead.
♻️ Proposed simplification
case tikvrpc.CmdPessimisticLock: lockReq := req.PessimisticLock() keys := make([][]byte, len(lockReq.Mutations)) for i, mutation := range lockReq.Mutations { keys[i] = append([]byte(nil), mutation.Key...) - if i == 0 { - requests = append(requests, requestSummary{ - cmd: req.Type, - keys: keys[:0], - op: mutation.Op, - }) - } else { - require.Equal(t, requests[len(requests)-1].op, mutation.Op) - } - requests[len(requests)-1].keys = append(requests[len(requests)-1].keys, keys[i]) + if i > 0 { + require.Equal(t, lockReq.Mutations[0].Op, mutation.Op) + } } + requests = append(requests, requestSummary{ + cmd: req.Type, + keys: keys, + op: lockReq.Mutations[0].Op, + }) return onLock(len(requests)-1, lockReq)🤖 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 `@txnkv/transaction/txn_test.go` around lines 323 - 336, Refactor the request recording logic around lockReq.Mutations so it first collects each mutation key and validates operation consistency during the loop, then creates and appends a single requestSummary after the loop. Do not store keys[:0] or append to the summary while keys is still being populated; ensure the summary owns the completed key collection.txnkv/transaction/2pc.go (1)
2439-2443: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueName the two results of
getTxnStateErrs.The function returns two values of the same
errortype. Callers attxnkv/transaction/txn.goLine 803 and Line 1400 depend on positional order. Named results document the order at the declaration and prevent a silent swap during a future edit.♻️ Proposed signature
-func (c *twoPhaseCommitter) getTxnStateErrs() (error, error) { +func (c *twoPhaseCommitter) getTxnStateErrs() (undeterminedErr error, fatalTxnErr error) { c.mu.RLock() defer c.mu.RUnlock() return c.mu.undeterminedErr, c.mu.fatalTxnErr }🤖 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 `@txnkv/transaction/2pc.go` around lines 2439 - 2443, Update the twoPhaseCommitter.getTxnStateErrs signature to use descriptive named error results that identify the undetermined and fatal transaction errors in positional order. Keep the existing locking and return expressions unchanged, ensuring callers such as txn.go continue receiving undeterminedErr first and fatalTxnErr second.
🤖 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 `@kv/kv.go`:
- Line 70: Add a clear Go doc comment immediately above the exported LockCtx
field AllowSharedLockUpgrade, documenting that it permits shared-lock upgrades,
the upgrade is rejected in aggressive/fair locking mode, and failure may leave
the transaction undetermined or fatally poisoned.
In `@txnkv/transaction/txn_test.go`:
- Around line 511-559: Add a subtest alongside
LockUpgradeConflictReturnsTypedErrorWithoutRetrySemantics that makes the upgrade
request return a Deadlock key error, configures lockCtx.OnDeadlock, and verifies
the callback is invoked with dl.IsRetryable set. Use the existing recorder
transaction and upgrade-lock setup, and assert the deadlock-specific branch
behavior through lockKeys.
In `@txnkv/transaction/txn.go`:
- Around line 1468-1479: The upgrade error handling in
txnkv/transaction/txn.go:1468-1479 must handle tikverr.ErrDeadlock before the
existing undetermined-result checks: set dl.IsRetryable via hashInKeys and
invoke lockCtx.OnDeadlock(dl), without adding pessimistic rollback so the
retained shared lock remains intact. Add a subtest in
txnkv/transaction/txn_test.go:511-559 that produces a Deadlock key error during
upgrade and verifies OnDeadlock is called and IsRetryable is set.
- Around line 1516-1519: In the loop using lockCtx.Values, remove the single-use
keyStr variable and inline the string conversion directly in the map lookup,
changing the access to use lockCtx.Values[string(key)] while preserving the
existing valExists and value-handling behavior.
- Around line 1778-1784: Guard the shared-lock upgrade path in the surrounding
transaction lock flow with txn.committer != nil && lockCtx.ForUpdateTS > 0
before calling lockKeysWithSharedLockUpgrade. Preserve the existing normal
pessimistic-lock branch, and when ForUpdateTS is zero update only the local
memory-buffer upgrade flag instead of sending upgrade RPCs.
- Around line 1565-1575: Update the primary-selection block around
selectPrimaryForPessimisticLock so it never uses upgradeKeys as a fallback
candidate when normalExclusiveKeys is empty. Require at least one normal
exclusive key before selecting the pessimistic primary, or explicitly return the
established error for a shared-only candidate before issuing the upgrade RPC.
---
Outside diff comments:
In `@go.mod`:
- Around line 18-38: Remove the github.com/pingcap/kvproto replace directive and
do not merge the fork. After pingcap/kvproto#1495 is merged, update the
github.com/pingcap/kvproto dependency to the corresponding upstream
pseudo-version containing LockUpgradeConflict and SharedLockLost, ensuring
downstream modules compile without relying on the local replacement.
---
Nitpick comments:
In `@error/error.go`:
- Around line 360-371: Add a concise comment in the error-conversion logic
around SharedLockLost, Conflict, and LockUpgradeConflict documenting that their
asymmetric ordering is intentional: SharedLockLost takes precedence over
Conflict, while Conflict takes precedence over LockUpgradeConflict so combined
errors follow the write-conflict handling path. Preserve the existing check
order and behavior.
In `@txnkv/transaction/2pc.go`:
- Around line 2439-2443: Update the twoPhaseCommitter.getTxnStateErrs signature
to use descriptive named error results that identify the undetermined and fatal
transaction errors in positional order. Keep the existing locking and return
expressions unchanged, ensuring callers such as txn.go continue receiving
undeterminedErr first and fatalTxnErr second.
In `@txnkv/transaction/txn_test.go`:
- Around line 323-336: Refactor the request recording logic around
lockReq.Mutations so it first collects each mutation key and validates operation
consistency during the loop, then creates and appends a single requestSummary
after the loop. Do not store keys[:0] or append to the summary while keys is
still being populated; ensure the summary owns the completed key collection.
In `@txnkv/transaction/txn.go`:
- Around line 1595-1604: Document the upgrade loop around
txn.lockPessimisticKeyGroup to state that it performs one sequential
PessimisticLock RPC per upgrade key and that failures can leave earlier keys
upgraded while later keys remain shared. Add a transaction test covering
multiple upgrade keys where the second upgrade fails, asserting the first is
exclusive, the second remains shared, and the error is returned.
- Around line 1481-1510: Extract the duplicated pessimistic-lock failure
handling from the current block and the legacy path into a shared helper,
preserving deadlock detection, callback invocation, rollbackForUpdateTS
selection, async rollback, and retry sleeps. Have both callers provide their
respective key collection (`allKeys` or `keys`) and an optional
aggressive-locking cleanup so the legacy path still updates `txn.lockedCnt` and
clears `aggressiveLockingContext.currentLockedKeys` without duplicating the
protocol.
- Around line 802-807: Add a concise explanatory comment immediately before the
txn.committer fatal-error gate, documenting that it must run before defer
txn.close() so fatal errors preserve transaction validity for caller Rollback,
while undetermined errors continue to the later post-close path; keep the
existing control flow unchanged.
🪄 Autofix
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 Plus
Run ID: 715b2c7b-1d02-4fd4-b301-2ff1cc9c350c
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (11)
error/error.goerror/error_test.gogo.modinternal/apicodec/codec_v2.gointernal/apicodec/codec_v2_test.gokv/kv.gotxnkv/transaction/2pc.gotxnkv/transaction/txn.gotxnkv/transaction/txn_test.goutil/redact/redact.goutil/redact/redact_test.go
| CheckExistence bool | ||
| LockOnlyIfExists bool | ||
| InShareMode bool | ||
| AllowSharedLockUpgrade bool |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add a doc comment to the new public field.
LockCtx is part of the public client API. AllowSharedLockUpgrade gates behavior that callers cannot infer from the name: an upgrade failure can leave the transaction undetermined or fatally poisoned, and the upgrade is rejected in aggressive/fair locking mode.
As per coding guidelines: "Exported identifiers must have clear Go doc comments when they are part of the public client API."
📝 Proposed doc comment
InShareMode bool
+ // AllowSharedLockUpgrade allows a key that this transaction already holds in
+ // shared mode to be upgraded to an exclusive lock. Upgrade requests are sent
+ // separately from ordinary exclusive locks. A deterministic upgrade failure
+ // keeps the existing shared lock. Any other upgrade failure marks the
+ // transaction result undetermined or fatal. Upgrades are rejected in
+ // aggressive/fair locking mode.
AllowSharedLockUpgrade bool📝 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.
| AllowSharedLockUpgrade bool | |
| InShareMode bool | |
| // AllowSharedLockUpgrade allows a key that this transaction already holds in | |
| // shared mode to be upgraded to an exclusive lock. Upgrade requests are sent | |
| // separately from ordinary exclusive locks. A deterministic upgrade failure | |
| // keeps the existing shared lock. Any other upgrade failure marks the | |
| // transaction result undetermined or fatal. Upgrades are rejected in | |
| // aggressive/fair locking mode. | |
| AllowSharedLockUpgrade bool |
🤖 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 `@kv/kv.go` at line 70, Add a clear Go doc comment immediately above the
exported LockCtx field AllowSharedLockUpgrade, documenting that it permits
shared-lock upgrades, the upgrade is rejected in aggressive/fair locking mode,
and failure may leave the transaction undetermined or fatally poisoned.
Source: Coding guidelines
| if isUpgrade { | ||
| var sharedLockLost *tikverr.ErrSharedLockLost | ||
| if errors.As(err, &sharedLockLost) { | ||
| txn.committer.setFatalTxnErr(err) | ||
| return 0, err | ||
| } | ||
| if isLockUpgradeResultUndetermined(err) { | ||
| txn.committer.setUndeterminedErr(err) | ||
| return 0, errors.WithStack(tikverr.ErrResultUndetermined) | ||
| } | ||
| return 0, err | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The upgrade error path omits the deadlock protocol, and no test detects it. isLockUpgradeResultUndetermined classifies *tikverr.ErrDeadlock as deterministic, so an upgrade that deadlocks returns the error directly without marking dl.IsRetryable and without calling lockCtx.OnDeadlock. The ordinary lock path performs both steps. No subtest covers a deadlock on an upgrade, so the gap is invisible.
txnkv/transaction/txn.go#L1468-L1479: in theisUpgradebranch, detect*tikverr.ErrDeadlock, setdl.IsRetryablethroughhashInKeys, and calllockCtx.OnDeadlock(dl). Do not add the pessimistic rollback, because the rollback would drop the retained shared lock.txnkv/transaction/txn_test.go#L511-L559: add a subtest that returns aDeadlockkey error for the upgrade request, then assert thatlockCtx.OnDeadlockruns and thatIsRetryableis set.
📍 Affects 2 files
txnkv/transaction/txn.go#L1468-L1479(this comment)txnkv/transaction/txn_test.go#L511-L559
🤖 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 `@txnkv/transaction/txn.go` around lines 1468 - 1479, The upgrade error
handling in txnkv/transaction/txn.go:1468-1479 must handle tikverr.ErrDeadlock
before the existing undetermined-result checks: set dl.IsRetryable via
hashInKeys and invoke lockCtx.OnDeadlock(dl), without adding pessimistic
rollback so the retained shared lock remains intact. Add a subtest in
txnkv/transaction/txn_test.go:511-559 that produces a Deadlock key error during
upgrade and verifies OnDeadlock is called and IsRetryable is set.
| if txn.committer.primaryKey == nil { | ||
| // Prefer selecting the primary from freshly requested exclusive locks | ||
| // when possible, so the primary is not just a shared-locked key that | ||
| // still depends on upgrade success. | ||
| assignedPrimaryKey = true | ||
| keysForPrimary := normalExclusiveKeys | ||
| if len(keysForPrimary) == 0 { | ||
| keysForPrimary = upgradeKeys | ||
| } | ||
| txn.selectPrimaryForPessimisticLock(keysForPrimary) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect primary selection, resetPrimary, and ttlManager interaction.
set -euo pipefail
ast-grep run --pattern $'func (txn *KVTxn) selectPrimaryForPessimisticLock($$$) {
$$$
}' --lang go txnkv/transaction/txn.go
ast-grep run --pattern $'func (txn *KVTxn) resetPrimary($$$) {
$$$
}' --lang go txnkv/transaction/txn.go
rg -n -C 3 'setPrimary|run\(|ttlManager' txnkv/transaction/txn.goRepository: tikv/client-go
Length of output: 8280
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "txnkv/transaction/txn.go relevant sections:"
sed -n '1530,1615p' txnkv/transaction/txn.go
sed -n '1780,1935p' txnkv/transaction/txn.go
sed -n '2038,2064p' txnkv/transaction/txn.go
echo
echo "LockCtx constructors/callers around pessimistic lock APIs:"
rg -n -C 4 '\bLockKeys\b|\bForUpdateTS\b|\bLockOnlyIfExists\b' . --glob '*.go' | head -n 220
echo
echo "lockPessimisticKeyGroup implementation:"
ast-grep run --pattern $'func (txn *KVTxn) lockPessimisticKeyGroup($$$) {
$$$
}' --lang go txnkv/transaction/txn.goRepository: tikv/client-go
Length of output: 27148
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find references to lockKeys/lockKeysWithSharedLockUpgrade and InShareMode:"
rg -n -C 5 '\blogInShareMode|InShareMode|lockKeysWithSharedLockUpgrade|lock.*Upgrade|PessimisticLockWakeUpMode_WakeUp'.*InShareMode|LockType|LockType\.' txnkv/transaction *.go integration_tests/shared_lock_test.go 2>/dev/null | head -n 360
echo
echo "Locate two-phase committer methods:"
rg -n -C 4 'func .*pessimisticLockMutations|func .*run\(|startTTLManager|stopTTLManager|reset\(' txnkv/transaction -g '*.go' | head -n 260Repository: tikv/client-go
Length of output: 287
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find references to lockKeys/lockKeysWithSharedLockUpgrade and InShareMode:"
rg -n -C 5 'lockKeysWithSharedLockUpgrade|InShareMode|PessimisticLockWakeUpMode|LockType|LockType\.' txnkv/transaction integration_tests 2>/dev/null | head -n 360
echo
echo "Locate two-phase committer methods:"
rg -n -C 4 'func .*pessimisticLockMutations|func .*run\(|startTTLManager|stopTTLManager|reset\(' txnkv/transaction -g '*.go' | head -n 260
echo
echo "Find upgrade-related tests or examples:"
rg -n -C 4 'SharedLock|shared lock|upgrade|ForUpdateTS|sharedlocked|locked.*shared|InShareMode|StartAggressiveLocking' integration_tests -g '*.go' | head -n 420Repository: tikv/client-go
Length of output: 50370
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "pessimistic.go around send/request handling:"
sed -n '120,205p' txnkv/transaction/pessimistic.go
sed -n '630,720p' txnkv/transaction/pessimistic.go
echo
echo "txn.go around lockKeysWithSharedLockUpgrade callers/input validation:"
sed -n '1750,1805p' txnkv/transaction/txn.go
sed -n '1620,1655p' txnkv/transaction/txn.go
sed -n '1935,1960p' txnkv/transaction/txn.go
echo
echo "2pc.go around shared lock precommit/shared lock validation:"
sed -n '690,735p' txnkv/transaction/2pc.goRepository: tikv/client-go
Length of output: 11166
Reject selecting an upgrade/key-only shared lock as the pessimistic primary.
selectPrimaryForPessimisticLock picks sortedKeys[0]; when normalExclusiveKeys is empty, upgradeKeys[0] is still a shared Op_SharedPessimisticLock request. The ordinary pessimistic fast path also rejects shared-mode pessimistic locking without an existing exclusive primary. Guard this fallback by requiring len(normalExclusiveKeys) > 0, or return an explicit error when the only candidate is shared-locked before sending the upgrade RPC.
🤖 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 `@txnkv/transaction/txn.go` around lines 1565 - 1575, Update the
primary-selection block around selectPrimaryForPessimisticLock so it never uses
upgradeKeys as a fallback candidate when normalExclusiveKeys is empty. Require
at least one normal exclusive key before selecting the pessimistic primary, or
explicitly return the established error for a shared-only candidate before
issuing the upgrade RPC.
| if len(upgradeKeys) > 0 { | ||
| if len(keys) > 0 { | ||
| keys = deduplicateKeys(keys) | ||
| } | ||
| upgradeKeys = deduplicateKeys(upgradeKeys) | ||
| return txn.lockKeysWithSharedLockUpgrade(ctx, lockCtx, keys, upgradeKeys) | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find ForUpdateTS handling and zero-value paths around pessimistic locking.
set -euo pipefail
rg -n -C 4 'ForUpdateTS\s*[><=]=?\s*0' --type=go
rg -n -C 3 'NewLockCtx\(' --type=go -g '!**/*_test.go'Repository: tikv/client-go
Length of output: 152
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Repository files matching txn.go:"
fd -a 'txn\.go$' . | sed 's#^\./##'
echo
echo "Locate lockKeysWithSharedLockUpgrade and nearby lines:"
for f in $(fd 'txn\.go$' .); do
if rg -q 'lockKeysWithSharedLockUpgrade|func .*lockKeys\(' "$f"; then
echo "FILE: $f"
wc -l "$f"
ast-grep outline "$f" --match 'lockKeysWithSharedLockUpgrade' --view expanded || true
ast-grep outline "$f" --match 'lockKeys' --view expanded || true
fi
doneRepository: tikv/client-go
Length of output: 754
🏁 Script executed:
#!/bin/bash
set -euo pipefail
for f in $(fd 'txn\.go$' .); do
echo "=== $f ==="
sed -n '1660,1965p' "$f" | cat -n -v | sed 's/^/LINE /'
doneRepository: tikv/client-go
Length of output: 15121
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "txn.go outline around NewLockCtx and lockCtx fields:"
rg -n -C 3 'NewLockCtx|type LockCtx|ForUpdateTS|AllowSharedLockUpgrade|LockOnlyIfExists|InShareMode' --type=go -g '!**/*_test.go'
echo
echo "txn.go lines 1500-1760:"
sed -n '1500,1760p' txnkv/transaction/txn.go | cat -n -v
echo
echo "search lockKeys calls:"
rg -n -C 4 '\.lockKeys\(|lockKeys\(' --type=go -g '!**/*_test.go'Repository: tikv/client-go
Length of output: 205
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find LockCtx-related definitions/usages:"
rg -n -C 4 'NewLockCtx|type LockCtx|ForUpdateTS|AllowSharedLockUpgrade|LockOnlyIfExists|InShareMode|lockKeys\(' txnkv examples -- '*.go'
echo
echo "Inspect lockKeysWithSharedLockUpgrade implementation:"
sed -n '1530,1610p' txnkv/transaction/txn.go | cat -n -vRepository: tikv/client-go
Length of output: 50370
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Inspect lockKeysWithSharedLockUpgrade implementation:"
sed -n '1544,1615p' txnkv/transaction/txn.go | cat -n -v
echo
echo "Find production-call-site paths with LockOnlyIfExists or shared locks around LockKeys:"
rg -n -C 5 'LockOnlyIfExists|AllowSharedLockUpgrade|LockKeys\(|LockKeysFunc\(' txnkv examples -- '*.go' ':!*_test.go'
echo
echo "Search public/client LockKeys helpers that set LockOnlyIfExists:"
rg -n -C 5 'LockOnlyIfExists|AllowSharedLockUpgrade|NewLockCtx' tikv/txnkv txnkv -- '*.go' ':!*_test.go'Repository: tikv/client-go
Length of output: 24503
Guard shared-lock upgrades with ForUpdateTS > 0.
lockKeysWithSharedLockUpgrade always creates/uses the committer and sends pessimistic-lock RPCs. Keep the existing branch at Line 1790 for normal pessimistic locks, then branch on txn.committer != nil && lockCtx.ForUpdateTS > 0 before sending upgrade RPCs. If ForUpdateTS == 0, only update the local memory-buffer upgrade flag, matching the normal fallback path.
🤖 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 `@txnkv/transaction/txn.go` around lines 1778 - 1784, Guard the shared-lock
upgrade path in the surrounding transaction lock flow with txn.committer != nil
&& lockCtx.ForUpdateTS > 0 before calling lockKeysWithSharedLockUpgrade.
Preserve the existing normal pessimistic-lock branch, and when ForUpdateTS is
zero update only the local memory-buffer upgrade flag instead of sending upgrade
RPCs.
Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Add a typed ErrLockUpgradeConflict, decode it from KeyError including API v2 key decoding, treat it as a known shared-lock upgrade failure, and keep the local go.work wiring committed so Task 5 remains buildable against the local kvproto update during rollout sequencing. Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Rename the shared-lock-upgrade helper to describe the undetermined-result case directly and invert its return value so the upgrade error branch reads in positive form without changing behavior. Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
For non-upgrade pessimistic lock groups, the rollback path always uses the same keys that were locked. The shared-lock-upgrade path returns before any pessimistic rollback, so the extra parameter was unused there and can be removed without changing behavior. Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Move the temporary lock key into the LockOnlyIfExists validation branch and rename it to lockKey so the variable name matches its diagnostic-only role. Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Replace the local go.work wiring with the pushed kvproto fork so this branch can build without host-specific workspace paths. Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Look up the returned lock value with the converted key directly so staticcheck can avoid retaining an unnecessary string conversion variable. Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
Mirror the root module's kvproto version and replacement in the independent integration module because dependency replace directives are not inherited. Regenerate its checksums with Go 1.25.12 so integration jobs can resolve the shared-lock RPC definitions. Signed-off-by: Wenxuan Zhang <wenxuangm@gmail.com>
afb7539 to
93081a8
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@integration_tests/go.mod`:
- Line 10: Scope the kvproto replacement in the integration test module to the
specific newer version requiring github.com/wfxr/kvproto, rather than replacing
every kvproto version. Preserve the older transitive version selected by
tikb/pd/client unless compatibility with the replacement is verified.
In `@txnkv/transaction/txn.go`:
- Around line 802-807: Move the fatalTxnErr check in Commit below the defer
txn.close() registration so returning fatalTxnErr still closes the transaction
and releases committer resources. Preserve the existing undetermined-error
handling and fatal-error return behavior.
🪄 Autofix
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 Plus
Run ID: a082e299-cebb-42f8-88bd-859e82ce6b24
⛔ Files ignored due to path filters (2)
go.sumis excluded by!**/*.sumintegration_tests/go.sumis excluded by!**/*.sum
📒 Files selected for processing (12)
error/error.goerror/error_test.gogo.modintegration_tests/go.modinternal/apicodec/codec_v2.gointernal/apicodec/codec_v2_test.gokv/kv.gotxnkv/transaction/2pc.gotxnkv/transaction/txn.gotxnkv/transaction/txn_test.goutil/redact/redact.goutil/redact/redact_test.go
🚧 Files skipped from review as they are similar to previous changes (10)
- kv/kv.go
- go.mod
- txnkv/transaction/2pc.go
- error/error_test.go
- internal/apicodec/codec_v2.go
- internal/apicodec/codec_v2_test.go
- util/redact/redact.go
- util/redact/redact_test.go
- txnkv/transaction/txn_test.go
- error/error.go
| github.com/pingcap/errors v0.11.5-0.20260508054701-306e305bcf41 | ||
| github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 | ||
| github.com/pingcap/kvproto v0.0.0-20260721064811-683dad8fa368 | ||
| github.com/pingcap/kvproto v0.0.0-20260806092442-d04fa0402753 |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
cd integration_tests
go list -m -f '{{.Path}}@{{.Version}}{{if .Replace}} => {{.Replace.Path}}@{{.Replace.Version}}{{end}}' \
github.com/pingcap/kvproto
go mod graph |
awk '$2 ~ /^github.com\/pingcap\/kvproto@/ { print $2 }' |
sort -u
go mod tidy -diff
go mod verifyRepository: tikv/client-go
Length of output: 144
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
echo "== files =="
git ls-files | sed -n '1,120p'
echo "== go.mod snippets =="
if [ -f integration_tests/go.mod ]; then
nl -ba integration_tests/go.mod | sed -n '1,40p;160,190p'
fi
echo "== go.sum kvproto entries =="
if [ -f integration_tests/go.sum ]; then
rg -n 'github.com/(pingcap|wfxr)/kvproto|github.com/tikv/client-go/v2' integration_tests/go.sum || true
fi
echo "== imports/usages =="
rg -n 'pingcap/kvproto|wfxr/kvproto|ErrShared|LockShared|shared.lock|shared lock' -g '*.go' -g 'go.mod' -g 'go.sum' integration_tests || trueRepository: tikv/client-go
Length of output: 3794
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
echo "== integration_tests/go.mod lines =="
sed -n '1,35p;150,190p' integration_tests/go.mod
echo "== integration_tests/go.sum kvproto references =="
awk '
/github\.com\/(pingcap|wfxr)\/kvproto/ || /github\.com\/tikv\/client-go/v2/ { print }
' integration_tests/go.sum | sed -n '1,120p'
echo "== shared-lock error schema definitions =="
rg -n 'ErrShared|Shared|shared_lock|shared lock|lock_shared|LockShared' integration_tests shared_lock_test lock_test internal error --glob '*.go' --glob 'go.mod' --glob 'go.sum' || true
echo "== go command availability =="
command -v go || true
go version || trueRepository: tikv/client-go
Length of output: 3462
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
echo "== go.sum kvproto references =="
grep -n 'github.com/\(pingcap\|wfxr\)/kvproto\|github.com/tikv/client-go/v2' integration_tests/go.sum || true
echo "== shared_lock references =="
grep -RInE 'ErrShared|Shared|shared_lock|shared lock|lock_shared|LockShared' integration_tests --include='*.go' --include='go.mod' --include='go.sum' || true
echo "== go command availability =="
command -v go || true
go version || trueRepository: tikv/client-go
Length of output: 7333
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
cd integration_tests
echo "== required module line =="
grep -n 'require github.com/pingcap/kvproto\|github.com/pingcap/kvproto' go.mod
echo "== replace directive line =="
grep -n '^replace github.com/pingcap/kvproto' go.mod
echo "== direct module info =="
go list -m -json github.com/pingcap/kvproto | sed -n '1,80p'
echo "== module graph kvproto lines =="
go mod graph | grep 'github.com/pingcap/kvproto@' | sort -u || true
echo "== tidy diff status =="
if ! go mod tidy -diff; then
true
fi
echo "== module verify status =="
if ! go mod verify; then
true
fiRepository: tikv/client-go
Length of output: 368
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
cd integration_tests
echo "== module info =="
go list -m -json github.com/pingcap/kvproto
echo "== module graph kvproto lines =="
go mod graph | grep 'github.com/pingcap/kvproto@' | sort -u || true
echo "== mod tidy diff exit status =="
go mod tidy; echo "exit=$?"
echo "== mod verify exit status =="
go mod verify; echo "exit=$?"Repository: tikv/client-go
Length of output: 22249
Limit the kvproto replacement to the versions that need it.
The unversioned replace directive also substitutes the older github.com/pingcap/kvproto@v0.0.0-20260622063236-b41e86365ce0 selected by tikb/pd/client, while the direct module target is the newer commit. Scope the replacement to the single version that requires github.com/wfxr/kvproto unless the older transitive version is also compatible.
🤖 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 `@integration_tests/go.mod` at line 10, Scope the kvproto replacement in the
integration test module to the specific newer version requiring
github.com/wfxr/kvproto, rather than replacing every kvproto version. Preserve
the older transitive version selected by tikb/pd/client unless compatibility
with the replacement is verified.
| if txn.committer != nil { | ||
| undeterminedErr, fatalTxnErr := txn.committer.getTxnStateErrs() | ||
| if undeterminedErr == nil && fatalTxnErr != nil { | ||
| return fatalTxnErr | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
The fatal-error return skips txn.close().
defer txn.close() is registered at Line 808, after this block. When fatalTxnErr != nil, Commit returns before that defer is registered. The transaction stays valid, and the committer resources, including the TTL manager, are not released. The undetermined path at Line 859 returns after the defer, so it does close. Move the fatal check below Line 808 for consistent cleanup.
🐛 Proposed fix
- if txn.committer != nil {
- undeterminedErr, fatalTxnErr := txn.committer.getTxnStateErrs()
- if undeterminedErr == nil && fatalTxnErr != nil {
- return fatalTxnErr
- }
- }
defer txn.close()
+
+ if txn.committer != nil {
+ undeterminedErr, fatalTxnErr := txn.committer.getTxnStateErrs()
+ if undeterminedErr == nil && fatalTxnErr != nil {
+ return fatalTxnErr
+ }
+ }📝 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.
| if txn.committer != nil { | |
| undeterminedErr, fatalTxnErr := txn.committer.getTxnStateErrs() | |
| if undeterminedErr == nil && fatalTxnErr != nil { | |
| return fatalTxnErr | |
| } | |
| } | |
| defer txn.close() | |
| if txn.committer != nil { | |
| undeterminedErr, fatalTxnErr := txn.committer.getTxnStateErrs() | |
| if undeterminedErr == nil && fatalTxnErr != nil { | |
| return fatalTxnErr | |
| } | |
| } |
🤖 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 `@txnkv/transaction/txn.go` around lines 802 - 807, Move the fatalTxnErr check
in Commit below the defer txn.close() registration so returning fatalTxnErr
still closes the transaction and releases committer resources. Preserve the
existing undetermined-error handling and fatal-error return behavior.
1 participant
Summary
Support upgrading a pessimistic shared lock to an exclusive lock when the caller explicitly opts in through
LockCtx.AllowSharedLockUpgrade.This is needed for TiDB foreign-key checking with
tidb_foreign_key_check_in_shared_lockenabled. In that flow, a transaction may first acquire a shared lock while checking the parent row and then update the same parent row later in the same transaction.Related to pingcap/tidb#68815.
Changes
This PR adds the client-side upgrade path and preserves transaction correctness across deterministic, fatal, and undetermined upgrade failures.
LockCtx.AllowSharedLockUpgradeto gate shared-to-exclusive lock upgrades explicitly.ErrLockUpgradeConflictandErrSharedLockLosthandling, including API v2 key decoding and error-key redaction.SharedLockLost, and mark the result undetermined when an upgrade failure leaves the remote lock state uncertain.Dependency and TODO
This draft depends on the protocol changes in pingcap/kvproto#1495. Until that PR merges,
go.modtemporarily replacesgithub.com/pingcap/kvprotowithgithub.com/wfxr/kvprotoat commitd04fa040275306661d8123933445bd72f7feaf7f. Go does not propagate dependency-modulereplacedirectives, so downstream projects that test this client-go branch must add the same temporary replacement in their main module.go.work/go.work.sumsetup with the pinned remote fork.github.com/pingcap/kvprototo the merged upstream commit, remove the temporarygithub.com/wfxr/kvprotoreplacement, rungo mod tidy, and rerun the full test suite.Testing
The branch now resolves kvproto without a local workspace, and the following checks pass:
GOWORK=off go mod tidy GOWORK=off go mod verify GOWORK=off go test --tags=intest ./... git diff --checkSummary by CodeRabbit
New Features
Bug Fixes
Tests