From 35a262e90d9b4b7160d5d0d8049b83b2a5d8340f Mon Sep 17 00:00:00 2001 From: Weizhen Wang Date: Tue, 17 Jan 2023 16:13:19 +0800 Subject: [PATCH] fix data race in the LockKeys (#672) 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 16fb1e4f1a..dd59a977e0 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -577,6 +577,17 @@ func (txn *KVTxn) LockKeysWithWaitTime(ctx context.Context, lockWaitTime int64, // 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 @@ -612,6 +623,11 @@ func (txn *KVTxn) LockKeys(ctx context.Context, lockCtx *tikv.LockCtx, keysInput } } }() + defer func() { + if fn != nil { + fn() + } + }() memBuf := txn.us.GetMemBuffer() // Avoid data race with concurrent updates to the memBuf