Handle snapshot in parallel if there is only one big region in raftstore v2

dbms/src/Common/CurrentMetrics.cpp

@@ -47,6 +47,7 @@
     M(DT_DeltaMergeTotalRows)                   \
     M(DT_DeltaIndexCacheSize)                   \
     M(RaftNumSnapshotsPendingApply)             \
+    M(RaftNumPrehandlingSubTasks)               \
     M(RateLimiterPendingWriteRequest)           \
     M(DT_SegmentReadTasks)                      \
     M(DT_SnapshotOfRead)                        \

dbms/src/Common/FailPoint.cpp

@@ -102,6 +102,8 @@ namespace DB
     M(exception_after_large_write_exceed)                    \
     M(proactive_flush_force_set_type)                        \
     M(exception_when_fetch_disagg_pages)                     \
+    M(force_set_parallel_prehandle_threshold)                \
+    M(force_raise_prehandle_exception)                       \
     M(force_agg_on_partial_block)
 
 #define APPLY_FOR_PAUSEABLE_FAILPOINTS_ONCE(M) \

dbms/src/Debug/MockRaftStoreProxy.cpp

@@ -111,6 +111,12 @@ void fn_gc_rust_ptr(RawVoidPtr ptr, RawRustPtrType type_)
     case RawObjType::MockAsyncWaker:
         delete reinterpret_cast<MockAsyncWaker *>(ptr);
         break;
+    case RawObjType::MockString:
+        delete reinterpret_cast<std::string *>(ptr);
+        break;
+    case RawObjType::MockVecOfString:
+        delete reinterpret_cast<RustStrWithViewVecInner *>(ptr);
+        break;
     }
 }
 
@@ -167,12 +173,22 @@ void fn_notify_compact_log(
     }
 }
 
-RaftstoreVer fn_get_cluster_raftstore_version(RaftStoreProxyPtr ptr, uint8_t, int64_t)
+static RaftstoreVer fn_get_cluster_raftstore_version(RaftStoreProxyPtr ptr, uint8_t, int64_t)
 {
     auto & x = as_ref(ptr);
     return x.cluster_ver;
 }
 
