From 0b1529a757c17a2785a0d2f7ebebc17186522bbe Mon Sep 17 00:00:00 2001 From: Wenxuan Zhang Date: Thu, 21 May 2026 17:39:20 +0800 Subject: [PATCH 1/7] txnkv/transaction, kv: support shared lock upgrade keys --- kv/kv.go | 1 + txnkv/transaction/txn.go | 232 ++++++++++++++++++++++++++++++++-- txnkv/transaction/txn_test.go | 197 ++++++++++++++++++++++++++++- 3 files changed, 420 insertions(+), 10 deletions(-) diff --git a/kv/kv.go b/kv/kv.go index 4bd086183e..0d07b75a7a 100644 --- a/kv/kv.go +++ b/kv/kv.go @@ -67,6 +67,7 @@ type LockCtx struct { CheckExistence bool LockOnlyIfExists bool InShareMode bool + AllowSharedLockUpgrade bool Values map[string]ReturnedValue MaxLockedWithConflictTS uint64 ValuesLock sync.Mutex diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index 52bb657fba..3982ae6214 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -850,6 +850,9 @@ func (txn *KVTxn) Commit(ctx context.Context) error { } txn.committer = committer } + if err := txn.getUndeterminedLockStateErr(); err != nil { + return err + } committer.SetDiskFullOpt(txn.diskFullOpt) committer.SetTxnSource(txn.txnSource) @@ -1384,6 +1387,199 @@ func (txn *KVTxn) LockKeysFunc(ctx context.Context, lockCtx *tikv.LockCtx, fn fu return txn.lockKeys(ctx, lockCtx, fn, keysInput...) } +func (txn *KVTxn) getUndeterminedLockStateErr() error { + if txn.committer != nil && txn.committer.getUndeterminedErr() != nil { + return errors.WithStack(tikverr.ErrResultUndetermined) + } + return nil +} + +func isKnownSharedLockUpgradeFailure(err error) bool { + if err == nil { + return false + } + if tikverr.IsErrWriteConflict(err) || tikverr.IsErrKeyExist(err) || + errors.Is(err, tikverr.ErrLockAcquireFailAndNoWaitSet) || + errors.Is(err, tikverr.ErrLockWaitTimeout) { + return true + } + var deadlock *tikverr.ErrDeadlock + if errors.As(err, &deadlock) { + return true + } + var assertionFailed *tikverr.ErrAssertionFailed + return errors.As(err, &assertionFailed) +} + +func (txn *KVTxn) lockPessimisticKeyGroup( + ctx context.Context, + lockCtx *tikv.LockCtx, + keys [][]byte, + rollbackKeys [][]byte, + isUpgrade bool, +) (int, error) { + bo := retry.NewBackofferWithVars(ctx, pessimisticLockMaxBackoff, txn.vars) + txn.committer.isFirstLock = txn.lockedCnt == 0 && len(keys) == 1 + err := txn.committer.pessimisticLockMutations(bo, lockCtx, kvrpcpb.PessimisticLockWakeUpMode_WakeUpModeNormal, &PlainMutations{keys: keys}) + if lockCtx.Stats != nil && bo.GetTotalSleep() > 0 { + atomic.AddInt64(&lockCtx.Stats.BackoffTime, int64(bo.GetTotalSleep())*int64(time.Millisecond)) + lockCtx.Stats.Mu.Lock() + lockCtx.Stats.Mu.BackoffTypes = append(lockCtx.Stats.Mu.BackoffTypes, bo.GetTypes()...) + lockCtx.Stats.Mu.Unlock() + } + if err != nil { + var unmarkKeys [][]byte + memBuf := txn.us.GetMemBuffer() + memBuf.RLock() + for _, key := range keys { + if txn.us.HasPresumeKeyNotExists(key) { + unmarkKeys = append(unmarkKeys, key) + } + } + memBuf.RUnlock() + for _, key := range unmarkKeys { + txn.us.UnmarkPresumeKeyNotExists(key) + } + + if isUpgrade { + if !isKnownSharedLockUpgradeFailure(err) { + txn.committer.setUndeterminedErr(err) + return 0, errors.WithStack(tikverr.ErrResultUndetermined) + } + return 0, err + } + + keyMayBeLocked := !tikverr.IsErrWriteConflict(err) && !tikverr.IsErrKeyExist(err) + if len(keys) > 1 || keyMayBeLocked { + dl, isDeadlock := errors.Cause(err).(*tikverr.ErrDeadlock) + if isDeadlock { + if hashInKeys(dl.DeadlockKeyHash, keys) { + dl.IsRetryable = true + } + if lockCtx.OnDeadlock != nil { + lockCtx.OnDeadlock(dl) + } + } + + rollbackForUpdateTS := lockCtx.ForUpdateTS + if lockCtx.MaxLockedWithConflictTS > rollbackForUpdateTS { + rollbackForUpdateTS = lockCtx.MaxLockedWithConflictTS + } + wg := txn.asyncPessimisticRollback(ctx, rollbackKeys, rollbackForUpdateTS) + + if isDeadlock { + logutil.Logger(ctx).Debug("deadlock error received", zap.Uint64("startTS", txn.startTS), zap.Stringer("deadlockInfo", dl)) + if dl.IsRetryable { + wg.Wait() + time.Sleep(time.Millisecond * 5) + if _, err := util.EvalFailpoint("SingleStmtDeadLockRetrySleep"); err == nil { + time.Sleep(300 * time.Millisecond) + } + } + } + } + return 0, err + } + + checkedExistence := lockCtx.CheckExistence + skippedLockKeys := 0 + memBuf := txn.us.GetMemBuffer() + for _, key := range keys { + valExists := true + keyStr := string(key) + if val, ok := lockCtx.Values[keyStr]; ok { + if lockCtx.ReturnValues || checkedExistence || val.LockedWithConflictTS != 0 { + if !val.Exists { + valExists = false + } + } + } + + if lockCtx.LockOnlyIfExists && !valExists { + skippedLockKeys++ + continue + } + + setValExists := tikv.SetKeyLockedValueExists + if !valExists { + setValExists = tikv.SetKeyLockedValueNotExists + } + memBuf.UpdateFlags(key, tikv.SetKeyLocked, tikv.DelNeedCheckExists, setValExists, tikv.SetKeyLockedInExclusiveMode) + } + if !isUpgrade { + txn.lockedCnt += len(keys) - skippedLockKeys + } + return skippedLockKeys, nil +} + +func (txn *KVTxn) lockKeysWithSharedLockUpgrade( + ctx context.Context, + lockCtx *tikv.LockCtx, + normalExclusiveKeys [][]byte, + upgradeKeys [][]byte, +) error { + if txn.committer == nil { + var sessionID uint64 + val := ctx.Value(util.SessionID) + if val != nil { + sessionID = val.(uint64) + } + var err error + txn.committer, err = newTwoPhaseCommitter(txn, sessionID) + if err != nil { + return err + } + } + + assignedPrimaryKey := false + totalKeys := len(normalExclusiveKeys) + len(upgradeKeys) + if txn.committer.primaryKey == nil { + assignedPrimaryKey = true + keysForPrimary := normalExclusiveKeys + if len(keysForPrimary) == 0 { + keysForPrimary = upgradeKeys + } + txn.selectPrimaryForPessimisticLock(keysForPrimary) + } + + txn.committer.forUpdateTS = lockCtx.ForUpdateTS + lockCtx.Stats = &util.LockKeysDetails{ + LockKeys: int32(totalKeys), + ResolveLock: util.ResolveLockDetail{}, + } + + lockedInThisCall := 0 + if len(normalExclusiveKeys) > 0 { + skipped, err := txn.lockPessimisticKeyGroup(ctx, lockCtx, normalExclusiveKeys, normalExclusiveKeys, false) + if err != nil { + if assignedPrimaryKey && lockedInThisCall == 0 { + txn.resetPrimary(false) + } + return err + } + lockedInThisCall += len(normalExclusiveKeys) - skipped + } + + for _, key := range upgradeKeys { + skipped, err := txn.lockPessimisticKeyGroup(ctx, lockCtx, [][]byte{key}, nil, true) + if err != nil { + if assignedPrimaryKey && lockedInThisCall == 0 { + txn.resetPrimary(false) + } + return err + } + lockedInThisCall += 1 - skipped + } + + if assignedPrimaryKey && lockCtx.LockOnlyIfExists { + if totalKeys != 1 { + panic("LockOnlyIfExists only assigns the primary key when locking only one key") + } + txn.unsetPrimaryKeyIfNeeded(lockCtx) + } + return nil +} + func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func(), keysInput ...[]byte) error { if txn.interceptor != nil { // User has called txn.SetRPCInterceptor() to explicitly set an interceptor, we @@ -1404,6 +1600,9 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() if err != nil { return err } + if err := txn.getUndeterminedLockStateErr(); err != nil { + return err + } defer func() { if lockCtx.InShareMode { @@ -1456,6 +1655,7 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() memBuf := txn.us.GetMemBuffer() // Avoid data race with concurrent updates to the memBuf memBuf.RLock() + upgradeKeys := make([][]byte, 0, len(keysInput)) for _, key := range keysInput { // The value of lockedMap is only used by pessimistic transactions. var valueExist, locked, lockedInShareMode, checkKeyExists bool @@ -1466,9 +1666,12 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() checkKeyExists = flags.HasNeedCheckExists() } - if lockedInShareMode && !lockCtx.InShareMode { - memBuf.RUnlock() - return errors.New("upgrading a shared lock to an exclusive lock is not supported") + upgradeCandidate := lockedInShareMode && !lockCtx.InShareMode + if upgradeCandidate { + if !txn.IsPessimistic() || !lockCtx.AllowSharedLockUpgrade { + memBuf.RUnlock() + return errors.New("upgrading a shared lock to an exclusive lock is not supported") + } } // If the key is locked in the current aggressive locking stage, override the information in memBuf. @@ -1484,7 +1687,9 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() } } - if !locked || isInLastAggressiveLockingStage { + if upgradeCandidate { + upgradeKeys = append(upgradeKeys, key) + } else if !locked || isInLastAggressiveLockingStage { // Locks acquired in the previous aggressive locking stage might need to be updated later in // `filterAggressiveLockedKeys`. keys = append(keys, key) @@ -1497,7 +1702,7 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() return txn.committer.extractKeyExistsErr(e) } } - if lockCtx.ReturnValues && locked { + if lockCtx.ReturnValues && locked && !upgradeCandidate { keyStr := string(key) // An already locked key can not return values, we add an entry to let the caller get the value // in other ways. @@ -1506,7 +1711,7 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() } memBuf.RUnlock() - if len(keys) == 0 { + if len(keys) == 0 && len(upgradeKeys) == 0 { return nil } if lockCtx.LockOnlyIfExists { @@ -1520,14 +1725,25 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() // It can't transform LockOnlyIfExists mode to normal mode. If so, it can add a lock to a key // which doesn't exist in tikv. TiDB should ensure that primary key must be set when it sends // a LockOnlyIfExists pessimistic lock request. - if (txn.committer == nil || txn.committer.primaryKey == nil) && len(keys) > 1 { + if (txn.committer == nil || txn.committer.primaryKey == nil) && len(keys)+len(upgradeKeys) > 1 { + lockKey := keys[0] + if len(keys) == 0 { + lockKey = upgradeKeys[0] + } return &tikverr.ErrLockOnlyIfExistsNoPrimaryKey{ StartTS: txn.startTS, ForUpdateTs: lockCtx.ForUpdateTS, - LockKey: keys[0], + LockKey: lockKey, } } } + if len(upgradeKeys) > 0 { + if len(keys) > 0 { + keys = deduplicateKeys(keys) + } + upgradeKeys = deduplicateKeys(upgradeKeys) + return txn.lockKeysWithSharedLockUpgrade(ctx, lockCtx, keys, upgradeKeys) + } keys = deduplicateKeys(keys) checkedExistence := false filteredAggressiveLockedKeysCount := 0 diff --git a/txnkv/transaction/txn_test.go b/txnkv/transaction/txn_test.go index b6d3b998ca..f56f31d113 100644 --- a/txnkv/transaction/txn_test.go +++ b/txnkv/transaction/txn_test.go @@ -23,6 +23,7 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/pkg/errors" "github.com/stretchr/testify/require" + tikverr "github.com/tikv/client-go/v2/error" "github.com/tikv/client-go/v2/internal/client" "github.com/tikv/client-go/v2/internal/locate" "github.com/tikv/client-go/v2/kv" @@ -273,8 +274,12 @@ func TestLockKeys(t *testing.T) { // `lockKeys` in exclusive mode again on k1 to upgrade the lock. lockCtx.InShareMode = false expectedLockType = kvrpcpb.Op_PessimisticLock - err = txn.lockKeys(context.TODO(), lockCtx, nil, key1) - require.ErrorContains(t, err, "upgrading a shared lock to an exclusive lock is not supported") + lockCtx.AllowSharedLockUpgrade = true + require.NoError(t, txn.lockKeys(context.TODO(), lockCtx, nil, key1)) + flags, err = txn.GetMemBuffer().GetFlags(key1) + require.NoError(t, err) + require.True(t, flags.HasLocked()) + require.False(t, flags.HasLockedInShareMode()) done := make(chan struct{}) go func() { @@ -289,6 +294,194 @@ func TestLockKeys(t *testing.T) { }) } +func TestSharedLockUpgrade(t *testing.T) { + type requestSummary struct { + keys [][]byte + op kvrpcpb.Op + } + + newRecorderTxn := func(t *testing.T, onLock func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error)) (*testTxn, *[]requestSummary) { + txn := newTestTxn(t, 1) + txn.SetPessimistic(true) + requests := make([]requestSummary, 0, 8) + txn.store.client.onSend = func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) { + require.Equal(t, tikvrpc.CmdPessimisticLock, req.Type) + 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{ + 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]) + } + return onLock(len(requests)-1, lockReq) + } + return txn, &requests + } + + lockSharedKey := func(t *testing.T, txn *testTxn, primaryKey, upgradeKey []byte) { + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + require.NoError(t, txn.lockKeys(context.TODO(), lockCtx, nil, primaryKey)) + lockCtx.InShareMode = true + require.NoError(t, txn.lockKeys(context.TODO(), lockCtx, nil, upgradeKey)) + flags, err := txn.GetMemBuffer().GetFlags(upgradeKey) + require.NoError(t, err) + require.True(t, flags.HasLocked()) + require.True(t, flags.HasLockedInShareMode()) + } + + keysAsStrings := func(keys [][]byte) []string { + ret := make([]string, len(keys)) + for i, key := range keys { + ret[i] = string(key) + } + return ret + } + + t.Run("GateOffRejectsLocally", func(t *testing.T) { + primaryKey := []byte("primary-key") + upgradeKey := []byte("upgrade-key") + txn, requests := newRecorderTxn(t, func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error) { + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{}}, nil + }) + lockSharedKey(t, txn, primaryKey, upgradeKey) + + requestCount := len(*requests) + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + err := txn.lockKeys(context.TODO(), lockCtx, nil, upgradeKey) + require.ErrorContains(t, err, "upgrading a shared lock to an exclusive lock is not supported") + require.Len(t, *requests, requestCount) + + flags, err := txn.GetMemBuffer().GetFlags(upgradeKey) + require.NoError(t, err) + require.True(t, flags.HasLockedInShareMode()) + }) + + t.Run("GateOnSendsUpgradeSeparatelyAndPromotesLocalFlags", func(t *testing.T) { + primaryKey := []byte("primary-key") + upgradeKey := []byte("upgrade-key") + normalKey1 := []byte("normal-key-1") + normalKey2 := []byte("normal-key-2") + txn, requests := newRecorderTxn(t, func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error) { + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{}}, nil + }) + lockSharedKey(t, txn, primaryKey, upgradeKey) + *requests = (*requests)[:0] + + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + lockCtx.AllowSharedLockUpgrade = true + require.NoError(t, txn.lockKeys(context.TODO(), lockCtx, nil, normalKey1, upgradeKey, normalKey2)) + + require.Len(t, *requests, 2) + require.Equal(t, kvrpcpb.Op_PessimisticLock, (*requests)[0].op) + require.ElementsMatch(t, []string{string(normalKey1), string(normalKey2)}, keysAsStrings((*requests)[0].keys)) + require.Equal(t, kvrpcpb.Op_PessimisticLock, (*requests)[1].op) + require.Equal(t, []string{string(upgradeKey)}, keysAsStrings((*requests)[1].keys)) + + flags, err := txn.GetMemBuffer().GetFlags(upgradeKey) + require.NoError(t, err) + require.True(t, flags.HasLocked()) + require.False(t, flags.HasLockedInShareMode()) + + committer, err := TxnProbe{KVTxn: txn.KVTxn}.NewCommitter(1) + require.NoError(t, err) + mutations := committer.MutationsOfKeys([][]byte{upgradeKey}) + prewriteReq := committer.BuildPrewriteRequest(1, 1, 1, mutations, 1).Req.(*kvrpcpb.PrewriteRequest) + require.Len(t, prewriteReq.Mutations, 1) + require.Equal(t, kvrpcpb.Op_Lock, prewriteReq.Mutations[0].Op) + }) + + t.Run("ExplicitUpgradeFailureKeepsExistingSharedHolderAndEarlierExclusiveLocks", func(t *testing.T) { + primaryKey := []byte("primary-key") + upgradeKey := []byte("upgrade-key") + normalKey := []byte("normal-key") + txn, requests := newRecorderTxn(t, func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error) { + if len(req.Mutations) == 1 && + req.Mutations[0].Op == kvrpcpb.Op_PessimisticLock && + string(req.Mutations[0].Key) == string(upgradeKey) { + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{ + Errors: []*kvrpcpb.KeyError{{ + Conflict: &kvrpcpb.WriteConflict{ + StartTs: 1, + ConflictTs: 2, + ConflictCommitTs: 3, + Key: upgradeKey, + Reason: kvrpcpb.WriteConflict_PessimisticRetry, + }, + }}, + }}, nil + } + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{}}, nil + }) + lockSharedKey(t, txn, primaryKey, upgradeKey) + *requests = (*requests)[:0] + + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + lockCtx.AllowSharedLockUpgrade = true + err := txn.lockKeys(context.TODO(), lockCtx, nil, normalKey, upgradeKey) + require.True(t, tikverr.IsErrWriteConflict(err)) + require.Len(t, *requests, 2) + require.Equal(t, []string{string(normalKey)}, keysAsStrings((*requests)[0].keys)) + require.Equal(t, []string{string(upgradeKey)}, keysAsStrings((*requests)[1].keys)) + + upgradeFlags, err := txn.GetMemBuffer().GetFlags(upgradeKey) + require.NoError(t, err) + require.True(t, upgradeFlags.HasLocked()) + require.True(t, upgradeFlags.HasLockedInShareMode()) + + normalFlags, err := txn.GetMemBuffer().GetFlags(normalKey) + require.NoError(t, err) + require.True(t, normalFlags.HasLocked()) + require.False(t, normalFlags.HasLockedInShareMode()) + + require.ElementsMatch(t, + []string{string(primaryKey), string(upgradeKey), string(normalKey)}, + keysAsStrings(TxnProbe{KVTxn: txn.KVTxn}.CollectLockedKeys())) + }) + + t.Run("OutcomeUnknownUpgradeFailureIsTransactionFatal", func(t *testing.T) { + primaryKey := []byte("primary-key") + upgradeKey := []byte("upgrade-key") + txn, requests := newRecorderTxn(t, func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error) { + if len(req.Mutations) == 1 && + req.Mutations[0].Op == kvrpcpb.Op_PessimisticLock && + string(req.Mutations[0].Key) == string(upgradeKey) { + return &tikvrpc.Response{}, nil + } + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{}}, nil + }) + lockSharedKey(t, txn, primaryKey, upgradeKey) + *requests = (*requests)[:0] + + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + lockCtx.AllowSharedLockUpgrade = true + err := txn.lockKeys(context.TODO(), lockCtx, nil, upgradeKey) + require.Error(t, err) + require.True(t, tikverr.IsErrorUndetermined(err)) + require.Len(t, *requests, 1) + require.Equal(t, []string{string(upgradeKey)}, keysAsStrings((*requests)[0].keys)) + + flags, err := txn.GetMemBuffer().GetFlags(upgradeKey) + require.NoError(t, err) + require.True(t, flags.HasLocked()) + require.True(t, flags.HasLockedInShareMode()) + + err = txn.LockKeys(context.TODO(), kv.NewLockCtx(2, kv.LockNoWait, time.Now()), []byte("later-key")) + require.Error(t, err) + require.True(t, tikverr.IsErrorUndetermined(err)) + + err = txn.Commit(context.TODO()) + require.Error(t, err) + require.True(t, tikverr.IsErrorUndetermined(err)) + }) +} + func TestSharedLockCommitterIncompatibilities(t *testing.T) { t.Run("RejectSharedLockPrimaryKey", func(t *testing.T) { key := []byte("shared-key") From 4da405178ce81c3b290375ed70f7456a06011fe0 Mon Sep 17 00:00:00 2001 From: Wenxuan Zhang Date: Thu, 21 May 2026 17:46:21 +0800 Subject: [PATCH 2/7] txnkv/transaction: reject upgrade in aggressive locking --- txnkv/transaction/txn.go | 18 +++++++++++------ txnkv/transaction/txn_test.go | 37 +++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index 3982ae6214..8adc8c7637 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -1672,6 +1672,10 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() memBuf.RUnlock() return errors.New("upgrading a shared lock to an exclusive lock is not supported") } + if txn.IsInAggressiveLockingMode() { + memBuf.RUnlock() + return errors.New("shared lock upgrade is not supported in aggressive/fair locking mode") + } } // If the key is locked in the current aggressive locking stage, override the information in memBuf. @@ -1714,26 +1718,28 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() if len(keys) == 0 && len(upgradeKeys) == 0 { return nil } + var lockOnlyIfExistsKey []byte + if len(keys) > 0 { + lockOnlyIfExistsKey = keys[0] + } else { + lockOnlyIfExistsKey = upgradeKeys[0] + } if lockCtx.LockOnlyIfExists { if !lockCtx.ReturnValues { return &tikverr.ErrLockOnlyIfExistsNoReturnValue{ StartTS: txn.startTS, ForUpdateTs: lockCtx.ForUpdateTS, - LockKey: keys[0], + LockKey: lockOnlyIfExistsKey, } } // It can't transform LockOnlyIfExists mode to normal mode. If so, it can add a lock to a key // which doesn't exist in tikv. TiDB should ensure that primary key must be set when it sends // a LockOnlyIfExists pessimistic lock request. if (txn.committer == nil || txn.committer.primaryKey == nil) && len(keys)+len(upgradeKeys) > 1 { - lockKey := keys[0] - if len(keys) == 0 { - lockKey = upgradeKeys[0] - } return &tikverr.ErrLockOnlyIfExistsNoPrimaryKey{ StartTS: txn.startTS, ForUpdateTs: lockCtx.ForUpdateTS, - LockKey: lockKey, + LockKey: lockOnlyIfExistsKey, } } } diff --git a/txnkv/transaction/txn_test.go b/txnkv/transaction/txn_test.go index f56f31d113..8a5aa5d4ca 100644 --- a/txnkv/transaction/txn_test.go +++ b/txnkv/transaction/txn_test.go @@ -397,6 +397,43 @@ func TestSharedLockUpgrade(t *testing.T) { require.Equal(t, kvrpcpb.Op_Lock, prewriteReq.Mutations[0].Op) }) + t.Run("RejectUpgradeInAggressiveLockingMode", func(t *testing.T) { + primaryKey := []byte("primary-key") + upgradeKey := []byte("upgrade-key") + txn, requests := newRecorderTxn(t, func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error) { + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{}}, nil + }) + lockSharedKey(t, txn, primaryKey, upgradeKey) + + requestCount := len(*requests) + txn.StartAggressiveLocking() + defer txn.CancelAggressiveLocking(context.Background()) + + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + lockCtx.AllowSharedLockUpgrade = true + err := txn.lockKeys(context.TODO(), lockCtx, nil, upgradeKey) + require.ErrorContains(t, err, "shared lock upgrade is not supported in aggressive/fair locking mode") + require.Len(t, *requests, requestCount) + require.True(t, txn.IsInAggressiveLockingMode()) + }) + + t.Run("UpgradeLockOnlyIfExistsRequiresReturnValues", func(t *testing.T) { + primaryKey := []byte("primary-key") + upgradeKey := []byte("upgrade-key") + txn, _ := newRecorderTxn(t, func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error) { + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{}}, nil + }) + lockSharedKey(t, txn, primaryKey, upgradeKey) + + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + lockCtx.AllowSharedLockUpgrade = true + lockCtx.LockOnlyIfExists = true + err := txn.lockKeys(context.TODO(), lockCtx, nil, upgradeKey) + var noReturnValueErr *tikverr.ErrLockOnlyIfExistsNoReturnValue + require.ErrorAs(t, err, &noReturnValueErr) + require.Equal(t, upgradeKey, noReturnValueErr.LockKey) + }) + t.Run("ExplicitUpgradeFailureKeepsExistingSharedHolderAndEarlierExclusiveLocks", func(t *testing.T) { primaryKey := []byte("primary-key") upgradeKey := []byte("upgrade-key") From fbb131c2a208103636e3fc28d4ff0f453fecdb84 Mon Sep 17 00:00:00 2001 From: Wenxuan Zhang Date: Fri, 22 May 2026 14:19:19 +0800 Subject: [PATCH 3/7] error, txnkv: preserve typed lock upgrade conflicts 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. --- error/error.go | 13 ++++++++ error/error_test.go | 27 +++++++++++++++ go.work | 6 ++++ go.work.sum | 53 ++++++++++++++++++++++++++++++ internal/apicodec/codec_v2.go | 6 ++++ internal/apicodec/codec_v2_test.go | 17 ++++++++++ txnkv/transaction/txn.go | 4 +++ txnkv/transaction/txn_test.go | 51 ++++++++++++++++++++++++++++ 8 files changed, 177 insertions(+) create mode 100644 go.work create mode 100644 go.work.sum diff --git a/error/error.go b/error/error.go index 792047b890..d44636d2e6 100644 --- a/error/error.go +++ b/error/error.go @@ -134,6 +134,15 @@ func (d *ErrDeadlock) Error() string { return d.String() } +// ErrLockUpgradeConflict wraps *kvrpcpb.LockUpgradeConflict to implement the error interface. +type ErrLockUpgradeConflict struct { + *kvrpcpb.LockUpgradeConflict +} + +func (e *ErrLockUpgradeConflict) Error() string { + return fmt.Sprintf("lock upgrade conflict { %s }", e.String()) +} + // PDError wraps *pdpb.Error to implement the error interface. type PDError struct { Err *pdpb.Error @@ -342,6 +351,10 @@ func ExtractKeyErr(keyErr *kvrpcpb.KeyError) error { return errors.WithStack(NewErrWriteConflict(keyErr.GetConflict())) } + if keyErr.LockUpgradeConflict != nil { + return errors.WithStack(&ErrLockUpgradeConflict{LockUpgradeConflict: keyErr.LockUpgradeConflict}) + } + if keyErr.Retryable != "" { return errors.WithStack(&ErrRetryable{Retryable: keyErr.Retryable}) } diff --git a/error/error_test.go b/error/error_test.go index 67eea837c8..0aeca92e69 100644 --- a/error/error_test.go +++ b/error/error_test.go @@ -1,13 +1,40 @@ package error import ( + stderrs "errors" "testing" "github.com/pingcap/errors" "github.com/pingcap/kvproto/pkg/kvrpcpb" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) +func TestExtractKeyErrLockUpgradeConflict(t *testing.T) { + keyErr := &kvrpcpb.KeyError{ + LockUpgradeConflict: &kvrpcpb.LockUpgradeConflict{ + Key: []byte("key"), + StartTs: 101, + OwnerStartTs: 202, + Reason: kvrpcpb.LockUpgradeConflict_SecondUpgrader, + }, + } + + err := ExtractKeyErr(keyErr) + require.Error(t, err) + require.False(t, IsErrWriteConflict(err)) + + var retryable *ErrRetryable + require.False(t, stderrs.As(err, &retryable)) + + var conflict *ErrLockUpgradeConflict + require.ErrorAs(t, err, &conflict) + require.Equal(t, []byte("key"), conflict.Key) + require.Equal(t, uint64(101), conflict.StartTs) + require.Equal(t, uint64(202), conflict.OwnerStartTs) + require.Equal(t, kvrpcpb.LockUpgradeConflict_SecondUpgrader, conflict.Reason) +} + func TestExtractDebugInfoStrFromKeyErr(t *testing.T) { origRedact := errors.RedactLogEnabled.Load() defer errors.RedactLogEnabled.Store(origRedact) diff --git a/go.work b/go.work new file mode 100644 index 0000000000..a521974339 --- /dev/null +++ b/go.work @@ -0,0 +1,6 @@ +go 1.25.9 + +use ( + . + /home/wenxuan/dev/pingcap/kvproto/.wt/feat/shared-lock-upgrade-next-gen +) diff --git a/go.work.sum b/go.work.sum new file mode 100644 index 0000000000..794366c1f8 --- /dev/null +++ b/go.work.sum @@ -0,0 +1,53 @@ +cel.dev/expr v0.24.0/go.mod h1:hLPLo1W4QUmuYdA72RBX06QTs6MXw941piREPl3Yfiw= +cloud.google.com/go/compute/metadata v0.7.0/go.mod h1:j5MvL9PprKL39t166CoB1uVHfQMs4tFQZZcKwksXUjo= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.29.0/go.mod h1:Cz6ft6Dkn3Et6l2v2a9/RpN7epQ1GtDlO6lj8bEcOvw= +github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/envoyproxy/go-control-plane v0.13.4/go.mod h1:kDfuBlDVsSj2MjrLEtRWtHlsWIFcGyB2RMO44Dc5GZA= +github.com/envoyproxy/go-control-plane/envoy v1.32.4/go.mod h1:Gzjc5k8JcJswLjAx1Zm+wSYE20UrLtt7JZMWiWQXQEw= +github.com/envoyproxy/go-control-plane/ratelimit v0.1.0/go.mod h1:Wk+tMFAFbCXaJPzVVHnPgRKdUdwW/KdbRt94AzgRee4= +github.com/envoyproxy/protoc-gen-validate v1.2.1/go.mod h1:d/C80l/jxXLdfEIhX1W2TmLfsJ31lvEjwamM4DxlWXU= +github.com/go-jose/go-jose/v4 v4.1.1/go.mod h1:BdsZGqgdO3b6tTc6LSE56wcDbMMLuPsw5d4ZD5f94kA= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/golang/glog v1.2.5/go.mod h1:6AhwSGph0fcJtXVM/PEHPqZlFeoLxhs7/t5UDAwmO+w= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/spiffe/go-spiffe/v2 v2.5.0/go.mod h1:P+NxobPc6wXhVtINNtFjNWGBTreew1GBUCwT2wPmb7g= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= +github.com/zeebo/errs v1.4.0/go.mod h1:sgbWHsvVuTPHcqJJGQ1WhI5KbWlHYz+2+2C/LSEtCw4= +go.opentelemetry.io/contrib/detectors/gcp v1.36.0/go.mod h1:IbBN8uAIIx734PTonTPxAxnjc2pQTxWNkwfstZ+6H2k= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U= +golang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos= +golang.org/x/mod v0.18.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= +golang.org/x/net v0.37.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8= +golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA= +golang.org/x/net v0.49.0/go.mod h1:/ysNB2EvaqvesRkuLAyjI1ycPZlQHM3q01F02UY/MV8= +golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= +golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= +golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/sys v0.40.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/term v0.40.0/go.mod h1:w2P8uVp06p2iyKKuvXIm7N/y0UCRt3UfJTfZ7oOpglM= +golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4= +golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA= +golang.org/x/tools v0.22.0/go.mod h1:aCwcsjqvq7Yqt6TNyX7QMU2enbQ/Gt0bo6krSeEri+c= +golang.org/x/tools v0.40.0/go.mod h1:Ik/tzLRlbscWpqqMRjyWYDisX8bG13FrdXp3o4Sr9lc= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A= +google.golang.org/grpc v1.71.0/go.mod h1:H0GRtasmQOh9LkFoCPDu3ZrwUtD1YGE+b2vYBYd/8Ec= +google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= diff --git a/internal/apicodec/codec_v2.go b/internal/apicodec/codec_v2.go index 070e75fe91..458c41166d 100644 --- a/internal/apicodec/codec_v2.go +++ b/internal/apicodec/codec_v2.go @@ -1033,6 +1033,12 @@ func (c *codecV2) decodeKeyError(keyError *kvrpcpb.KeyError) (*kvrpcpb.KeyError, } } } + if keyError.LockUpgradeConflict != nil { + keyError.LockUpgradeConflict.Key, err = c.DecodeKey(keyError.LockUpgradeConflict.Key) + if err != nil { + return nil, err + } + } if keyError.CommitTsExpired != nil { keyError.CommitTsExpired.Key, err = c.DecodeKey(keyError.CommitTsExpired.Key) if err != nil { diff --git a/internal/apicodec/codec_v2_test.go b/internal/apicodec/codec_v2_test.go index f92ba59ee0..5608e2855d 100644 --- a/internal/apicodec/codec_v2_test.go +++ b/internal/apicodec/codec_v2_test.go @@ -530,6 +530,23 @@ func (suite *testCodecV2Suite) TestDecodeKeyError() { }, decoded.DebugInfo.MvccInfo[0].Mvcc.Lock.Secondaries) }, }, + { + name: "LockUpgradeConflict", + err: &kvrpcpb.KeyError{ + LockUpgradeConflict: &kvrpcpb.LockUpgradeConflict{ + Key: append(keyspacePrefix, []byte("key1")...), + StartTs: 11, + OwnerStartTs: 22, + Reason: kvrpcpb.LockUpgradeConflict_DuplicateInFlight, + }, + }, + validate: func(decoded *kvrpcpb.KeyError) { + re.Equal([]byte("key1"), decoded.LockUpgradeConflict.Key) + re.Equal(uint64(11), decoded.LockUpgradeConflict.StartTs) + re.Equal(uint64(22), decoded.LockUpgradeConflict.OwnerStartTs) + re.Equal(kvrpcpb.LockUpgradeConflict_DuplicateInFlight, decoded.LockUpgradeConflict.Reason) + }, + }, } codec := suite.codec diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index 8adc8c7637..d93e98ad30 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -1407,6 +1407,10 @@ func isKnownSharedLockUpgradeFailure(err error) bool { if errors.As(err, &deadlock) { return true } + var lockUpgradeConflict *tikverr.ErrLockUpgradeConflict + if errors.As(err, &lockUpgradeConflict) { + return true + } var assertionFailed *tikverr.ErrAssertionFailed return errors.As(err, &assertionFailed) } diff --git a/txnkv/transaction/txn_test.go b/txnkv/transaction/txn_test.go index 8a5aa5d4ca..360c5fcbaf 100644 --- a/txnkv/transaction/txn_test.go +++ b/txnkv/transaction/txn_test.go @@ -16,6 +16,7 @@ package transaction import ( "context" + stderrs "errors" "testing" "time" @@ -482,6 +483,56 @@ func TestSharedLockUpgrade(t *testing.T) { keysAsStrings(TxnProbe{KVTxn: txn.KVTxn}.CollectLockedKeys())) }) + t.Run("LockUpgradeConflictReturnsTypedErrorWithoutRetrySemantics", func(t *testing.T) { + primaryKey := []byte("primary-key") + upgradeKey := []byte("upgrade-key") + txn, requests := newRecorderTxn(t, func(callIndex int, req *kvrpcpb.PessimisticLockRequest) (*tikvrpc.Response, error) { + if len(req.Mutations) == 1 && + req.Mutations[0].Op == kvrpcpb.Op_PessimisticLock && + string(req.Mutations[0].Key) == string(upgradeKey) { + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{ + Errors: []*kvrpcpb.KeyError{{ + LockUpgradeConflict: &kvrpcpb.LockUpgradeConflict{ + Key: upgradeKey, + StartTs: 1, + OwnerStartTs: 2, + Reason: kvrpcpb.LockUpgradeConflict_SecondUpgrader, + }, + }}, + }}, nil + } + return &tikvrpc.Response{Resp: &kvrpcpb.PessimisticLockResponse{}}, nil + }) + lockSharedKey(t, txn, primaryKey, upgradeKey) + *requests = (*requests)[:0] + + lockCtx := kv.NewLockCtx(2, kv.LockNoWait, time.Now()) + lockCtx.AllowSharedLockUpgrade = true + err := txn.lockKeys(context.TODO(), lockCtx, nil, upgradeKey) + require.Error(t, err) + require.Len(t, *requests, 1) + require.Equal(t, []string{string(upgradeKey)}, keysAsStrings((*requests)[0].keys)) + require.False(t, tikverr.IsErrWriteConflict(err)) + require.False(t, tikverr.IsErrorUndetermined(err)) + + var retryable *tikverr.ErrRetryable + require.False(t, stderrs.As(err, &retryable)) + var deadlock *tikverr.ErrDeadlock + require.False(t, stderrs.As(err, &deadlock)) + + var conflict *tikverr.ErrLockUpgradeConflict + require.ErrorAs(t, err, &conflict) + require.Equal(t, []byte("upgrade-key"), conflict.Key) + require.Equal(t, uint64(1), conflict.StartTs) + require.Equal(t, uint64(2), conflict.OwnerStartTs) + require.Equal(t, kvrpcpb.LockUpgradeConflict_SecondUpgrader, conflict.Reason) + + flags, getErr := txn.GetMemBuffer().GetFlags(upgradeKey) + require.NoError(t, getErr) + require.True(t, flags.HasLocked()) + require.True(t, flags.HasLockedInShareMode()) + }) + t.Run("OutcomeUnknownUpgradeFailureIsTransactionFatal", func(t *testing.T) { primaryKey := []byte("primary-key") upgradeKey := []byte("upgrade-key") From 7c485409d66b3e5c1fbc3902509a9c07a81dc6c0 Mon Sep 17 00:00:00 2001 From: Wenxuan Zhang Date: Thu, 28 May 2026 11:02:59 +0800 Subject: [PATCH 4/7] refactor(transaction): rename lock upgrade outcome helper 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. --- txnkv/transaction/txn.go | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index d93e98ad30..acb48251dd 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -1394,25 +1394,35 @@ func (txn *KVTxn) getUndeterminedLockStateErr() error { return nil } -func isKnownSharedLockUpgradeFailure(err error) bool { +// isLockUpgradeResultUndetermined reports whether a shared-to-exclusive +// upgrade error leaves the remote lock state uncertain. +// +// Upgrade requests are special because the transaction already owns the shared +// lock locally. If TiKV returns a deterministic conflict or validation error, +// we know the new exclusive lock was not granted, so the caller should return +// that error directly and keep the original shared-lock state. Any error +// outside that known-safe set is treated as undetermined and must poison the +// transaction to avoid continuing with divergent local/remote lock state. +func isLockUpgradeResultUndetermined(err error) bool { if err == nil { return false } - if tikverr.IsErrWriteConflict(err) || tikverr.IsErrKeyExist(err) || + if tikverr.IsErrWriteConflict(err) || + tikverr.IsErrKeyExist(err) || errors.Is(err, tikverr.ErrLockAcquireFailAndNoWaitSet) || errors.Is(err, tikverr.ErrLockWaitTimeout) { - return true + return false } var deadlock *tikverr.ErrDeadlock if errors.As(err, &deadlock) { - return true + return false } var lockUpgradeConflict *tikverr.ErrLockUpgradeConflict if errors.As(err, &lockUpgradeConflict) { - return true + return false } var assertionFailed *tikverr.ErrAssertionFailed - return errors.As(err, &assertionFailed) + return !errors.As(err, &assertionFailed) } func (txn *KVTxn) lockPessimisticKeyGroup( @@ -1446,7 +1456,7 @@ func (txn *KVTxn) lockPessimisticKeyGroup( } if isUpgrade { - if !isKnownSharedLockUpgradeFailure(err) { + if isLockUpgradeResultUndetermined(err) { txn.committer.setUndeterminedErr(err) return 0, errors.WithStack(tikverr.ErrResultUndetermined) } From 2aa5c41c211db9896a037460e83880d832523ce6 Mon Sep 17 00:00:00 2001 From: Wenxuan Zhang Date: Thu, 28 May 2026 11:49:13 +0800 Subject: [PATCH 5/7] refactor(transaction): remove redundant rollbackKeys parameter 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. --- txnkv/transaction/txn.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index acb48251dd..00832cd52d 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -1429,7 +1429,6 @@ func (txn *KVTxn) lockPessimisticKeyGroup( ctx context.Context, lockCtx *tikv.LockCtx, keys [][]byte, - rollbackKeys [][]byte, isUpgrade bool, ) (int, error) { bo := retry.NewBackofferWithVars(ctx, pessimisticLockMaxBackoff, txn.vars) @@ -1479,7 +1478,7 @@ func (txn *KVTxn) lockPessimisticKeyGroup( if lockCtx.MaxLockedWithConflictTS > rollbackForUpdateTS { rollbackForUpdateTS = lockCtx.MaxLockedWithConflictTS } - wg := txn.asyncPessimisticRollback(ctx, rollbackKeys, rollbackForUpdateTS) + wg := txn.asyncPessimisticRollback(ctx, keys, rollbackForUpdateTS) if isDeadlock { logutil.Logger(ctx).Debug("deadlock error received", zap.Uint64("startTS", txn.startTS), zap.Stringer("deadlockInfo", dl)) @@ -1564,7 +1563,7 @@ func (txn *KVTxn) lockKeysWithSharedLockUpgrade( lockedInThisCall := 0 if len(normalExclusiveKeys) > 0 { - skipped, err := txn.lockPessimisticKeyGroup(ctx, lockCtx, normalExclusiveKeys, normalExclusiveKeys, false) + skipped, err := txn.lockPessimisticKeyGroup(ctx, lockCtx, normalExclusiveKeys, false) if err != nil { if assignedPrimaryKey && lockedInThisCall == 0 { txn.resetPrimary(false) @@ -1575,7 +1574,7 @@ func (txn *KVTxn) lockKeysWithSharedLockUpgrade( } for _, key := range upgradeKeys { - skipped, err := txn.lockPessimisticKeyGroup(ctx, lockCtx, [][]byte{key}, nil, true) + skipped, err := txn.lockPessimisticKeyGroup(ctx, lockCtx, [][]byte{key}, true) if err != nil { if assignedPrimaryKey && lockedInThisCall == 0 { txn.resetPrimary(false) From 988e697e35d379753035038f28e156eac95af7a0 Mon Sep 17 00:00:00 2001 From: Wenxuan Zhang Date: Thu, 28 May 2026 14:58:56 +0800 Subject: [PATCH 6/7] refactor(txnkv): narrow LockOnlyIfExists err key scope Move the temporary lock key into the LockOnlyIfExists validation branch and rename it to lockKey so the variable name matches its diagnostic-only role. --- txnkv/transaction/txn.go | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index 00832cd52d..14b42dbd39 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -1731,18 +1731,18 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() if len(keys) == 0 && len(upgradeKeys) == 0 { return nil } - var lockOnlyIfExistsKey []byte - if len(keys) > 0 { - lockOnlyIfExistsKey = keys[0] - } else { - lockOnlyIfExistsKey = upgradeKeys[0] - } if lockCtx.LockOnlyIfExists { + var lockKey []byte + if len(keys) > 0 { + lockKey = keys[0] + } else { + lockKey = upgradeKeys[0] + } if !lockCtx.ReturnValues { return &tikverr.ErrLockOnlyIfExistsNoReturnValue{ StartTS: txn.startTS, ForUpdateTs: lockCtx.ForUpdateTS, - LockKey: lockOnlyIfExistsKey, + LockKey: lockKey, } } // It can't transform LockOnlyIfExists mode to normal mode. If so, it can add a lock to a key @@ -1752,7 +1752,7 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func() return &tikverr.ErrLockOnlyIfExistsNoPrimaryKey{ StartTS: txn.startTS, ForUpdateTs: lockCtx.ForUpdateTS, - LockKey: lockOnlyIfExistsKey, + LockKey: lockKey, } } } From 2741a6a7258d1a37fc5b4a6afc5b3041f7e27195 Mon Sep 17 00:00:00 2001 From: Wenxuan Zhang Date: Thu, 28 May 2026 16:16:33 +0800 Subject: [PATCH 7/7] chore: add comment --- txnkv/transaction/txn.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index 14b42dbd39..ba133dc87f 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -1547,6 +1547,9 @@ func (txn *KVTxn) lockKeysWithSharedLockUpgrade( assignedPrimaryKey := false totalKeys := len(normalExclusiveKeys) + len(upgradeKeys) 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 {