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

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,13 @@ 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:
+        auto * inner = reinterpret_cast<RustStrWithViewVecInner *>(ptr);
+        delete inner;
+        break;
     }
 }
 
@@ -167,12 +174,20 @@ 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;
 }
 
+static RustStrWithView fn_get_config_json(RaftStoreProxyPtr, uint64_t)
+{
+    std::string * s = new std::string("{\"raftstore\":{\"snap-handle-pool-size\":4}}");
+    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 +199,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;
@@ -941,7 +957,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 +994,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)

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,38 @@ 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;
+    std::vector<std::string> * 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()};
+    }
+    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 +116,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/BoundedSSTFilesToBlockInputStream.cpp

@@ -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,12 +44,17 @@ 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),
         _raw_child->opts.gc_safepoint,
         is_common_handle);
+    LOG_INFO(
+        &Poco::Logger::get("BoundedSSTFilesToBlockInputStream"),
+        "create bounded sst file stream, split_id={}",
+        split_id);
 }
 
 void BoundedSSTFilesToBlockInputStream::readPrefix()
@@ -71,6 +77,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;

dbms/src/Storages/DeltaMerge/PKSquashingBlockInputStream.h

@@ -18,6 +18,7 @@
 #include <Interpreters/sortBlock.h>
 #include <Storages/DeltaMerge/DeltaMergeHelpers.h>
 #include <Storages/DeltaMerge/RowKeyRange.h>
+#include <Storages/DeltaMerge/SSTFilesToBlockInputStream.h>
 
 namespace DB
 {
@@ -32,10 +33,15 @@ template <bool need_extra_sort>
 class PKSquashingBlockInputStream final : public IBlockInputStream
 {
 public:
-    PKSquashingBlockInputStream(BlockInputStreamPtr child, ColId pk_column_id_, bool is_common_handle_)
+    PKSquashingBlockInputStream(
+        BlockInputStreamPtr child,
+        ColId pk_column_id_,
+        bool is_common_handle_,
+        size_t split_id_ = DM::SSTScanSoftLimit::HEAD_SPLIT)
         : sorted_input_stream(child)
         , pk_column_id(pk_column_id_)
         , is_common_handle(is_common_handle_)
+        , split_id(split_id_)
     {
         assert(sorted_input_stream != nullptr);
         cur_block = {};
@@ -86,7 +92,8 @@ class PKSquashingBlockInputStream final : public IBlockInputStream
             }
 #endif
 
-            const size_t cut_offset = findCutOffsetInNextBlock(cur_block, next_block, pk_column_id, is_common_handle);
+            const size_t cut_offset
+                = findCutOffsetInNextBlock(split_id, cur_block, next_block, pk_column_id, is_common_handle);
             if (unlikely(cut_offset == 0))
                 // There is no pk overlap between `cur_block` and `next_block`, or `next_block` is empty, just return `cur_block`.
                 return finializeBlock(std::move(cur_block));
@@ -123,11 +130,13 @@ class PKSquashingBlockInputStream final : public IBlockInputStream
 
 private:
     static size_t findCutOffsetInNextBlock(
+        size_t split_id,
         const Block & cur_block,
         const Block & next_block,
         const ColId pk_column_id,
         bool is_common_handle)
     {
+        UNUSED(split_id);
         assert(cur_block);
         if (!next_block)
             return 0;
@@ -138,6 +147,7 @@ class PKSquashingBlockInputStream final : public IBlockInputStream
         auto next_col = getByColumnId(next_block, pk_column_id).column;
         RowKeyColumnContainer next_rowkey_column(next_col, is_common_handle);
         size_t cut_offset = 0;
+
         for (/* */; cut_offset < next_col->size(); ++cut_offset)
         {
             const auto next_pk = next_rowkey_column.getRowKeyValue(cut_offset);
@@ -147,8 +157,12 @@ class PKSquashingBlockInputStream final : public IBlockInputStream
                 {
                     if (unlikely(next_pk < last_curr_pk))
                         throw Exception(
-                            "InputStream is not sorted, pk in next block is smaller than current block: "
-                                + next_pk.toDebugString() + " < " + last_curr_pk.toDebugString(),
+                            fmt::format(
+                                "InputStream is not sorted, pk in next block {} is smaller than current block {}, "
+                                "split_id={}",
+                                next_pk.toDebugString(),
+                                last_curr_pk.toDebugString(),
+                                split_id),
                             ErrorCodes::LOGICAL_ERROR);
                 }
                 break;
@@ -180,6 +194,7 @@ class PKSquashingBlockInputStream final : public IBlockInputStream
 
     bool first_read = true;
     const bool is_common_handle;
+    size_t split_id;
 };
 
 } // namespace DM