diff --git a/store/driver/txn/snapshot.go b/store/driver/txn/snapshot.go index 6064db02cdeba..ff8f4b79cc918 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.ReplicaRead: + s.KVSnapshot.SetReplicaRead(val.(tikvstore.ReplicaReadType)) case tikvstore.TaskID: s.KVSnapshot.SetTaskID(val.(uint64)) default: diff --git a/store/driver/txn/txn_driver.go b/store/driver/txn/txn_driver.go index 739b983fd3d99..10eef3dbfd7d8 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.ReplicaRead: + txn.KVTxn.GetSnapshot().SetReplicaRead(val.(tikvstore.ReplicaReadType)) case tikvstore.TaskID: txn.KVTxn.GetSnapshot().SetTaskID(val.(uint64)) case tikvstore.InfoSchema: diff --git a/store/tikv/snapshot.go b/store/tikv/snapshot.go index a37e1d8343c5e..d77be2c21ff35 100644 --- a/store/tikv/snapshot.go +++ b/store/tikv/snapshot.go @@ -565,10 +565,6 @@ func (s *KVSnapshot) IterReverse(k []byte) (unionstore.Iterator, error) { // value of this option. Only ReplicaRead is supported for snapshot func (s *KVSnapshot) SetOption(opt int, val interface{}) { switch opt { - case kv.ReplicaRead: - s.mu.Lock() - s.mu.replicaRead = val.(kv.ReplicaReadType) - s.mu.Unlock() case kv.CollectRuntimeStats: s.mu.Lock() s.mu.stats = val.(*SnapshotRuntimeStats) @@ -589,10 +585,6 @@ func (s *KVSnapshot) SetOption(opt int, val interface{}) { // DelOption deletes an option. func (s *KVSnapshot) DelOption(opt int) { switch opt { - case kv.ReplicaRead: - s.mu.Lock() - s.mu.replicaRead = kv.ReplicaReadLeader - s.mu.Unlock() case kv.CollectRuntimeStats: s.mu.Lock() s.mu.stats = nil @@ -611,6 +603,13 @@ func (s *KVSnapshot) SetKeyOnly(b bool) { s.keyOnly = b } +// SetReplicaRead sets up the replica read type. +func (s *KVSnapshot) SetReplicaRead(readType kv.ReplicaReadType) { + s.mu.Lock() + defer s.mu.Unlock() + s.mu.replicaRead = readType +} + // SetIsolationLevel sets the isolation level used to scan data from tikv. func (s *KVSnapshot) SetIsolationLevel(level IsoLevel) { s.isolationLevel = level