+// Must call `RustGcHelper` to gc the returned pointer in the end.
+static RustStrWithView fn_get_config_json(RaftStoreProxyPtr, uint64_t)
+{
+    auto * s = new std::string(R"({"raftstore":{"snap-handle-pool-size":4}})");
+    GCMonitor::instance().add(RawObjType::MockString, 1);
+    return RustStrWithView{
+        .buff = cppStringAsBuff(*s),
+        .inner = RawRustPtr{.ptr = s, .type = static_cast<RawRustPtrType>(RawObjType::MockString)}};
+}
+
 TiFlashRaftProxyHelper MockRaftStoreProxy::SetRaftStoreProxyFFIHelper(RaftStoreProxyPtr proxy_ptr)
 {
     TiFlashRaftProxyHelper res{};
@@ -184,6 +200,7 @@ TiFlashRaftProxyHelper MockRaftStoreProxy::SetRaftStoreProxyFFIHelper(RaftStoreP
     res.fn_get_region_local_state = fn_get_region_local_state;
     res.fn_notify_compact_log = fn_notify_compact_log;
     res.fn_get_cluster_raftstore_version = fn_get_cluster_raftstore_version;
+    res.fn_get_config_json = fn_get_config_json;
     {
         // make sure such function pointer will be set at most once.
         static std::once_flag flag;
@@ -897,6 +914,7 @@ void MockRaftStoreProxy::Cf::finish_file(SSTFormatKind kind)
     {
         kv_list.emplace_back(kv.first, kv.second);
     }
+    std::sort(kv_list.begin(), kv_list.end(), [](const auto & a, const auto & b) { return a.first < b.first; });
     auto & mmp = MockSSTReader::getMockSSTData();
     mmp[MockSSTReader::Key{region_id_str, type}] = std::move(kv_list);
     kvs.clear();
@@ -941,7 +959,7 @@ void MockRaftStoreProxy::Cf::insert_raw(std::string key, std::string val)
     kvs.emplace_back(std::move(key), std::move(val));
 }
 
-RegionPtr MockRaftStoreProxy::snapshot(
+std::tuple<RegionPtr, PrehandleResult> MockRaftStoreProxy::snapshot(
     KVStore & kvs,
     TMTContext & tmt,
     UInt64 region_id,
@@ -978,19 +996,19 @@ RegionPtr MockRaftStoreProxy::snapshot(
     SSTViewVec snaps{ssts.data(), ssts.size()};
     auto prehandle_result = kvs.preHandleSnapshotToFiles(new_kv_region, snaps, index, term, deadline_index, tmt);
 
-    auto rg = RegionPtrWithSnapshotFiles{new_kv_region, std::move(prehandle_result.ingest_ids)};
+    auto rg = RegionPtrWithSnapshotFiles{new_kv_region, std::vector(prehandle_result.ingest_ids)};
     if (cancel_after_prehandle)
     {
         kvs.releasePreHandledSnapshot(rg, tmt);
-        return kvs.getRegion(region_id);
+        return std::make_tuple(kvs.getRegion(region_id), prehandle_result);
     }
     kvs.checkAndApplyPreHandledSnapshot<RegionPtrWithSnapshotFiles>(rg, tmt);
     region->updateAppliedIndex(index);
     // PreHandledSnapshotWithFiles will do that, however preHandleSnapshotToFiles will not.
     new_kv_region->setApplied(index, term);
 
     // Region changes during applying snapshot, must re-get.
-    return kvs.getRegion(region_id);
+    return std::make_tuple(kvs.getRegion(region_id), prehandle_result);
 }
 
 TableID MockRaftStoreProxy::bootstrapTable(Context & ctx, KVStore & kvs, TMTContext & tmt, bool drop_at_first)
@@ -1041,7 +1059,10 @@ bool GCMonitor::checkClean()
     for (auto && d : data)
     {
         if (d.second)
+        {
+            LOG_INFO(&Poco::Logger::get("GCMonitor"), "checkClean {} has {}", magic_enum::enum_name(d.first), d.second);
             return false;
+        }
     }
     return true;
 }

dbms/src/Debug/MockRaftStoreProxy.h

@@ -260,7 +260,7 @@ struct MockRaftStoreProxy : MutexLockWrap
         bool freezed;
     };
 
-    RegionPtr snapshot(
+    std::tuple<RegionPtr, PrehandleResult> snapshot(
         KVStore & kvs,
         TMTContext & tmt,
         UInt64 region_id,
@@ -312,6 +312,8 @@ enum class RawObjType : uint32_t
     None,
     MockReadIndexTask,
     MockAsyncWaker,
+    MockString,
+    MockVecOfString
 };
 
 struct GCMonitor
@@ -344,4 +346,21 @@ std::vector<std::pair<std::string, std::string>> regionRangeToEncodeKeys(Types &
     return ranges_str;
 }
 
+struct RustStrWithViewVecInner
+{
+    std::vector<std::string> * vec;
+    BaseBuffView * buffs;
+    ~RustStrWithViewVecInner()
+    {
+        delete vec;
+        delete[] buffs;
+    }
+};
+
+inline BaseBuffView * createBaseBuffViewArray(size_t len)
+{
+    void * raw_memory = operator new[](len * sizeof(BaseBuffView));
+    BaseBuffView * ptr = static_cast<BaseBuffView *>(raw_memory);
+    return ptr;
+}
 } // namespace DB

dbms/src/Debug/MockSSTReader.cpp

@@ -15,6 +15,9 @@
 #include "MockSSTReader.h"
 
 #include <Common/Exception.h>
+#include <Debug/MockRaftStoreProxy.h>
+
+#include <ext/scope_guard.h>
 
 namespace DB
 {
@@ -69,6 +72,39 @@ void fn_seek(SSTReaderPtr ptr, ColumnFamilyType ct, EngineIteratorSeekType et, B
     auto * reader = reinterpret_cast<MockSSTReader *>(ptr.inner);
     reader->ffi_seek(ptr, ct, et, bf);
 }
+static uint64_t fn_approx_size(SSTReaderPtr ptr, ColumnFamilyType)
+{
+    auto * reader = reinterpret_cast<MockSSTReader *>(ptr.inner);
+    return reader->ffi_approx_size();
+}
+
+static RustStrWithViewVec fn_get_split_keys(SSTReaderPtr ptr, uint64_t splits_count)
+{
+    auto * reader = reinterpret_cast<MockSSTReader *>(ptr.inner);
+    auto length = reader->length();
+    RUNTIME_CHECK(splits_count > 1);
+    auto size_per_split = length / splits_count;
+    auto * vec = new std::vector<std::string>();
+    auto split_key_count = splits_count - 1;
+    auto it = reader->getBegin();
+    BaseBuffView * buffs = createBaseBuffViewArray(split_key_count);
+    for (size_t i = 0; i < split_key_count; i++)
+    {
+        for (size_t j = 0; j < size_per_split; j++)
+        {
+            it++;
+        }
+        vec->push_back(it->first);
+        new (&buffs[i]) BaseBuffView{.data = vec->back().data(), .len = vec->back().size()};
+    }
+    GCMonitor::instance().add(RawObjType::MockVecOfString, 1);
+    return RustStrWithViewVec{
+        .buffs = buffs,
+        .len = split_key_count,
+        .inner = RawRustPtr{
+            .ptr = new RustStrWithViewVecInner{.vec = vec, .buffs = buffs},
+            .type = static_cast<RawRustPtrType>(RawObjType::MockVecOfString)}};
+}
 
 SSTReaderInterfaces make_mock_sst_reader_interface()
 {
@@ -81,6 +117,8 @@ SSTReaderInterfaces make_mock_sst_reader_interface()
         .fn_gc = fn_gc,
         .fn_kind = fn_kind,
         .fn_seek = fn_seek,
+        .fn_approx_size = fn_approx_size,
+        .fn_get_split_keys = fn_get_split_keys,
     };
 }
 } // namespace DB

dbms/src/Debug/MockSSTReader.h

@@ -37,10 +37,13 @@ struct MockSSTReader
     using Key = std::pair<std::string, ColumnFamilyType>;
     struct Data : std::vector<std::pair<std::string, std::string>>
     {
-        DISALLOW_COPY(Data);
         Data(Data &&) = default;
         Data & operator=(Data &&) = default;
         Data() = default;
+        static Data copyFrom(const Data & other) { return other; }
+
+    private:
+        Data(const Data &) = default;
     };
 
     explicit MockSSTReader(const Data & data_, SSTFormatKind kind_)
@@ -49,7 +52,14 @@ struct MockSSTReader
         , end(data_.end())
         , remained(iter != end)
         , kind(kind_)
-    {}
+        , len(data_.size())
+    {
+        total_bytes = 0;
+        for (auto it = data_.begin(); it != data_.end(); it++)
+        {
+            total_bytes += it->second.size();
+        }
+    }
 
     static SSTReaderPtr ffi_get_cf_file_reader(const Data & data_, SSTFormatKind kind_)
     {
@@ -96,6 +106,14 @@ struct MockSSTReader
         }
     }
 
