From e797f02939858d37e5256c5819d14343e649cd81 Mon Sep 17 00:00:00 2001 From: Weizhen Wang Date: Fri, 6 Jan 2023 13:31:13 +0800 Subject: [PATCH] fix data race in the LockKeys (#655) Signed-off-by: Weizhen Wang --- txnkv/transaction/txn.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index a8bb345db..4727270a6 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -832,6 +832,17 @@ func (txn *KVTxn) filterAggressiveLockedKeys(lockCtx *tikv.LockCtx, allKeys [][] // LockKeys tries to lock the entries with the keys in KV store. // lockCtx is the context for lock, lockCtx.lockWaitTime in ms func (txn *KVTxn) LockKeys(ctx context.Context, lockCtx *tikv.LockCtx, keysInput ...[]byte) error { + return txn.lockKeys(ctx, lockCtx, nil, keysInput...) +} + +// LockKeysFunc tries to lock the entries with the keys in KV store. +// lockCtx is the context for lock, lockCtx.lockWaitTime in ms +// fn is a function which run before the lock is released. +func (txn *KVTxn) LockKeysFunc(ctx context.Context, lockCtx *tikv.LockCtx, fn func(), keysInput ...[]byte) error { + return txn.lockKeys(ctx, lockCtx, fn, keysInput...) +} + +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 // need to bind it to ctx so that the internal client can perceive and execute @@ -869,6 +880,11 @@ func (txn *KVTxn) LockKeys(ctx context.Context, lockCtx *tikv.LockCtx, keysInput } } }() + defer func() { + if fn != nil { + fn() + } + }() if !txn.IsPessimistic() && txn.aggressiveLockingContext != nil { return errors.New("trying to perform aggressive locking in optimistic transaction")