From 9417f095c4bf4908a47b6c7701d609e167d3924d Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Wed, 19 May 2021 11:08:59 +0800
Subject: [PATCH] store/tikv: extract methods for LockCtx

Signed-off-by: disksing <i@disksing.com>
---
 executor/batch_point_get.go  |  9 +++------
 executor/point_get.go        | 14 ++++----------
 store/tikv/kv/kv.go          | 28 ++++++++++++++++++++++++++++
 store/tikv/tests/2pc_test.go |  3 +--
 4 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/executor/batch_point_get.go b/executor/batch_point_get.go
index 2137884c69745..d9defde9e0558 100644
--- a/executor/batch_point_get.go
+++ b/executor/batch_point_get.go
@@ -29,7 +29,6 @@ import (
 	"github.com/pingcap/tidb/sessionctx/variable"
 	driver "github.com/pingcap/tidb/store/driver/txn"
 	"github.com/pingcap/tidb/store/tikv"
-	tikvstore "github.com/pingcap/tidb/store/tikv/kv"
 	"github.com/pingcap/tidb/store/tikv/oracle"
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/types"
@@ -418,8 +417,7 @@ func LockKeys(ctx context.Context, seCtx sessionctx.Context, lockWaitTime int64,
 	txnCtx := seCtx.GetSessionVars().TxnCtx
 	lctx := newLockCtx(seCtx.GetSessionVars(), lockWaitTime)
 	if txnCtx.IsPessimistic {
-		lctx.ReturnValues = true
-		lctx.Values = make(map[string]tikvstore.ReturnedValue, len(keys))
+		lctx.InitReturnValues(len(keys))
 	}
 	err := doLockKeys(ctx, seCtx, lctx, keys...)
 	if err != nil {
@@ -429,9 +427,8 @@ func LockKeys(ctx context.Context, seCtx sessionctx.Context, lockWaitTime int64,
 		// When doLockKeys returns without error, no other goroutines access the map,
 		// it's safe to read it without mutex.
 		for _, key := range keys {
-			rv := lctx.Values[string(key)]
-			if !rv.AlreadyLocked {
-				txnCtx.SetPessimisticLockCache(key, rv.Value)
+			if v, ok := lctx.GetValueNotLocked(key); ok {
+				txnCtx.SetPessimisticLockCache(key, v)
 			}
 		}
 	}
diff --git a/executor/point_get.go b/executor/point_get.go
index fc8326555bf01..140da5f7eed93 100644
--- a/executor/point_get.go
+++ b/executor/point_get.go
@@ -31,7 +31,6 @@ import (
 	plannercore "github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/store/tikv"
-	tikvstore "github.com/pingcap/tidb/store/tikv/kv"
 	"github.com/pingcap/tidb/store/tikv/oracle"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/table/tables"
@@ -322,19 +321,14 @@ func (e *PointGetExecutor) lockKeyIfNeeded(ctx context.Context, key []byte) erro
 	if e.lock {
 		seVars := e.ctx.GetSessionVars()
 		lockCtx := newLockCtx(seVars, e.lockWaitTime)
-		lockCtx.ReturnValues = true
-		lockCtx.Values = map[string]tikvstore.ReturnedValue{}
+		lockCtx.InitReturnValues(1)
 		err := doLockKeys(ctx, e.ctx, lockCtx, key)
 		if err != nil {
 			return err
 		}
-		lockCtx.ValuesLock.Lock()
-		defer lockCtx.ValuesLock.Unlock()
-		for key, val := range lockCtx.Values {
-			if !val.AlreadyLocked {
-				seVars.TxnCtx.SetPessimisticLockCache(kv.Key(key), val.Value)
-			}
-		}
+		lockCtx.IterateValuesNotLocked(func(k, v []byte) {
+			seVars.TxnCtx.SetPessimisticLockCache(kv.Key(k), v)
+		})
 		if len(e.handleVal) > 0 {
 			seVars.TxnCtx.SetPessimisticLockCache(e.idxKey, e.handleVal)
 		}
diff --git a/store/tikv/kv/kv.go b/store/tikv/kv/kv.go
index 8ba36a749db4f..eea6fd796888b 100644
--- a/store/tikv/kv/kv.go
+++ b/store/tikv/kv/kv.go
@@ -29,3 +29,31 @@ type LockCtx struct {
 	Stats                 *util.LockKeysDetails
 	ResourceGroupTag      []byte
 }
+
+// InitReturnValues creates the map to store returned value.
+func (ctx *LockCtx) InitReturnValues(valueLen int) {
+	ctx.ReturnValues = true
+	ctx.Values = make(map[string]ReturnedValue, valueLen)
+}
+
+// GetValueNotLocked returns a value if the key is not already locked.
+// (nil, false) means already locked.
+func (ctx *LockCtx) GetValueNotLocked(key []byte) ([]byte, bool) {
+	rv := ctx.Values[string(key)]
+	if !rv.AlreadyLocked {
+		return rv.Value, true
+	}
+	return nil, false
+}
+
+// IterateValuesNotLocked applies f to all key-values that are not already
+// locked.
+func (ctx *LockCtx) IterateValuesNotLocked(f func([]byte, []byte)) {
+	ctx.ValuesLock.Lock()
+	defer ctx.ValuesLock.Unlock()
+	for key, val := range ctx.Values {
+		if !val.AlreadyLocked {
+			f([]byte(key), val.Value)
+		}
+	}
+}
diff --git a/store/tikv/tests/2pc_test.go b/store/tikv/tests/2pc_test.go
index 5589752043b2b..9e45423bdae04 100644
--- a/store/tikv/tests/2pc_test.go
+++ b/store/tikv/tests/2pc_test.go
@@ -713,8 +713,7 @@ func (s *testCommitterSuite) TestPessimisticLockReturnValues(c *C) {
 	txn = s.begin(c)
 	txn.SetPessimistic(true)
 	lockCtx := &kv.LockCtx{ForUpdateTS: txn.StartTS(), WaitStartTime: time.Now()}
-	lockCtx.ReturnValues = true
-	lockCtx.Values = map[string]kv.ReturnedValue{}
+	lockCtx.InitReturnValues(2)
 	c.Assert(txn.LockKeys(context.Background(), lockCtx, key, key2), IsNil)
 	c.Assert(lockCtx.Values, HasLen, 2)
 	c.Assert(lockCtx.Values[string(key)].Value, BytesEquals, key)