Repository: tikv/tikv (language: rust)

Issue:
In-memory Engine: Region merging leads to a load split ripple effect and causes latency spikes
## Bug Report

<!-- Thanks for your bug report! Don't worry if you can't fill out all the sections. -->

The latency spike in the following metrics is caused by a critical region (13511195) being evicted from IME due to region merge. The eviction causes a ripple effect, leading to load-based region splits and resulting in the server\_is\_busy error, which spikes latency.

![Image](https://github.com/user-attachments/assets/22c1a9be-1fc7-418d-bb5a-4bd2f34dbc46)

| Time | TiKV 1 | TiKV 2 |
|--------|--------|--------|
| 2025/01/07 11:49:19.140 | region 13511195 starts evict due to CommitMerge |  |
| 2025/01/07 11:49:19.493 | region 13511195 evict done |  |
| 2025/01/07 11:49:32.560 | region 13517458 load-splits from region 13511195 |  |
| 2025/01/07 11:49:43.672 | region 13517462 load-splits from region 13517458 |  |
| 2025/01/07 11:50:03.789 | region 13517472 load-splits from region 13517462 |  |
| 2025/01/07 11:50:15.858 | region 13517484 load-splits from region 13517472 |  |
| 2025/01/07 11:52:36.829 |  | region 13517484 became leader due to transfer leader |
| 2025/01/07 11:53:00.760 |  | region 13517523 load-splits from region 13517484 |
| 2025/01/07 11:53:12.829 |  | region 13517532 load-splits from region 13517523 |
| 2025/01/07 11:55:55.644 |  | ime auto load evict triggered |
| 2025/01/07 11:55:55.645 |  | region 13517484 / 13517532 are cached |


### What version of TiKV are you using?
<!-- You can run `tikv-server --version` -->

nightly 2025-01-08

### Steps to reproduce
<!-- If possible, provide a recipe for reproducing the error. A complete runnable program is good. -->

Run read-only high MVCC amplification workload.

### What did you expect?

No spike.

### What did happened?

Spike.


Discussion / hints:

Fix patch:
diff --git a/components/engine_traits/src/region_cache_engine.rs b/components/engine_traits/src/region_cache_engine.rs
index feeaad5ca9b..556dbcc410b 100644
--- a/components/engine_traits/src/region_cache_engine.rs
+++ b/components/engine_traits/src/region_cache_engine.rs
@@ -2,6 +2,8 @@
 
 use std::{
     fmt::{self, Debug},
+    future::Future,
+    pin::Pin,
     result,
 };
 
@@ -22,7 +24,9 @@ pub enum FailedReason {
     EpochNotMatch,
 }
 
-#[derive(Debug, PartialEq)]
+pub type OnEvictFinishedCallback =
+    Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
+
 pub enum RegionEvent {
     Split {
         source: CacheRegion,
@@ -35,6 +39,7 @@ pub enum RegionEvent {
     Eviction {
         region: CacheRegion,
         reason: EvictReason,
+        on_evict_finished: Option<OnEvictFinishedCallback>,
     },
     // range eviction triggered by delete_range
     // we should evict all cache regions that overlaps with this range
@@ -44,6 +49,44 @@ pub enum RegionEvent {
     },
 }
 
+impl fmt::Debug for RegionEvent {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            RegionEvent::Split {
+                source,
+                new_regions,
+            } => f
+                .debug_struct("RegionEvent::Split")
+                .field("source", source)
+                .field("new_regions", new_regions)
+                .finish(),
+            RegionEvent::TryLoad {
+                region,
+                for_manual_range,
+            } => f
+                .debug_struct("RegionEvent::TryLoad")
+                .field("region", region)
+                .field("for_manual_range", for_manual_range)
+                .finish(),
+            RegionEvent::Eviction {
+                region,
+                reason,
+                on_evict_finished,
+            } => f
+                .debug_struct("RegionEvent::Eviction")
+                .field("region", region)
+                .field("reason", reason)
+                .field("on_evict_finished", &on_evict_finished.is_some())
+                .finish(),
+            RegionEvent::EvictByRange { range, reason } => f
+                .debug_struct("RegionEvent::EvictByRange")
+                .field("range", range)
+                .field("reason", reason)
+                .finish(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum EvictReason {
     LoadFailed,
@@ -52,12 +95,14 @@ pub enum EvictReason {
     BecomeFollower,
     AutoEvict,
     DeleteRange,
+    PrepareMerge,
     Merge,
     Disabled,
     ApplySnapshot,
     Flashback,
     Manual,
-    PeerDestroy,
+    DestroyPeer,
+    IngestSST,
 }
 
 /// RegionCacheEngine works as a region cache caching some regions (in Memory or
diff --git a/components/hybrid_engine/src/observer/load_eviction.rs b/components/hybrid_engine/src/observer/load_eviction.rs
index 289288aebab..198003ee7a1 100644
--- a/components/hybrid_engine/src/observer/load_eviction.rs
+++ b/components/hybrid_engine/src/observer/load_eviction.rs
@@ -74,32 +74,44 @@ impl LoadEvictionObserver {
         // Evicting the cache for region splits is not worthwhile and may cause
         // performance regression due to frequent loading and evicting of
         // hot regions.
-        if apply.pending_handle_ssts.is_some()
-            || (state.modified_region.is_some()
-                && matches!(
-                    cmd.request.get_admin_request().get_cmd_type(),
-                    AdminCmdType::PrepareMerge
-                        | AdminCmdType::CommitMerge
-                        | AdminCmdType::RollbackMerge
-                ))
-        {
+        if apply.pending_handle_ssts.is_some() {
             let cache_region = CacheRegion::from_region(ctx.region());
             debug!(
-                "ime evict range due to apply commands";
+                "ime evict region due to ingest sst";
                 "region" => ?cache_region,
-                "is_ingest_sst" => apply.pending_handle_ssts.is_some(),
-                "admin_command" => ?cmd.request.get_admin_request().get_cmd_type(),
             );
-            self.evict_region(cache_region, EvictReason::Merge)
+            self.evict_region(cache_region, EvictReason::IngestSST)
+        }
+        let cmd_type = cmd.request.get_admin_request().get_cmd_type();
+        if let Some(modified_region) = &state.modified_region {
+            if matches!(cmd_type, AdminCmdType::PrepareMerge) {
+                let cache_region = CacheRegion::from_region(ctx.region());
+                debug!(
+                    "ime evict region due to apply commands";
+                    "region" => ?cache_region,
+                    "admin_command" => ?cmd_type
+                );
+                self.evict_region(cache_region, EvictReason::PrepareMerge)
+            } else if matches!(
+                cmd_type,
+                AdminCmdType::CommitMerge | AdminCmdType::RollbackMerge
+            ) {
+                let cache_region = CacheRegion::from_region(ctx.region());
+                debug!(
+                    "ime evict region followed by reloading due to apply commands";
+                    "region" => ?cache_region,
+                    "admin_command" => ?cmd_type
+                );
+                self.merge_region(cache_region.clone(), modified_region.clone())
+            }
         }
         // there are new_regions, this must be a split event.
         if !state.new_regions.is_empty() {
-            let cmd_type = cmd.request.get_admin_request().get_cmd_type();
             assert!(cmd_type == AdminCmdType::BatchSplit || cmd_type == AdminCmdType::Split);
             debug!(
                 "ime handle region split";
                 "region_id" => ctx.region().get_id(),
-                "admin_command" => ?cmd.request.get_admin_request().get_cmd_type(),
+                "admin_command" => ?cmd_type,
                 "region" => ?state.modified_region.as_ref().unwrap(),
                 "new_regions" => ?state.new_regions,
             );
@@ -122,6 +134,25 @@ impl LoadEvictionObserver {
         });
     }
 
+    // Merged regions are evicted first and will be reloaded immediately to
+    // minimize the impact of the merge operations on the cache hit ratio.
+    fn merge_region(&self, region: CacheRegion, modified_region: Region) {
+        let cache_engine = self.cache_engine.clone();
+        let modified_region = CacheRegion::from_region(&modified_region);
+        self.cache_engine.on_region_event(RegionEvent::Eviction {
+            region: region.clone(),
+            reason: EvictReason::Merge,
+            on_evict_finished: Some(Box::new(move || {
+                Box::pin(async move {
+                    cache_engine.on_region_event(RegionEvent::TryLoad {
+                        region: modified_region,
+                        for_manual_range: false,
+                    });
+                })
+            })),
+        });
+    }
+
     // Try to load region. It will be loaded if it's overlapped with maunal range
     fn try_load_region(&self, region: CacheRegion) {
         self.cache_engine.on_region_event(RegionEvent::TryLoad {
@@ -131,8 +162,11 @@ impl LoadEvictionObserver {
     }
 
     fn evict_region(&self, region: CacheRegion, reason: EvictReason) {
-        self.cache_engine
-            .on_region_event(RegionEvent::Eviction { region, reason });
+        self.cache_engine.on_region_event(RegionEvent::Eviction {
+            region,
+            reason,
+            on_evict_finished: None,
+        });
     }
 }
 
@@ -293,7 +327,8 @@ impl DestroyPeerObserver for LoadEvictionObserver {
         }
         self.cache_engine.on_region_event(RegionEvent::Eviction {
             region: CacheRegion::from_region(&region),
-            reason: EvictReason::PeerDestroy,
+            reason: EvictReason::DestroyPeer,
+            on_evict_finished: None,
         });
     }
 }
@@ -338,6 +373,40 @@ mod tests {
         request
     }
 
+    fn assert_eq_region_events(got: &RegionEvent, expected: &RegionEvent) {
+        match (got, expected) {
+            (
+                RegionEvent::TryLoad {
+                    region: got_region,
+                    for_manual_range: got_for_manual_range,
+                },
+                RegionEvent::TryLoad {
+                    region: expected_region,
+                    for_manual_range: expected_for_manual_range,
+                },
+            ) => {
+                assert_eq!(got_region, expected_region);
+                assert_eq!(got_for_manual_range, expected_for_manual_range);
+            }
+            (
+                RegionEvent::Eviction {
+                    region: got_region,
+                    reason: got_reason,
+                    on_evict_finished: _,
+                },
+                RegionEvent::Eviction {
+                    region: expected_region,
+                    reason: expected_reason,
+                    on_evict_finished: _,
+                },
+            ) => {
+                assert_eq!(got_region, expected_region);
+                assert_eq!(got_reason, expected_reason);
+            }
+            _ => panic!("unexpected region event"),
+        }
+    }
+
     #[test]
     fn test_do_not_evict_region_region_split() {
         let cache_engine = Arc::new(MockRegionCacheEngine::default());
@@ -361,7 +430,7 @@ mod tests {
         let response = RaftCmdResponse::default();
         let cmd = Cmd::new(0, 0, request, response);
 
-        // Must not evict range for region split.
+        // Must not evict region for region split.
         observer.post_exec_cmd(&mut ctx, &cmd, &RegionState::default(), &mut apply);
         assert!(&cache_engine.region_events.lock().unwrap().is_empty());
     }
@@ -400,9 +469,10 @@ mod tests {
         let cached_region = CacheRegion::from_region(&region);
         let expected = RegionEvent::Eviction {
             region: cached_region,
-            reason: EvictReason::Merge,
+            reason: EvictReason::IngestSST,
+            on_evict_finished: None,
         };
-        assert_eq!(&cache_engine.region_events.lock().unwrap()[0], &expected);
+        assert_eq_region_events(&cache_engine.region_events.lock().unwrap()[0], &expected);
     }
 
     #[test]
@@ -421,6 +491,6 @@ mod tests {
             region: cached_region,
             for_manual_range: true,
         };
-        assert_eq!(&cache_engine.region_events.lock().unwrap()[0], &expected);
+        assert_eq_region_events(&cache_engine.region_events.lock().unwrap()[0], &expected);
     }
 }
diff --git a/components/in_memory_engine/src/background.rs b/components/in_memory_engine/src/background.rs
index 59e27e5ef83..2ee71146538 100644
--- a/components/in_memory_engine/src/background.rs
+++ b/components/in_memory_engine/src/background.rs
@@ -9,8 +9,8 @@ use crossbeam::{
 };
 use engine_rocks::{RocksEngine, RocksSnapshot};
 use engine_traits::{
-    CacheRegion, EvictReason, IterOptions, Iterable, Iterator, MiscExt, RangeHintService,
-    SnapshotMiscExt, CF_DEFAULT, CF_WRITE, DATA_CFS,
+    CacheRegion, EvictReason, IterOptions, Iterable, Iterator, MiscExt, OnEvictFinishedCallback,
+    RangeHintService, SnapshotMiscExt, CF_DEFAULT, CF_WRITE, DATA_CFS,
 };
 use fail::fail_point;
 use keys::{origin_end_key, origin_key};
@@ -49,7 +49,7 @@ use crate::{
     region_label::{
         LabelRule, RegionLabelChangedCallback, RegionLabelRulesManager, RegionLabelServiceBuilder,
     },
-    region_manager::{AsyncFnOnce, CacheRegionMeta, RegionState},
+    region_manager::{CacheRegionMeta, RegionState},
     region_stats::{RegionStatsManager, DEFAULT_EVICT_MIN_DURATION},
     write_batch::RegionCacheWriteBatchEntry,
     InMemoryEngineConfig, RegionCacheMemoryEngine,
@@ -1150,7 +1150,7 @@ impl Runnable for BackgroundRunner {
 
                             let evict_fn = |evict_region: &CacheRegion,
                                             evict_reason: EvictReason,
-                                            cb: Option<Box<dyn AsyncFnOnce + Send + Sync>>|
+                                            cb: Option<OnEvictFinishedCallback>|
                              -> Vec<CacheRegion> {
                                 core.engine.region_manager.evict_region(
                                     evict_region,
@@ -2779,16 +2779,8 @@ pub mod tests {
         let k = format!("zk{:08}", 15).into_bytes();
         let region1 = CacheRegion::new(1, 0, DATA_MIN_KEY, k.clone());
         let region2 = CacheRegion::new(2, 0, k, DATA_MAX_KEY);
-        engine
-            .core
-            .region_manager()
-            .load_region(region1.clone())
-            .unwrap();
-        engine
-            .core
-            .region_manager()
-            .load_region(region2.clone())
-            .unwrap();
+        engine.load_region(region1.clone()).unwrap();
+        engine.load_region(region2.clone()).unwrap();
         engine.prepare_for_apply(&region1, false);
         engine.prepare_for_apply(&region2, false);
 
diff --git a/components/in_memory_engine/src/engine.rs b/components/in_memory_engine/src/engine.rs
index 872feec2706..d82474e26cb 100644
--- a/components/in_memory_engine/src/engine.rs
+++ b/components/in_memory_engine/src/engine.rs
@@ -14,8 +14,8 @@ use crossbeam_skiplist::{
 };
 use engine_rocks::RocksEngine;
 use engine_traits::{
-    CacheRegion, EvictReason, FailedReason, KvEngine, RegionCacheEngine, RegionCacheEngineExt,
-    RegionEvent, CF_DEFAULT, CF_LOCK, CF_WRITE, DATA_CFS,
+    CacheRegion, EvictReason, FailedReason, KvEngine, OnEvictFinishedCallback, RegionCacheEngine,
+    RegionCacheEngineExt, RegionEvent, CF_DEFAULT, CF_LOCK, CF_WRITE, DATA_CFS,
 };
 use fail::fail_point;
 use kvproto::metapb::Region;
@@ -31,9 +31,7 @@ use crate::{
     },
     memory_controller::MemoryController,
     read::RegionCacheSnapshot,
-    region_manager::{
-        AsyncFnOnce, LoadFailedReason, RegionCacheStatus, RegionManager, RegionState,
-    },
+    region_manager::{LoadFailedReason, RegionCacheStatus, RegionManager, RegionState},
     statistics::Statistics,
     InMemoryEngineConfig, InMemoryEngineContext,
 };
@@ -398,6 +396,7 @@ impl RegionCacheMemoryEngine {
     }
 
     pub fn load_region(&self, cache_region: CacheRegion) -> result::Result<(), LoadFailedReason> {
+        fail_point!("ime_on_load_region", |_| { Ok(()) });
         self.core.region_manager().load_region(cache_region)
     }
 
@@ -415,12 +414,12 @@ impl RegionCacheMemoryEngine {
         &self,
         region: &CacheRegion,
         evict_reason: EvictReason,
-        cb: Option<Box<dyn AsyncFnOnce + Send + Sync>>,
+        on_evict_finished: Option<OnEvictFinishedCallback>,
     ) {
-        let deletable_regions = self
-            .core
-            .region_manager
-            .evict_region(region, evict_reason, cb);
+        let deletable_regions =
+            self.core
+                .region_manager
+                .evict_region(region, evict_reason, on_evict_finished);
         if !deletable_regions.is_empty() {
             // The region can be deleted directly.
             if let Err(e) = self
@@ -545,8 +544,12 @@ impl RegionCacheEngine for RegionCacheMemoryEngine {
 impl RegionCacheEngineExt for RegionCacheMemoryEngine {
     fn on_region_event(&self, event: RegionEvent) {
         match event {
-            RegionEvent::Eviction { region, reason } => {
-                self.evict_region(&region, reason, None);
+            RegionEvent::Eviction {
+                region,
+                reason,
+                on_evict_finished,
+            } => {
+                self.evict_region(&region, reason, on_evict_finished);
             }
             RegionEvent::TryLoad {
                 region,
@@ -572,7 +575,7 @@ impl RegionCacheEngineExt for RegionCacheMemoryEngine {
                             );
                         }
                     }
-                } else if let Err(e) = self.core.region_manager().load_region(region.clone()) {
+                } else if let Err(e) = self.load_region(region.clone()) {
                     warn!(
                         "ime load region failed";
                         "error" => ?e,
@@ -829,6 +832,7 @@ pub mod tests {
         engine.on_region_event(RegionEvent::Eviction {
             region: new_regions[0].clone(),
             reason: EvictReason::AutoEvict,
+            on_evict_finished: None,
         });
 
         // trigger split again
@@ -855,6 +859,7 @@ pub mod tests {
         engine.on_region_event(RegionEvent::Eviction {
             region: cache_region,
             reason: EvictReason::AutoEvict,
+            on_evict_finished: None,
         });
 
         // engine should become inactive after all regions are evicted.
diff --git a/components/in_memory_engine/src/lib.rs b/components/in_memory_engine/src/lib.rs
index 351907ce24a..e4011f20235 100644
--- a/components/in_memory_engine/src/lib.rs
+++ b/components/in_memory_engine/src/lib.rs
@@ -5,7 +5,6 @@
 #![allow(internal_features)]
 #![feature(core_intrinsics)]
 #![feature(slice_pattern)]
-#![feature(trait_alias)]
 
 use std::sync::Arc;
 
diff --git a/components/in_memory_engine/src/metrics.rs b/components/in_memory_engine/src/metrics.rs
index 9b037c1c914..bca5821c68d 100644
--- a/components/in_memory_engine/src/metrics.rs
+++ b/components/in_memory_engine/src/metrics.rs
@@ -32,6 +32,7 @@ make_auto_flush_static_metric! {
     }
 
     pub label_enum EvictReasonType {
+        prepare_merge,
         merge,
         auto_evict,
         load_failed,
@@ -44,6 +45,7 @@ make_auto_flush_static_metric! {
         flashback,
         manual,
         destroy_peer,
+        ingest_sst,
     }
 
     pub label_enum OperationType {
@@ -246,6 +248,9 @@ pub(crate) fn observe_eviction_duration(secs: f64, evict_reason: EvictReason) {
         EvictReason::MemoryLimitReached => IN_MEMORY_ENGINE_EVICTION_DURATION_HISTOGRAM_STATIC
             .memory_limit_reached
             .observe(secs),
+        EvictReason::PrepareMerge => IN_MEMORY_ENGINE_EVICTION_DURATION_HISTOGRAM_STATIC
+            .prepare_merge
+            .observe(secs),
         EvictReason::Merge => IN_MEMORY_ENGINE_EVICTION_DURATION_HISTOGRAM_STATIC
             .merge
             .observe(secs),
@@ -261,9 +266,12 @@ pub(crate) fn observe_eviction_duration(secs: f64, evict_reason: EvictReason) {
         EvictReason::Manual => IN_MEMORY_ENGINE_EVICTION_DURATION_HISTOGRAM_STATIC
             .manual
             .observe(secs),
-        EvictReason::PeerDestroy => IN_MEMORY_ENGINE_EVICTION_DURATION_HISTOGRAM_STATIC
+        EvictReason::DestroyPeer => IN_MEMORY_ENGINE_EVICTION_DURATION_HISTOGRAM_STATIC
             .destroy_peer
             .observe(secs),
+        EvictReason::IngestSST => IN_MEMORY_ENGINE_EVICTION_DURATION_HISTOGRAM_STATIC
+            .ingest_sst
+            .observe(secs),
     }
 }
 
diff --git a/components/in_memory_engine/src/region_manager.rs b/components/in_memory_engine/src/region_manager.rs
index 15a0933efa0..8557c5cc063 100644
--- a/components/in_memory_engine/src/region_manager.rs
+++ b/components/in_memory_engine/src/region_manager.rs
@@ -7,8 +7,6 @@ use std::{
         Bound::{self, Excluded, Unbounded},
     },
     fmt::Debug,
-    future::Future,
-    pin::Pin,
     result,
     sync::{
         atomic::{AtomicBool, Ordering},
@@ -17,7 +15,7 @@ use std::{
 };
 
 use collections::HashMap;
-use engine_traits::{CacheRegion, EvictReason, FailedReason};
+use engine_traits::{CacheRegion, EvictReason, FailedReason, OnEvictFinishedCallback};
 use futures::executor::block_on;
 use parking_lot::RwLock;
 use strum::EnumCount;
@@ -25,8 +23,6 @@ use tikv_util::{info, smoother::Smoother, time::Instant, warn};
 
 use crate::{metrics::observe_eviction_duration, read::RegionCacheSnapshotMeta};
 
-pub(crate) trait AsyncFnOnce = FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send>>;
-
 #[derive(PartialEq, Eq, Debug, Clone, Copy, Default, Hash, EnumCount)]
 #[repr(usize)]
 pub enum RegionState {
@@ -154,7 +150,7 @@ struct EvictInfo {
     start: Instant,
     reason: EvictReason,
     // called when eviction finishes
-    cb: Option<Box<dyn AsyncFnOnce + Send + Sync>>,
+    on_evict_finished: Option<OnEvictFinishedCallback>,
 }
 
 impl Debug for EvictInfo {
@@ -240,7 +236,7 @@ impl CacheRegionMeta {
         &mut self,
         state: RegionState,
         reason: EvictReason,
-        cb: Option<Box<dyn AsyncFnOnce + Send + Sync>>,
+        on_evict_finished: Option<OnEvictFinishedCallback>,
     ) {
         use RegionState::*;
         assert_matches!(self.state, Loading | Active | LoadingCanceled);
@@ -249,7 +245,7 @@ impl CacheRegionMeta {
         self.evict_info = Some(EvictInfo {
             start: Instant::now_coarse(),
             reason,
-            cb,
+            on_evict_finished,
         });
     }
 
@@ -855,7 +851,7 @@ impl RegionManager {
         &self,
         evict_region: &CacheRegion,
         evict_reason: EvictReason,
-        mut cb: Option<Box<dyn AsyncFnOnce + Send + Sync>>,
+        mut on_evict_finished: Option<OnEvictFinishedCallback>,
     ) -> Vec<CacheRegion> {
         info!(
             "ime try to evict region";
@@ -883,7 +879,7 @@ impl RegionManager {
                 evict_reason,
                 &mut regions_map,
                 if rid == evict_region.id {
-                    cb.take()
+                    on_evict_finished.take()
                 } else {
                     None
                 },
@@ -901,7 +897,7 @@ impl RegionManager {
         evict_region: &CacheRegion,
         evict_reason: EvictReason,
         regions_map: &mut RegionMetaMap,
-        cb: Option<Box<dyn AsyncFnOnce + Send + Sync>>,
+        on_evict_finished: Option<OnEvictFinishedCallback>,
     ) -> Option<CacheRegion> {
         let meta = regions_map.mut_region_meta(id).unwrap();
         let prev_state = meta.state;
@@ -928,7 +924,7 @@ impl RegionManager {
         }
 
         if prev_state == RegionState::Active {
-            meta.mark_evict(RegionState::PendingEvict, evict_reason, cb);
+            meta.mark_evict(RegionState::PendingEvict, evict_reason, on_evict_finished);
         } else {
             meta.set_state(RegionState::LoadingCanceled)
         };
@@ -969,7 +965,7 @@ impl RegionManager {
                     evict_info.start.saturating_elapsed_secs(),
                     evict_info.reason,
                 );
-                if let Some(cb) = evict_info.cb {
+                if let Some(cb) = evict_info.on_evict_finished {
                     cbs.push(cb);
                 }
 
@@ -1005,7 +1001,7 @@ impl RegionManager {
         }
     }
 
-    pub fn load_region(&self, cache_region: CacheRegion) -> Result<(), LoadFailedReason> {
+    pub(crate) fn load_region(&self, cache_region: CacheRegion) -> Result<(), LoadFailedReason> {
         self.regions_map.write().load_region(cache_region)
     }
 
diff --git a/components/in_memory_engine/src/region_stats.rs b/components/in_memory_engine/src/region_stats.rs
index 96028470da0..5c5ffdc6ad9 100644
--- a/components/in_memory_engine/src/region_stats.rs
+++ b/components/in_memory_engine/src/region_stats.rs
@@ -11,7 +11,7 @@ use std::{
 
 use collections::HashMap;
 use crossbeam::sync::ShardedLock;
-use engine_traits::{CacheRegion, EvictReason};
+use engine_traits::{CacheRegion, EvictReason, OnEvictFinishedCallback};
 use kvproto::metapb::Region;
 use parking_lot::Mutex;
 use pd_client::RegionStat;
@@ -21,9 +21,8 @@ use tikv_util::{config::VersionTrack, info, worker::Scheduler};
 use tokio::sync::mpsc;
 
 use crate::{
-    memory_controller::MemoryController,
-    region_manager::{AsyncFnOnce, CopRequestsSma},
-    BackgroundTask, InMemoryEngineConfig,
+    memory_controller::MemoryController, region_manager::CopRequestsSma, BackgroundTask,
+    InMemoryEngineConfig,
 };
 
 /// Do not evict a region if has been cached for less than this duration.
@@ -319,11 +318,7 @@ impl RegionStatsManager {
         cached_regions: HashMap<u64, Arc<StdMutex<CopRequestsSma>>>,
         memory_controller: &MemoryController,
     ) where
-        F: FnMut(
-            &CacheRegion,
-            EvictReason,
-            Option<Box<dyn AsyncFnOnce + Send + Sync>>,
-        ) -> Vec<CacheRegion>,
+        F: FnMut(&CacheRegion, EvictReason, Option<OnEvictFinishedCallback>) -> Vec<CacheRegion>,
     {
         let regions_activity = self.get_cached_region_stats(&cached_regions);
         if regions_activity.is_empty() {
@@ -716,7 +711,7 @@ pub mod tests {
         let cbs2 = cbs.clone();
         let evict_fn = move |evict_region: &CacheRegion,
                              _: EvictReason,
-                             cb: Option<Box<dyn AsyncFnOnce + Send + Sync>>|
+                             cb: Option<OnEvictFinishedCallback>|
               -> Vec<CacheRegion> {
             evicted_regions2.lock().push(evict_region.id);
             cbs2.lock().push(cb.unwrap());
diff --git a/components/in_memory_engine/src/write_batch.rs b/components/in_memory_engine/src/write_batch.rs
index e7ddfc1aa45..f2e4b6aeed9 100644
--- a/components/in_memory_engine/src/write_batch.rs
+++ b/components/in_memory_engine/src/write_batch.rs
@@ -1023,7 +1023,7 @@ mod tests {
 
         let r1 = CacheRegion::new(1, 0, b"zk00", b"zk10");
 
-        engine.core().region_manager().load_region(r1).unwrap();
+        engine.load_region(r1).unwrap();
 
         // load a region with a newer epoch and small range, should trigger replace.
         let mut r_new = Region::default();
@@ -1057,11 +1057,7 @@ mod tests {
         let mut wb = RegionCacheWriteBatch::from(&engine);
         wb.prepare_for_region(&r);
 
-        engine
-            .core()
-            .region_manager()
-            .load_region(cache_region.clone())
-            .unwrap();
+        engine.load_region(cache_region.clone()).unwrap();
         wb.prepare_for_region(&r);
         wb.put(b"k1", b"val1").unwrap();
         wb.put(b"k2", b"val2").unwrap();


Test patch:
diff --git a/tests/failpoints/cases/test_in_memory_engine.rs b/tests/failpoints/cases/test_in_memory_engine.rs
index aeb2b04454f..40dbde4235b 100644
--- a/tests/failpoints/cases/test_in_memory_engine.rs
+++ b/tests/failpoints/cases/test_in_memory_engine.rs
@@ -216,18 +216,12 @@ fn test_load() {
         let r2 = cluster.get_region(&split_keys[1]);
         let region_cache_engine = cluster.sim.rl().get_region_cache_engine(1);
         region_cache_engine
-            .core()
-            .region_manager()
             .load_region(CacheRegion::from_region(&r))
             .unwrap();
         region_cache_engine
-            .core()
-            .region_manager()
             .load_region(CacheRegion::from_region(&r1))
             .unwrap();
         region_cache_engine
-            .core()
-            .region_manager()
             .load_region(CacheRegion::from_region(&r2))
             .unwrap();
     }
@@ -281,8 +275,6 @@ fn test_load_with_split() {
         // it correctly.
         let cache_range = new_region(1, "", "");
         region_cache_engine
-            .core()
-            .region_manager()
             .load_region(CacheRegion::from_region(&cache_range))
             .unwrap();
     }
@@ -370,8 +362,6 @@ fn test_load_with_split2() {
     // try to load a region with old epoch and bigger range,
     // it should be updated to the real region range.
     region_cache_engine
-        .core()
-        .region_manager()
         .load_region(CacheRegion::new(r_split.id, 0, DATA_MIN_KEY, DATA_MAX_KEY))
         .unwrap();
 
@@ -437,11 +427,7 @@ fn test_load_with_eviction() {
     // it correctly.
     let r = cluster.get_region(b"");
     let cache_range = CacheRegion::from_region(&r);
-    region_cache_engine
-        .core()
-        .region_manager()
-        .load_region(cache_range)
-        .unwrap();
+    region_cache_engine.load_region(cache_range).unwrap();
 
     let product1 = ProductTable::new();
     let product2 = ProductTable::new();
@@ -522,8 +508,7 @@ fn test_evictions_after_transfer_leader() {
         .unwrap_err();
 }
 
-#[test]
-fn test_eviction_after_merge() {
+fn test_commit_merge(test_reload: bool) {
     let mut cluster = new_server_cluster_with_hybrid_engine(0, 1);
     cluster.run();
     let r = cluster.get_region(b"");
@@ -551,11 +536,163 @@ fn test_eviction_after_merge() {
         .snapshot(range2.clone(), 100, 100)
         .unwrap();
 
+    if !test_reload {
+        // Prevent reload.
+        fail::cfg("ime_on_load_region", "return").unwrap();
+    }
+
     let pd_client = Arc::clone(&cluster.pd_client);
     pd_client.must_merge(r.get_id(), r2.get_id());
 
     region_cache_engine.snapshot(range1, 100, 100).unwrap_err();
     region_cache_engine.snapshot(range2, 100, 100).unwrap_err();
+
+    if test_reload {
+        // Get the latest region `r2`.
+        let r2 = cluster.get_region(b"key1");
+        let range2 = CacheRegion::from_region(&r2);
+        // `r2` should be reloaded after merge completed.
+        eventually(Duration::from_millis(100), Duration::from_secs(5), || {
+            region_cache_engine
+                .snapshot(range2.clone(), 100, 100)
+                .is_ok()
+        });
+    }
+}
+
+#[test]
+fn test_eviction_after_merge() {
+    let test_reload = false;
+    test_commit_merge(test_reload);
+}
+
+// IME should reload cached region automatically after commit merge to
+// minimize the impact to coprocessor requests.
+#[test]
+fn test_reload_after_merge() {
+    let test_reload = true;
+    test_commit_merge(test_reload);
+}
+
+fn test_rollback_merge(test_reload: bool) {
+    let mut cluster = new_server_cluster_with_hybrid_engine(0, 3);
+    configure_for_merge(&mut cluster.cfg);
+    let pd_client = Arc::clone(&cluster.pd_client);
+    pd_client.disable_default_operator();
+
+    cluster.run_conf_change();
+
+    let region = pd_client.get_region(b"k1").unwrap();
+    cluster.must_split(&region, b"k2");
+    let left = pd_client.get_region(b"k1").unwrap();
+    let right = pd_client.get_region(b"k2").unwrap();
+
+    pd_client.must_add_peer(left.get_id(), new_peer(2, 2));
+    pd_client.must_add_peer(right.get_id(), new_peer(2, 4));
+
+    cluster.must_put(b"k1", b"v1");
+    cluster.must_put(b"k3", b"v3");
+
+    let mut region = pd_client.get_region(b"k1").unwrap();
+    let target_region = pd_client.get_region(b"k3").unwrap();
+
+    let schedule_merge_fp = "on_schedule_merge";
+    fail::cfg(schedule_merge_fp, "return()").unwrap();
+
+    let (tx, rx) = sync_channel(1);
+    fail::cfg_callback("on_apply_res_prepare_merge", move || {
+        tx.send(()).unwrap();
+    })
+    .unwrap();
+
+    cluster.merge_region(region.get_id(), target_region.get_id(), Callback::None);
+    // PrepareMerge is applied.
+    rx.recv().unwrap();
+
+    let leader = cluster.leader_of_region(left.get_id()).unwrap();
+
+    let region_cache_engine = cluster.sim.rl().get_region_cache_engine(leader.store_id);
+
+    // After prepare merge, version becomes 2 + 1 = 3;
+    region.mut_region_epoch().set_version(3);
+    // load region after PrepareMerge.
+    let cache_region = CacheRegion::from_region(&region);
+    region_cache_engine
+        .core()
+        .region_manager()
+        .new_region(cache_region.clone());
+    eventually(Duration::from_millis(100), Duration::from_secs(5), || {
+        region_cache_engine
+            .snapshot(cache_region.clone(), 100, 100)
+            .is_ok()
+    });
+    if !test_reload {
+        // Prevent reload after rollback merge.
+        fail::cfg("ime_on_load_region", "return").unwrap();
+    }
+
+    // Add a peer to trigger rollback.
+    pd_client.must_add_peer(right.get_id(), new_peer(3, 5));
+    cluster.must_put(b"k4", b"v4");
+    must_get_equal(&cluster.get_engine(3), b"k4", b"v4");
+
+    let mut region = pd_client.get_region(b"k1").unwrap();
+    // After split and prepare_merge, version becomes 1 + 2 = 3;
+    assert_eq!(region.get_region_epoch().get_version(), 3);
+    // After ConfChange and prepare_merge, conf version becomes 1 + 2 = 3;
+    assert_eq!(region.get_region_epoch().get_conf_ver(), 3);
+    fail::remove(schedule_merge_fp);
+    // Wait till rollback.
+    cluster.must_put(b"k11", b"v11");
+
+    // After rollback, version becomes 3 + 1 = 4;
+    region.mut_region_epoch().set_version(4);
+    for i in 1..3 {
+        must_get_equal(&cluster.get_engine(i), b"k11", b"v11");
+        let state = cluster.region_local_state(region.get_id(), i);
+        assert_eq!(state.get_state(), PeerState::Normal);
+        assert_eq!(*state.get_region(), region);
+    }
+
+    if !test_reload {
+        // After rollback, IME cached region is evicted.
+        eventually(Duration::from_millis(100), Duration::from_secs(5), || {
+            let region_map = region_cache_engine
+                .core()
+                .region_manager()
+                .regions_map()
+                .read();
+            region_map.regions().is_empty()
+        });
+    } else {
+        // Get the latest region.
+        let region = pd_client.get_region(b"k1").unwrap();
+        let cache_region = CacheRegion::from_region(&region);
+        // Region should be reloaded after merge completed.
+        eventually(Duration::from_millis(100), Duration::from_secs(5), || {
+            region_cache_engine
+                .snapshot(cache_region.clone(), 100, 100)
+                .is_ok()
+        });
+    }
+}
+
+// IME should also handle RollbackMerge event, we also try to evict the region
+// on merge rollback for simplicity. If region is loaded after PrepareMerge and
+// the merge is rollbacked, IME should track this rollback because it will also
+// change epoch version.
+#[test]
+fn test_eviction_after_rollback_merge() {
+    let test_reload = false;
+    test_rollback_merge(test_reload);
+}
+
+// IME should reload cached region automatically after merge rollback to
+// minimize the impact to coprocessor requests.
+#[test]
+fn test_reload_after_rollback_merge() {
+    let test_reload = true;
+    test_rollback_merge(test_reload);
 }
 
 #[test]
@@ -740,8 +877,6 @@ fn test_delete_range() {
             // it correctly.
             let cache_range = new_region(1, "", "");
             region_cache_engine
-                .core()
-                .region_manager()
                 .load_region(CacheRegion::from_region(&cache_range))
                 .unwrap();
         }
@@ -808,8 +943,6 @@ fn test_evict_on_flashback() {
     {
         let region_cache_engine = cluster.sim.rl().get_region_cache_engine(1);
         region_cache_engine
-            .core()
-            .region_manager()
             .load_region(CacheRegion::from_region(&r))
             .unwrap();
     }
@@ -853,8 +986,6 @@ fn test_load_during_flashback() {
     {
         let region_cache_engine = cluster.sim.rl().get_region_cache_engine(1);
         region_cache_engine
-            .core()
-            .region_manager()
             .load_region(CacheRegion::from_region(&r))
             .unwrap();
     }
@@ -888,13 +1019,9 @@ fn test_apply_prepared_but_not_write() {
     // load both regions at store 1
     let region_cache_engine = cluster.sim.rl().get_region_cache_engine(1);
     region_cache_engine
-        .core()
-        .region_manager()
         .load_region(CacheRegion::from_region(&r1))
         .unwrap();
     region_cache_engine
-        .core()
-        .region_manager()
         .load_region(CacheRegion::from_region(&r2))
         .unwrap();
 
@@ -989,11 +1116,7 @@ fn test_eviction_when_destroy_peer() {
     {
         let region_cache_engine = cluster.sim.rl().get_region_cache_engine(1);
         let cache_region = CacheRegion::from_region(&r);
-        region_cache_engine
-            .core()
-            .region_manager()
-            .load_region(cache_region)
-            .unwrap();
+        region_cache_engine.load_region(cache_region).unwrap();
     }
 
     must_copr_load_data(&mut cluster, &t1, 1);
@@ -1061,99 +1184,6 @@ fn test_eviction_when_destroy_uninitialized_peer() {
     cluster.must_region_exist(region.get_id(), 2);
 }
 
-// IME should also handle RollbackMerge event, we also try to evict the region
-// on merge rollback for simplicity. If region is loaded after PrepareMerge and
-// the merge is rollbacked, IME should track this rollback because it will also
-// change epoch version.
-#[test]
-fn test_region_rollback_merge() {
-    let mut cluster = new_server_cluster_with_hybrid_engine(0, 3);
-    configure_for_merge(&mut cluster.cfg);
-    let pd_client = Arc::clone(&cluster.pd_client);
-    pd_client.disable_default_operator();
-
-    cluster.run_conf_change();
-
-    let region = pd_client.get_region(b"k1").unwrap();
-    cluster.must_split(&region, b"k2");
-    let left = pd_client.get_region(b"k1").unwrap();
-    let right = pd_client.get_region(b"k2").unwrap();
-
-    pd_client.must_add_peer(left.get_id(), new_peer(2, 2));
-    pd_client.must_add_peer(right.get_id(), new_peer(2, 4));
-
-    cluster.must_put(b"k1", b"v1");
-    cluster.must_put(b"k3", b"v3");
-
-    let mut region = pd_client.get_region(b"k1").unwrap();
-    let target_region = pd_client.get_region(b"k3").unwrap();
-
-    let schedule_merge_fp = "on_schedule_merge";
-    fail::cfg(schedule_merge_fp, "return()").unwrap();
-
-    let (tx, rx) = sync_channel(1);
-    fail::cfg_callback("on_apply_res_prepare_merge", move || {
-        tx.send(()).unwrap();
-    })
-    .unwrap();
-
-    cluster.merge_region(region.get_id(), target_region.get_id(), Callback::None);
-    // PrepareMerge is applied.
-    rx.recv().unwrap();
-
-    let leader = cluster.leader_of_region(left.get_id()).unwrap();
-
-    let region_cache_engine = cluster.sim.rl().get_region_cache_engine(leader.store_id);
-
-    // After prepare merge, version becomes 2 + 1 = 3;
-    region.mut_region_epoch().set_version(3);
-    // load region after PrepareMerge.
-    {
-        let cache_region = CacheRegion::from_region(&region);
-        region_cache_engine
-            .core()
-            .region_manager()
-            .new_region(cache_region);
-    }
-
-    // Add a peer to trigger rollback.
-    pd_client.must_add_peer(right.get_id(), new_peer(3, 5));
-    cluster.must_put(b"k4", b"v4");
-    must_get_equal(&cluster.get_engine(3), b"k4", b"v4");
-
-    let mut region = pd_client.get_region(b"k1").unwrap();
-    // After split and prepare_merge, version becomes 1 + 2 = 3;
-    assert_eq!(region.get_region_epoch().get_version(), 3);
-    // After ConfChange and prepare_merge, conf version becomes 1 + 2 = 3;
-    assert_eq!(region.get_region_epoch().get_conf_ver(), 3);
-    fail::remove(schedule_merge_fp);
-    // Wait till rollback.
-    cluster.must_put(b"k11", b"v11");
-
-    // After rollback, version becomes 3 + 1 = 4;
-    region.mut_region_epoch().set_version(4);
-    for i in 1..3 {
-        must_get_equal(&cluster.get_engine(i), b"k11", b"v11");
-        let state = cluster.region_local_state(region.get_id(), i);
-        assert_eq!(state.get_state(), PeerState::Normal);
-        assert_eq!(*state.get_region(), region);
-    }
-
-    // after rollback, IME cached region is evicted.
-    test_util::eventually(
-        Duration::from_millis(10),
-        Duration::from_millis(1000),
-        || {
-            let region_map = region_cache_engine
-                .core()
-                .region_manager()
-                .regions_map()
-                .read();
-            region_map.regions().is_empty()
-        },
-    );
-}
-
 // IME must not panic when pre-load an uninitialized peer.
 #[test]
 fn test_transfer_leader_pre_load_uninitialized_peer() {