From ef433cc60f2bcba44e242e02596381dce06d3b55 Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Thu, 29 Apr 2021 19:46:33 +0800
Subject: [PATCH] store/tikv: remove use of TaskID transaction option in
 store/tikv

Signed-off-by: disksing <i@disksing.com>
---
 store/driver/txn/snapshot.go   |  2 ++
 store/driver/txn/txn_driver.go |  2 ++
 store/tikv/snapshot.go         | 12 ++++++++----
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/store/driver/txn/snapshot.go b/store/driver/txn/snapshot.go
index 3d41ef8cfa1ae..d87ed3d2f574d 100644
--- a/store/driver/txn/snapshot.go
+++ b/store/driver/txn/snapshot.go
@@ -73,6 +73,8 @@ func (s *tikvSnapshot) SetOption(opt int, val interface{}) {
 		s.KVSnapshot.SetNotFillCache(val.(bool))
 	case tikvstore.SnapshotTS:
 		s.KVSnapshot.SetSnapshotTS(val.(uint64))
+	case tikvstore.TaskID:
+		s.KVSnapshot.SetTaskID(val.(uint64))
 	default:
 		s.KVSnapshot.SetOption(opt, val)
 	}
diff --git a/store/driver/txn/txn_driver.go b/store/driver/txn/txn_driver.go
index 968d0f8701233..97ed943dc1083 100644
--- a/store/driver/txn/txn_driver.go
+++ b/store/driver/txn/txn_driver.go
@@ -144,6 +144,8 @@ func (txn *tikvTxn) SetOption(opt int, val interface{}) {
 		txn.SetPessimistic(val.(bool))
 	case tikvstore.SnapshotTS:
 		txn.KVTxn.GetSnapshot().SetSnapshotTS(val.(uint64))
+	case tikvstore.TaskID:
+		txn.KVTxn.GetSnapshot().SetTaskID(val.(uint64))
 	case tikvstore.CommitHook:
 		txn.SetCommitCallback(val.(func(string, error)))
 	default:
diff --git a/store/tikv/snapshot.go b/store/tikv/snapshot.go
index 0b09a408d7dc9..96b4f478355cb 100644
--- a/store/tikv/snapshot.go
+++ b/store/tikv/snapshot.go
@@ -570,10 +570,6 @@ func (s *KVSnapshot) SetOption(opt int, val interface{}) {
 		s.mu.Lock()
 		s.mu.replicaRead = val.(kv.ReplicaReadType)
 		s.mu.Unlock()
-	case kv.TaskID:
-		s.mu.Lock()
-		s.mu.taskID = val.(uint64)
-		s.mu.Unlock()
 	case kv.CollectRuntimeStats:
 		s.mu.Lock()
 		s.mu.stats = val.(*SnapshotRuntimeStats)
@@ -628,6 +624,14 @@ func (s *KVSnapshot) SetPriority(pri Priority) {
 	s.priority = pri
 }
 
+// SetTaskID marks current task's unique ID to allow TiKV to schedule
+// tasks more fairly.
+func (s *KVSnapshot) SetTaskID(id uint64) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.mu.taskID = id
+}
+
 // SnapCacheHitCount gets the snapshot cache hit count. Only for test.
 func (s *KVSnapshot) SnapCacheHitCount() int {
 	return int(atomic.LoadInt64(&s.mu.hitCnt))