+    size_t ffi_approx_size() { return total_bytes; }
+
+    size_t length() const { return len; }
+
+    Data::const_iterator getBegin() const { return begin; }
+
+    Data::const_iterator getEnd() const { return end; }
+
     static std::map<Key, MockSSTReader::Data> & getMockSSTData() { return MockSSTData; }
 
 private:
@@ -104,6 +122,8 @@ struct MockSSTReader
     Data::const_iterator end;
     bool remained;
     SSTFormatKind kind;
+    size_t total_bytes;
+    size_t len;
 
     // (region_id, cf) -> Data
     static std::map<Key, MockSSTReader::Data> MockSSTData;

dbms/src/Storages/DeltaMerge/CMakeLists.txt

@@ -29,6 +29,7 @@ add_headers_and_sources(delta_merge ./Delta)
 add_headers_and_sources(delta_merge ./ReadThread)
 add_headers_and_sources(delta_merge ./Remote)
 add_headers_and_sources(delta_merge ./Remote/DataStore)
+add_headers_and_sources(delta_merge ./Decode)
 
 add_library(delta_merge ${delta_merge_headers} ${delta_merge_sources})
 target_link_libraries(delta_merge PRIVATE dbms page)

dbms/src/Storages/DeltaMerge/BoundedSSTFilesToBlockInputStream.cpp → dbms/src/Storages/DeltaMerge/Decode/BoundedSSTFilesToBlockInputStream.cpp

@@ -15,10 +15,10 @@
 #include <Interpreters/Context.h>
 #include <Poco/File.h>
 #include <Storages/DeltaMerge/DMVersionFilterBlockInputStream.h>
+#include <Storages/DeltaMerge/Decode/SSTFilesToBlockInputStream.h>
 #include <Storages/DeltaMerge/DeltaMergeStore.h>
 #include <Storages/DeltaMerge/File/DMFileBlockOutputStream.h>
 #include <Storages/DeltaMerge/PKSquashingBlockInputStream.h>
-#include <Storages/DeltaMerge/SSTFilesToBlockInputStream.h>
 #include <Storages/KVStore/Decode/PartitionStreams.h>
 #include <Storages/KVStore/FFI/ProxyFFI.h>
 #include <Storages/StorageDeltaMerge.h>
@@ -31,7 +31,8 @@ namespace DM
 BoundedSSTFilesToBlockInputStream::BoundedSSTFilesToBlockInputStream( //
     SSTFilesToBlockInputStreamPtr child,
     const ColId pk_column_id_,
-    const DecodingStorageSchemaSnapshotConstPtr & schema_snap)
+    const DecodingStorageSchemaSnapshotConstPtr & schema_snap,
+    size_t split_id)
     : pk_column_id(pk_column_id_)
     , _raw_child(std::move(child))
 {
@@ -43,7 +44,8 @@ BoundedSSTFilesToBlockInputStream::BoundedSSTFilesToBlockInputStream( //
     auto stream = std::make_shared<PKSquashingBlockInputStream</*need_extra_sort=*/true>>(
         _raw_child,
         pk_column_id,
-        is_common_handle);
+        is_common_handle,
+        split_id);
     mvcc_compact_stream = std::make_unique<DMVersionFilterBlockInputStream<DM_VERSION_FILTER_MODE_COMPACT>>(
         stream,
         *(schema_snap->column_defines),
@@ -71,6 +73,12 @@ SSTFilesToBlockInputStream::ProcessKeys BoundedSSTFilesToBlockInputStream::getPr
     return _raw_child->process_keys;
 }
 
+
+size_t BoundedSSTFilesToBlockInputStream::getSplitId() const
+{
+    return _raw_child->getSplitId();
+}
+
 RegionPtr BoundedSSTFilesToBlockInputStream::getRegion() const
 {
     return _raw_child->region;