Repository: tikv/tikv (language: rust)

Issue:
scheduler can read locks by rocksdb cursor
## Development Task

When handling commit and resolve_lock operations, TiKV reads the related locks before writing. The read is performed through multiple individual get operations, which results in numerous slow seek operations in RocksDB. Below is a CPU profile and resolve-lock scheduler metrics during the lock resolution period.

![Image](https://github.com/user-attachments/assets/1460f810-318b-4fa7-a0cd-794f59003b51)

![Image](https://github.com/user-attachments/assets/f0851c7d-ea27-4dfe-934a-3775117e767d)

The root cause of this behavior is that the reader is not running in scan mode (`self.lock_cursor` is none). So the `near_seek` operation doesn't work.

https://github.com/tikv/tikv/blob/caf82d2864d1501af733bf60c99a7732e105144b/src/storage/mvcc/reader/reader.rs#L243-L254

## Usages need to be optimized

https://github.com/tikv/tikv/blob/caf82d2864d1501af733bf60c99a7732e105144b/src/storage/txn/commands/resolve_lock.rs#L86-L89

https://github.com/tikv/tikv/blob/caf82d2864d1501af733bf60c99a7732e105144b/src/storage/txn/commands/commit.rs#L58-L61

https://github.com/tikv/tikv/blob/caf82d2864d1501af733bf60c99a7732e105144b/src/storage/txn/commands/prewrite.rs#L540-L543

https://github.com/tikv/tikv/blob/f0fc694e1f73871b694963191a8277aacbf372f3/src/storage/txn/commands/rollback.rs#L53-L56

https://github.com/tikv/tikv/blob/f0fc694e1f73871b694963191a8277aacbf372f3/src/storage/txn/commands/acquire_pessimistic_lock.rs#L91-L94

https://github.com/tikv/tikv/blob/f0fc694e1f73871b694963191a8277aacbf372f3/src/storage/txn/commands/check_secondary_locks.rs#L155-L158

https://github.com/tikv/tikv/blob/f0fc694e1f73871b694963191a8277aacbf372f3/src/storage/txn/commands/flush.rs#L79-L82


Discussion / hints:

Fix patch:
diff --git a/src/storage/mvcc/reader/reader.rs b/src/storage/mvcc/reader/reader.rs
index 5e26f964521..3bc50091048 100644
--- a/src/storage/mvcc/reader/reader.rs
+++ b/src/storage/mvcc/reader/reader.rs
@@ -122,6 +122,18 @@ impl<S: EngineSnapshot> SnapshotReader<S> {
     pub fn take_statistics(&mut self) -> Statistics {
         std::mem::take(&mut self.reader.statistics)
     }
+
+    #[inline(always)]
+    pub fn setup_with_hint_items<T>(&mut self, items: &mut [T], key_of: fn(&T) -> &Key) {
+        // enable scan mode if there are multiple items, so that we don't need to seek
+        // for every key.
+        if items.len() > 1 {
+            items.sort_by(|a, b| key_of(a).cmp(key_of(b)));
+            self.reader.scan_mode = Some(ScanMode::Forward);
+            self.reader
+                .set_range(Some(key_of(items.first().unwrap()).clone()), None);
+        }
+    }
 }
 
 pub struct MvccReader<S: EngineSnapshot> {
diff --git a/src/storage/txn/commands/acquire_pessimistic_lock.rs b/src/storage/txn/commands/acquire_pessimistic_lock.rs
index 1b18b7f27bc..da8367ed49a 100644
--- a/src/storage/txn/commands/acquire_pessimistic_lock.rs
+++ b/src/storage/txn/commands/acquire_pessimistic_lock.rs
@@ -86,12 +86,12 @@ impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for AcquirePessimisticLock
             ))));
         }
 
-        let (start_ts, ctx, keys) = (self.start_ts, self.ctx, self.keys);
+        let (start_ts, ctx, mut keys) = (self.start_ts, self.ctx, self.keys);
         let mut txn = MvccTxn::new(start_ts, context.concurrency_manager);
-        let mut reader = ReaderWithStats::new(
-            SnapshotReader::new_with_ctx(start_ts, snapshot, &ctx),
-            context.statistics,
-        );
+
+        let mut snapshot_reader = SnapshotReader::new_with_ctx(start_ts, snapshot, &ctx);
+        snapshot_reader.setup_with_hint_items(&mut keys, |k| &k.0);
+        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);
 
         let total_keys = keys.len();
         let mut res = PessimisticLockResults::with_capacity(total_keys);
