KVStore: Fix spurious region overlap when two region are both applying snapshots

dbms/src/Storages/KVStore/MultiRaft/ApplySnapshot.cpp

@@ -80,8 +80,8 @@ void KVStore::checkAndApplyPreHandledSnapshot(const RegionPtrWrap & new_region,
             // engine may delete data unsafely.
             auto region_lock = region_manager.genRegionTaskLock(old_region->id());
             old_region->setStateApplying();
-            tmt.getRegionTable().tryWriteBlockByRegion(old_region);
-            tryFlushRegionCacheInStorage(tmt, *old_region, log);
+            // It is not worthy to call `tryWriteBlockByRegion` and `tryFlushRegionCacheInStorage` here,
+            // even if the written data is useful, it could be overwritten later in `onSnapshot`.
             persistRegion(*old_region, region_lock, PersistRegionReason::ApplySnapshotPrevRegion, "");
         }
     }
@@ -95,23 +95,49 @@ void KVStore::checkAndApplyPreHandledSnapshot(const RegionPtrWrap & new_region,
             if (overlapped_region.first != region_id)
             {
                 auto state = getProxyHelper()->getRegionLocalState(overlapped_region.first);
-                if (state.state() != raft_serverpb::PeerState::Tombstone)
+                auto extra_msg = fmt::format("state={}, tiflash_state={}, new_region_state={}", state.ShortDebugString(),
+                    overlapped_region.second->mutMeta().getRegionState().getBase().ShortDebugString(),
+                    new_region->mutMeta().getRegionState().getBase().ShortDebugString());
+                if (state.state() == raft_serverpb::PeerState::Tombstone)
                 {
-                    throw Exception(
-                        ErrorCodes::LOGICAL_ERROR,
-                        "range of region_id={} is overlapped with region_id={}, state: {}",
+                    LOG_INFO(
+                        log,
+                        "range of region_id={} is overlapped with `Tombstone` region_id={}, {}",
                         region_id,
                         overlapped_region.first,
-                        state.ShortDebugString());
+                        extra_msg);
+                    handleDestroy(overlapped_region.first, tmt, task_lock);
+                }
+                else if (state.state() == raft_serverpb::PeerState::Applying)
+                {
+                    auto r = RegionRangeKeys {
+                        TiKVKey::copyFrom(state.region().start_key()),
+                        TiKVKey::copyFrom(state.region().end_key()),
+                    };
+                    if(RegionsRangeIndex::isRangeOverlapped(new_range->comparableKeys(), r.comparableKeys())) {
+                        throw Exception(
+                            ErrorCodes::LOGICAL_ERROR,
+                            "range of region_id={} is overlapped with `Applying` region_id={}, {}",
+                            region_id,
+                            overlapped_region.first,
+                            extra_msg);
+                    } else {
+                        LOG_INFO(
+                            log,
+                            "range of region_id={} is overlapped with `Applying` region_id={}, {}",
+                            region_id,
+                            overlapped_region.first,
+                            extra_msg);
+                    }
                 }
                 else
                 {
-                    LOG_INFO(
-                        log,
-                        "range of region_id={} is overlapped with `Tombstone` region_id={}",
+                    throw Exception(
+                        ErrorCodes::LOGICAL_ERROR,
+                        "range of region_id={} is overlapped with region_id={}, {}",
                         region_id,
-                        overlapped_region.first);
-                    handleDestroy(overlapped_region.first, tmt, task_lock);
+                        overlapped_region.first,
+                        extra_msg);
                 }
             }
         }
@@ -196,9 +222,9 @@ void KVStore::onSnapshot(
                             new_key_range.toDebugString(),
                             keyspace_id,
                             table_id);
-                        dm_storage->deleteRange(old_key_range, context.getSettingsRef());
+                        dm_storage->deleteRange(new_key_range, context.getSettingsRef());
                         // We must flush the deletion to the disk here, because we only flush new range when persisting this region later.
-                        dm_storage->flushCache(context, old_key_range, /*try_until_succeed*/ true);
+                        dm_storage->flushCache(context, new_key_range, /*try_until_succeed*/ true);
                     }
                 }
                 if constexpr (std::is_same_v<RegionPtrWrap, RegionPtrWithSnapshotFiles>)

dbms/src/Storages/KVStore/MultiRaft/RegionsRangeIndex.cpp

@@ -171,4 +171,17 @@ RegionsRangeIndex::RootMap::iterator RegionsRangeIndex::split(const TiKVRangeKey
         return do_split(begin_it, new_start);
 }
 
+
+bool RegionsRangeIndex::isRangeOverlapped(const RegionRangeKeys::RegionRange & a, const RegionRangeKeys::RegionRange & b) {
+    auto start = a.first.compare(b.first);
+    if(start == 0) {
+        return true;
+    } else if (start < 0) {
+        return a.second.compare(b.second) >= 0;
+    } else {
+        return b.second.compare(a.second) >= 0;
+    }
+}
+
+
 } // namespace DB

dbms/src/Storages/KVStore/MultiRaft/RegionsRangeIndex.h

@@ -71,6 +71,8 @@ class RegionsRangeIndex : private boost::noncopyable
     // TODO this friend class decl is not working.
     RootMap::iterator split(const TiKVRangeKey & new_start);
 
+    static bool isRangeOverlapped(const RegionRangeKeys::RegionRange &, const RegionRangeKeys::RegionRange &);
+
 private:
     friend class ::DB::tests::KVStoreTestBase;
     friend class ::DB::tests::RegionKVStoreOldTest;

dbms/src/Storages/KVStore/tests/gtest_new_kvstore.cpp

@@ -1164,4 +1164,42 @@ try
 }
 CATCH
 
+TEST_F(RegionKVStoreTest, ApplyShrinkedRegion)
+try
+{
+    auto & ctx = TiFlashTestEnv::getGlobalContext();
+    ASSERT_NE(proxy_helper->sst_reader_interfaces.fn_key, nullptr);
+    UInt64 region_id = 1;
+    TableID table_id;
+    {
+        initStorages();
+        KVStore & kvs = getKVS();
+        table_id = proxy_instance->bootstrapTable(ctx, kvs, ctx.getTMTContext());
+        LOG_INFO(&Poco::Logger::get("Test"), "generated table_id {}", table_id);
+        proxy_instance->bootstrapWithRegion(kvs, ctx.getTMTContext(), region_id, std::nullopt);
+        auto kvr1 = kvs.getRegion(region_id);
+        auto r1 = proxy_instance->getRegion(region_id);
+        {
+            // Multiple files
+            MockSSTReader::getMockSSTData().clear();
+            MockSSTGenerator default_cf{902, 800, ColumnFamilyType::Default};
+            default_cf.insert(1, "v1");
+            default_cf.finish_file();
+            default_cf.insert(2, "v2");
+            default_cf.finish_file();
+            default_cf.insert(3, "v3");
+            default_cf.insert(4, "v4");
+            default_cf.finish_file();
+            default_cf.insert(5, "v5");
+            default_cf.insert(6, "v6");
+            default_cf.finish_file();
+            default_cf.insert(7, "v7");
+            default_cf.finish_file();
+            default_cf.freeze();
+            validateSSTGeneration(kvs, proxy_instance, region_id, default_cf, ColumnFamilyType::Default, 5, 7);
+        }
+    }
+}
+CATCH
+
 } // namespace DB::tests