From e96f8e6e720ca11a0ad21a68eeb505b15194ed29 Mon Sep 17 00:00:00 2001 From: Ti Chi Robot Date: Tue, 17 Jan 2023 19:09:49 +0800 Subject: [PATCH] session: fix data race in the LazyTxn.LockKeys (#40350) (#40389) close pingcap/tidb#40355 --- go.mod | 2 +- go.sum | 4 ++-- kv/interface_mock_test.go | 7 +++++++ kv/kv.go | 4 ++++ session/txn.go | 26 +++++++++++++++++--------- store/driver/txn/txn_driver.go | 6 ++++++ 6 files changed, 37 insertions(+), 12 deletions(-) diff --git a/go.mod b/go.mod index 39f37fd88565f..95872b570f544 100644 --- a/go.mod +++ b/go.mod @@ -63,7 +63,7 @@ require ( github.com/spf13/pflag v1.0.5 github.com/stretchr/testify v1.7.2-0.20220504104629-106ec21d14df github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 - github.com/tikv/client-go/v2 v2.0.1-0.20221012074928-624e0ed3cc67 + github.com/tikv/client-go/v2 v2.0.1-0.20230117081319-35a262e90d9b github.com/tikv/pd/client v0.0.0-20220307081149-841fa61e9710 github.com/twmb/murmur3 v1.1.3 github.com/uber/jaeger-client-go v2.22.1+incompatible diff --git a/go.sum b/go.sum index c677fbec891d0..40babd153241d 100644 --- a/go.sum +++ b/go.sum @@ -755,8 +755,8 @@ github.com/stretchr/testify v1.7.2-0.20220504104629-106ec21d14df/go.mod h1:6Fq8o github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69rRypqCw= github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJfDRtkanvQPiooDH8HvJ2FBh+iKT/OmiQQ= github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU= -github.com/tikv/client-go/v2 v2.0.1-0.20221012074928-624e0ed3cc67 h1:ofj1E/98HrrEpXi+RmjBhnBY3+/9ObyorIkloGbvEa0= -github.com/tikv/client-go/v2 v2.0.1-0.20221012074928-624e0ed3cc67/go.mod h1:VTlli8fRRpcpISj9I2IqroQmcAFfaTyBquiRhofOcDs= +github.com/tikv/client-go/v2 v2.0.1-0.20230117081319-35a262e90d9b h1:IUH/4BrP9BJm7to+XUJslcwaZONuIEwwClBnlrO7zJM= +github.com/tikv/client-go/v2 v2.0.1-0.20230117081319-35a262e90d9b/go.mod h1:VTlli8fRRpcpISj9I2IqroQmcAFfaTyBquiRhofOcDs= github.com/tikv/pd/client v0.0.0-20220307081149-841fa61e9710 h1:jxgmKOscXSjaFEKQGRyY5qOpK8hLqxs2irb/uDJMtwk= github.com/tikv/pd/client v0.0.0-20220307081149-841fa61e9710/go.mod h1:AtvppPwkiyUgQlR1W9qSqfTB+OsOIu19jDCOxOsPkmU= github.com/tklauser/go-sysconf v0.3.9 h1:JeUVdAOWhhxVcU6Eqr/ATFHgXk/mmiItdKeJPev3vTo= diff --git a/kv/interface_mock_test.go b/kv/interface_mock_test.go index 26f9acf1d6cfb..3ea9fe5aa37a5 100644 --- a/kv/interface_mock_test.go +++ b/kv/interface_mock_test.go @@ -52,6 +52,13 @@ func (t *mockTxn) LockKeys(_ context.Context, _ *LockCtx, _ ...Key) error { return nil } +func (t *mockTxn) LockKeysFunc(_ context.Context, _ *LockCtx, fn func(), _ ...Key) error { + if fn != nil { + fn() + } + return nil +} + func (t *mockTxn) SetOption(opt int, val interface{}) { t.opts[opt] = val } diff --git a/kv/kv.go b/kv/kv.go index 58aecb5195891..aa52227118c71 100644 --- a/kv/kv.go +++ b/kv/kv.go @@ -199,6 +199,10 @@ type Transaction interface { // LockKeys tries to lock the entries with the keys in KV store. // Will block until all keys are locked successfully or an error occurs. LockKeys(ctx context.Context, lockCtx *LockCtx, keys ...Key) error + // LockKeysFunc tries to lock the entries with the keys in KV store. + // Will block until all keys are locked successfully or an error occurs. + // fn is called before LockKeys unlocks the keys. + LockKeysFunc(ctx context.Context, lockCtx *LockCtx, fn func(), keys ...Key) error // SetOption sets an option with a value, when val is nil, uses the default // value of this option. SetOption(opt int, val interface{}) diff --git a/session/txn.go b/session/txn.go index 0620d59d1cf3c..8d359f94e60f4 100644 --- a/session/txn.go +++ b/session/txn.go @@ -378,6 +378,11 @@ func (txn *LazyTxn) Rollback() error { // LockKeys Wrap the inner transaction's `LockKeys` to record the status func (txn *LazyTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keys ...kv.Key) error { + return txn.LockKeysFunc(ctx, lockCtx, nil, keys...) +} + +// LockKeysFunc Wrap the inner transaction's `LockKeys` to record the status +func (txn *LazyTxn) LockKeysFunc(ctx context.Context, lockCtx *kv.LockCtx, fn func(), keys ...kv.Key) error { failpoint.Inject("beforeLockKeys", func() {}) t := time.Now() @@ -389,15 +394,18 @@ func (txn *LazyTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keys ...k txn.mu.TxnInfo.BlockStartTime.Time = t txn.mu.Unlock() - err := txn.Transaction.LockKeys(ctx, lockCtx, keys...) - - txn.mu.Lock() - defer txn.mu.Unlock() - txn.mu.TxnInfo.State = originState - txn.mu.TxnInfo.BlockStartTime.Valid = false - txn.mu.TxnInfo.EntriesCount = uint64(txn.Transaction.Len()) - txn.mu.TxnInfo.EntriesSize = uint64(txn.Transaction.Size()) - return err + lockFunc := func() { + if fn != nil { + fn() + } + txn.mu.Lock() + defer txn.mu.Unlock() + txn.mu.TxnInfo.State = originState + txn.mu.TxnInfo.BlockStartTime.Valid = false + txn.mu.TxnInfo.EntriesCount = uint64(txn.Transaction.Len()) + txn.mu.TxnInfo.EntriesSize = uint64(txn.Transaction.Size()) + } + return txn.Transaction.LockKeysFunc(ctx, lockCtx, lockFunc, keys...) } func (txn *LazyTxn) reset() { diff --git a/store/driver/txn/txn_driver.go b/store/driver/txn/txn_driver.go index fc127be0013f2..b12697a2efc6b 100644 --- a/store/driver/txn/txn_driver.go +++ b/store/driver/txn/txn_driver.go @@ -75,6 +75,12 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keysInput return txn.extractKeyErr(err) } +func (txn *tikvTxn) LockKeysFunc(ctx context.Context, lockCtx *kv.LockCtx, fn func(), keysInput ...kv.Key) error { + keys := toTiKVKeys(keysInput) + err := txn.KVTxn.LockKeysFunc(ctx, lockCtx, fn, keys...) + return txn.extractKeyErr(err) +} + func (txn *tikvTxn) Commit(ctx context.Context) error { err := txn.KVTxn.Commit(ctx) return txn.extractKeyErr(err)