diff --git a/src/storage/txn/commands/check_secondary_locks.rs b/src/storage/txn/commands/check_secondary_locks.rs
index 6664f16bdc6..bfc845ba691 100644
--- a/src/storage/txn/commands/check_secondary_locks.rs
+++ b/src/storage/txn/commands/check_secondary_locks.rs
@@ -141,7 +141,7 @@ fn check_status_from_lock<S: Snapshot>(
 }
 
 impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for CheckSecondaryLocks {
-    fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
+    fn process_write(mut self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
         // It is not allowed for commit to overwrite a protected rollback. So we update
         // max_ts to prevent this case from happening.
         let region_id = self.ctx.get_region_id();
@@ -152,10 +152,10 @@ impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for CheckSecondaryLocks {
             })?;
 
         let mut txn = MvccTxn::new(self.start_ts, context.concurrency_manager);
-        let mut reader = ReaderWithStats::new(
-            SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx),
-            context.statistics,
-        );
+
+        let mut snapshot_reader = SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx);
+        snapshot_reader.setup_with_hint_items(&mut self.keys, |k| k);
+        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);
         let mut released_locks = ReleasedLocks::new();
         let mut result = SecondaryLocksStatus::Locked(Vec::new());
 
diff --git a/src/storage/txn/commands/commit.rs b/src/storage/txn/commands/commit.rs
index f863dcfc105..6f515fccb32 100644
--- a/src/storage/txn/commands/commit.rs
+++ b/src/storage/txn/commands/commit.rs
@@ -55,15 +55,15 @@ impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for Commit {
             }));
         }
         let mut txn = MvccTxn::new(self.lock_ts, context.concurrency_manager);
-        let mut reader = ReaderWithStats::new(
-            SnapshotReader::new_with_ctx(self.lock_ts, snapshot, &self.ctx),
-            context.statistics,
-        );
+        let mut keys = self.keys;
+        let mut snapshot_reader = SnapshotReader::new_with_ctx(self.lock_ts, snapshot, &self.ctx);
+        snapshot_reader.setup_with_hint_items(&mut keys, |k| k);
+        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);
 
-        let rows = self.keys.len();
+        let rows = keys.len();
         // Pessimistic txn needs key_hashes to wake up waiters
         let mut released_locks = ReleasedLocks::new();
-        for k in self.keys {
+        for k in keys {
             released_locks.push(commit(&mut txn, &mut reader, k, self.commit_ts)?);
         }
 
diff --git a/src/storage/txn/commands/flush.rs b/src/storage/txn/commands/flush.rs
index 519786781de..faf8d4689d6 100644
--- a/src/storage/txn/commands/flush.rs
+++ b/src/storage/txn/commands/flush.rs
@@ -2,8 +2,8 @@
 
 use std::mem;
 
-use kvproto::kvrpcpb::{AssertionLevel, ExtraOp, PrewriteRequestPessimisticAction};
 // #[PerformanceCriticalPath]
+use kvproto::kvrpcpb::{AssertionLevel, ExtraOp, PrewriteRequestPessimisticAction};
 use txn_types::{insert_old_value_if_resolved, Mutation, OldValues, TimeStamp, TxnExtra};
 
 use crate::storage::{
@@ -76,10 +76,10 @@ impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for Flush {
         }
         let rows = self.mutations.len();
         let mut txn = MvccTxn::new(self.start_ts, context.concurrency_manager);
-        let mut reader = ReaderWithStats::new(
-            SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx),
-            context.statistics,
-        );
+
+        let mut snapshot_reader = SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx);
+        snapshot_reader.setup_with_hint_items(&mut self.mutations, |m| m.key());
+        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);
         let mut old_values = Default::default();
 
         let res = self.flush(&mut txn, &mut reader, &mut old_values, context.extra_op);
