fix data race in the LockKeys (#672)

Signed-off-by: Weizhen Wang <wangweizhen@pingcap.com>

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