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/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..ba133dc87f 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,215 @@ 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 +} + +// 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) || + errors.Is(err, tikverr.ErrLockAcquireFailAndNoWaitSet) || + errors.Is(err, tikverr.ErrLockWaitTimeout) { + return false + } + var deadlock *tikverr.ErrDeadlock + if errors.As(err, &deadlock) { + return false + } + var lockUpgradeConflict *tikverr.ErrLockUpgradeConflict + if errors.As(err, &lockUpgradeConflict) { + return false + } + var assertionFailed *tikverr.ErrAssertionFailed + return !errors.As(err, &assertionFailed) +} + +func (txn *KVTxn) lockPessimisticKeyGroup( + ctx context.Context, + lockCtx *tikv.LockCtx, + keys [][]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 isLockUpgradeResultUndetermined(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, keys, 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 { + // 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) + } + + 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, 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}, 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 +1616,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 +1671,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 +1682,16 @@ 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 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. @@ -1484,7 +1707,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 +1722,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,28 +1731,41 @@ 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 { + 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: keys[0], + LockKey: lockKey, } } // 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 { 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..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" @@ -23,6 +24,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 +275,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 +295,281 @@ 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("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") + 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("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") + 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")