diff --git a/src/storage/txn/commands/prewrite.rs b/src/storage/txn/commands/prewrite.rs
index d9c7a65921b..252f12ca99e 100644
--- a/src/storage/txn/commands/prewrite.rs
+++ b/src/storage/txn/commands/prewrite.rs
@@ -537,10 +537,9 @@ impl<K: PrewriteKind> Prewriter<K> {
         self.check_max_ts_synced(&snapshot)?;
 
         let mut txn = MvccTxn::new(self.start_ts, context.concurrency_manager);
-        let mut reader = ReaderWithStats::new(
-            SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx),
-            context.statistics,
-        );
+        let mut snapshot_reader = SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx);
+        snapshot_reader.setup_with_hint_items(&mut self.mutations, |m| m.key());
+        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);
         // Set extra op here for getting the write record when check write conflict in
         // prewrite.
 
@@ -616,7 +615,6 @@ impl<K: PrewriteKind> Prewriter<K> {
 
         // If there are other errors, return other error prior to `AssertionFailed`.
         let mut assertion_failure = None;
-
         for m in mem::take(&mut self.mutations) {
             let pessimistic_action = m.pessimistic_action();
             let expected_for_update_ts = m.pessimistic_expected_for_update_ts();
@@ -885,6 +883,7 @@ trait MutationLock {
     fn pessimistic_action(&self) -> PrewriteRequestPessimisticAction;
     fn pessimistic_expected_for_update_ts(&self) -> Option<TimeStamp>;
     fn into_mutation(self) -> Mutation;
+    fn key(&self) -> &Key;
 }
 
 impl MutationLock for Mutation {
@@ -899,6 +898,10 @@ impl MutationLock for Mutation {
     fn into_mutation(self) -> Mutation {
         self
     }
+
+    fn key(&self) -> &Key {
+        self.key()
+    }
 }
 
 #[derive(Debug)]
@@ -926,6 +929,10 @@ impl MutationLock for PessimisticMutation {
     fn into_mutation(self) -> Mutation {
         self.mutation
     }
+
+    fn key(&self) -> &Key {
+        self.mutation.key()
+    }
 }
 
 impl PessimisticMutation {
@@ -1200,7 +1207,7 @@ mod tests {
         .unwrap();
         let d = perf.delta();
         assert_eq!(1, statistic.write.seek);
-        assert_eq!(d.internal_delete_skipped_count, 0);
+        assert_eq!(d.internal_delete_skipped_count, 40);
     }
 
     #[test]
diff --git a/src/storage/txn/commands/resolve_lock.rs b/src/storage/txn/commands/resolve_lock.rs
index b4749bab4df..6023f7aade9 100644
--- a/src/storage/txn/commands/resolve_lock.rs
+++ b/src/storage/txn/commands/resolve_lock.rs
@@ -80,13 +80,13 @@ impl CommandExt for ResolveLock {
 
 impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for ResolveLock {
     fn process_write(mut self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
-        let (ctx, txn_status, key_locks) = (self.ctx, self.txn_status, self.key_locks);
+        let (ctx, txn_status, mut key_locks) = (self.ctx, self.txn_status, self.key_locks);
 
         let mut txn = MvccTxn::new(TimeStamp::zero(), context.concurrency_manager);
-        let mut reader = ReaderWithStats::new(
-            SnapshotReader::new_with_ctx(TimeStamp::zero(), snapshot, &ctx),
-            context.statistics,
-        );
+
+        let mut snapshot_reader = SnapshotReader::new_with_ctx(TimeStamp::zero(), snapshot, &ctx);
+        snapshot_reader.setup_with_hint_items(&mut key_locks, |k| &k.0);
+        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);
 
         let mut scan_key = self.scan_key.take();
         let rows = key_locks.len();
diff --git a/src/storage/txn/commands/rollback.rs b/src/storage/txn/commands/rollback.rs
index 8f845ad57eb..0dc88c62822 100644
--- a/src/storage/txn/commands/rollback.rs
+++ b/src/storage/txn/commands/rollback.rs
@@ -48,12 +48,12 @@ impl CommandExt for Rollback {
 }
 
 impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for Rollback {
-    fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
+    fn process_write(mut self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
         let mut txn = MvccTxn::new(self.start_ts, context.concurrency_manager);
-        let mut reader = ReaderWithStats::new(
-            SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx),
-            context.statistics,
-        );
+
+        let mut snapshot_reader = SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx);
+        snapshot_reader.setup_with_hint_items(&mut self.keys, |k| k);
+        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);
 
         let rows = self.keys.len();
         let mut released_locks = ReleasedLocks::new();


Test patch: