raft: Fix the max concurrency of IngestSST (relase-8.5) (#10610)

close #10605

* Parse `raftstore.apply-low-priority-pool-size` in `KVStore::fetchProxyConfig`
* `KVStore::getMaxPrehandleSubtaskSize` return different max concurrency according to the job_type

Signed-off-by: JaySon-Huang <tshent@qq.com>

Author: JaySon-Huang

dbms/src/Debug/MockKVStore/MockRaftStoreProxy.h

@@ -274,7 +274,8 @@ struct MockRaftStoreProxy : MutexLockWrap
         , table_id(1)
         , cluster_ver(RaftstoreVer::V1)
     {
-        proxy_config_string = R"({"raftstore":{"snap-handle-pool-size":4},"server":{"engine-addr":"123"}})";
+        proxy_config_string
+            = R"({"raftstore":{"snap-handle-pool-size":4,"apply-low-priority-pool-size":8},"server":{"engine-addr":"123"}})";
     }
 
     LoggerPtr log;

dbms/src/Storages/KVStore/KVStore.cpp

@@ -133,22 +133,28 @@ void KVStore::fetchProxyConfig(const TiFlashRaftProxyHelper * proxy_helper)
             Poco::JSON::Parser parser;
             auto obj = parser.parse(cpp_string);
             auto ptr = obj.extract<Poco::JSON::Object::Ptr>();
+
             auto raftstore = ptr->getObject("raftstore");
             proxy_config_summary.snap_handle_pool_size = raftstore->getValue<uint64_t>("snap-handle-pool-size");
+            proxy_config_summary.apply_low_priority_pool_size
+                = raftstore->getValue<uint64_t>("apply-low-priority-pool-size");
+
             auto server = ptr->getObject("server");
             proxy_config_summary.engine_addr = server->getValue<std::string>("engine-addr");
+
             LOG_INFO(
                 log,
-                "Parsed proxy config: snap_handle_pool_size={} engine_addr={}",
+                "Parsed proxy config: snap_handle_pool_size={} apply_low_priority_pool_size={} engine_addr={}",
                 proxy_config_summary.snap_handle_pool_size,
+                proxy_config_summary.apply_low_priority_pool_size,
                 proxy_config_summary.engine_addr);
             proxy_config_summary.valid = true;
         }
         catch (...)
         {
             proxy_config_summary.valid = false;
-            // we don't care
-            LOG_WARNING(log, "Can't parse config from proxy {}", cpp_string);
+            // ignore the error and log a error message
+            tryLogCurrentWarningException(log, fmt::format("Can't parse config from proxy {}", cpp_string));
         }
     }
 }

dbms/src/Storages/KVStore/KVStore.h

@@ -113,7 +113,10 @@ static_assert(magic_enum::enum_count<PersistRegionReason>() == sizeof(PersistReg
 struct ProxyConfigSummary
 {
     bool valid = false;
+    // The max concurrency of PreHandleSnapshot tasks in proxy.
     size_t snap_handle_pool_size = 0;
+    // The max concurrency of IngestSST tasks in proxy.
+    size_t apply_low_priority_pool_size = 0;
     std::string engine_addr;
 };
 
@@ -144,7 +147,7 @@ class KVStore final : private boost::noncopyable
     metapb::Store & debugMutStoreMeta();
     FileUsageStatistics getFileUsageStatistics() const;
     // Proxy will validate and refit the config items from the toml file.
-    const ProxyConfigSummary & getProxyConfigSummay() const { return proxy_config_summary; }
+    const ProxyConfigSummary & getProxyConfigSummary() const { return proxy_config_summary; }
     void reportThreadAllocInfo(std::string_view, ReportThreadAllocateInfoType type, uint64_t value);
     static void reportThreadAllocBatch(std::string_view, ReportThreadAllocateInfoBatch data);
     JointThreadInfoJeallocMapPtr getJointThreadInfoJeallocMap() const { return joint_memory_allocation_map; }
@@ -229,7 +232,8 @@ class KVStore final : private boost::noncopyable
     size_t getOngoingPrehandleTaskCount() const;
     size_t getOngoingPrehandleSubtaskCount() const;
     EngineStoreApplyRes handleIngestSST(UInt64 region_id, SSTViewVec, UInt64 index, UInt64 term, TMTContext & tmt);
-    size_t getMaxParallelPrehandleSize() const;
+    size_t getMaxParallelPrehandleSize(DM::FileConvertJobType job_type) const;
+    size_t getMaxPrehandleSubtaskSize(DM::FileConvertJobType job_type) const;
 
 public: // Raft Read
     void addReadIndexEvent(Int64 f) { read_index_event_flag += f; }

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

@@ -122,10 +122,11 @@ void PreHandlingTrace::waitForSubtaskResources(uint64_t region_id, size_t parall
         .Observe(watch.elapsedSeconds());
     LOG_INFO(
         log,
-        "Prehandle resource acquired after {:.3f} seconds, region_id={} parallel={}",
+        "Prehandle resource acquired after {:.3f} seconds, region_id={} parallel={} limit={}",
         watch.elapsedSeconds(),
         region_id,
-        parallel);
+        parallel,
+        parallel_subtask_limit);
     CurrentMetrics::sub(CurrentMetrics::RaftNumWaitedParallelPrehandlingTasks);
 }
 
@@ -301,20 +302,41 @@ PrehandleResult KVStore::preHandleSnapshotToFiles(
     return PrehandleResult{};
 }
 
-size_t KVStore::getMaxParallelPrehandleSize() const
+size_t KVStore::getMaxParallelPrehandleSize(DM::FileConvertJobType job_type) const
 {
-    const auto & proxy_config = getProxyConfigSummay();
-    size_t total_concurrency = 0;
-    if (proxy_config.valid)
+    return getMaxPrehandleSubtaskSize(job_type);
+}
+
+size_t KVStore::getMaxPrehandleSubtaskSize(DM::FileConvertJobType job_type) const
+{
+    const auto & proxy_config = getProxyConfigSummary();
+    switch (job_type)
+    {
+    case DM::FileConvertJobType::ApplySnapshot:
     {
-        total_concurrency = proxy_config.snap_handle_pool_size;
+        if (proxy_config.valid && proxy_config.snap_handle_pool_size > 0)
+        {
+            return proxy_config.snap_handle_pool_size;
+        }
+        else
+        {
+            auto cpu_num = std::thread::hardware_concurrency();
+            return static_cast<size_t>(std::clamp(cpu_num * 0.7, 2.0, 16.0));
+        }
     }
-    else
+    case DM::FileConvertJobType::IngestSST:
     {
-        auto cpu_num = std::thread::hardware_concurrency();
-        total_concurrency = static_cast<size_t>(std::clamp(cpu_num * 0.7, 2.0, 16.0));
+        if (proxy_config.valid && proxy_config.apply_low_priority_pool_size > 0)
+        {
+            return proxy_config.apply_low_priority_pool_size;
+        }
+        else
+        {
+            auto cpu_num = std::thread::hardware_concurrency();
+            return std::max(1, static_cast<size_t>(cpu_num));
+        }
+    }
     }
-    return total_concurrency;
 }
 
 // If size is 0, do not parallel prehandle for this snapshot, which is regular.
@@ -323,7 +345,8 @@ static inline std::pair<std::vector<std::string>, size_t> getSplitKey(
     LoggerPtr log,
     KVStore * kvstore,
     RegionPtr new_region,
-    std::shared_ptr<DM::SSTFilesToBlockInputStream> sst_stream)
+    std::shared_ptr<DM::SSTFilesToBlockInputStream> sst_stream,
+    DM::FileConvertJobType job_type)
 {
     // We don't use this is the single snapshot is small, due to overhead in decoding.
     constexpr size_t default_parallel_prehandle_threshold = 1 * 1024 * 1024 * 1024;
@@ -356,7 +379,12 @@ static inline std::pair<std::vector<std::string>, size_t> getSplitKey(
     // so we must add 1 here.
     auto ongoing_count = kvstore->getOngoingPrehandleSubtaskCount() + 1;
     uint64_t want_split_parts = 0;
-    auto total_concurrency = kvstore->getMaxParallelPrehandleSize();
+    // If total_concurrency is 4, and prehandle-pool is sized 8,
+    // and if there are 4 ongoing snapshots, then we will not parallel prehandling any new snapshot.
+    // This is because in serverless, too much parallelism causes performance reduction.
+    // So, if there is already enough parallelism that is used to prehandle,
+    // it is not necessary to manually split a snapshot.
+    auto total_concurrency = kvstore->getMaxParallelPrehandleSize(job_type);
     if (total_concurrency + 1 > ongoing_count)
     {
         // Current thread takes 1 which is in `ongoing_count`.
@@ -385,7 +413,7 @@ static inline std::pair<std::vector<std::string>, size_t> getSplitKey(
     if (split_keys.size() + 1 < want_split_parts)
     {
         // If there are too few split keys, the `split_keys` itself may be not be uniformly distributed,
-        // it is even better that we still handle it sequantially.
+        // it is even better that we still handle it sequentially.
         split_keys.clear();
         LOG_INFO(
             log,
@@ -717,8 +745,11 @@ PrehandleResult KVStore::preHandleSSTsToDTFiles(
             };
 
             // `split_keys` do not begin with 'z'.
-            auto [split_keys, approx_bytes] = getSplitKey(log, this, new_region, sst_stream);
-            prehandling_trace.waitForSubtaskResources(region_id, split_keys.size() + 1, getMaxParallelPrehandleSize());
+            auto [split_keys, approx_bytes] = getSplitKey(log, this, new_region, sst_stream, job_type);
+            prehandling_trace.waitForSubtaskResources(
+                region_id,
+                split_keys.size() + 1,
+                getMaxPrehandleSubtaskSize(job_type));
             ReadFromStreamResult result;
             if (split_keys.empty())
             {

dbms/src/Storages/KVStore/TMTContext.cpp

@@ -173,15 +173,13 @@ void TMTContext::initS3GCManager(const TiFlashRaftProxyHelper * proxy_helper)
     if (S3::ClientFactory::instance().isEnabled() && !context.getSharedContextDisagg()->isDisaggregatedComputeMode())
     {
         kvstore->fetchProxyConfig(proxy_helper);
-        if (kvstore->getProxyConfigSummay().valid)
+        if (kvstore->getProxyConfigSummary().valid)
         {
-            LOG_INFO(
-                Logger::get(),
-                "Build s3gc manager from proxy's conf engine_addr={}",
-                kvstore->getProxyConfigSummay().engine_addr);
+            auto engine_addr = kvstore->getProxyConfigSummary().engine_addr;
+            LOG_INFO(Logger::get(), "Build s3gc manager from proxy's conf engine_addr={}", engine_addr);
             s3gc_owner = OwnerManager::createS3GCOwner(
                 context,
-                /*id*/ kvstore->getProxyConfigSummay().engine_addr,
+                /*id*/ engine_addr,
                 etcd_client);
         }
         else if (!raftproxy_config.pd_addrs.empty())

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

@@ -871,7 +871,8 @@ try
 {
     auto & ctx = TiFlashTestEnv::getGlobalContext();
     proxy_instance->cluster_ver = RaftstoreVer::V2;
-    proxy_instance->proxy_config_string = R"({"raftstore":{"snap-handle-pool-size":3},"server":{"engine-addr":"123"}})";
+    proxy_instance->proxy_config_string
+        = R"({"raftstore":{"snap-handle-pool-size":3, "apply-low-priority-pool-size":8},"server":{"engine-addr":"123"}})";
     KVStore & kvs = getKVS();
     kvs.fetchProxyConfig(proxy_helper.get());
     ASSERT_NE(proxy_helper->sst_reader_interfaces.fn_key, nullptr);
@@ -912,7 +913,7 @@ try
         MockSSTReader::getMockSSTData().clear();
         DB::FailPointHelper::enablePauseFailPoint(DB::FailPoints::pause_before_prehandle_subtask, 100);
         std::vector<std::thread> ths;
-        auto runId = [&](size_t ths_id) {
+        auto run_id = [&](size_t ths_id) {
             auto [value_write, value_default] = proxy_instance->generateTiKVKeyValue(111, 999);
             MockSSTGenerator default_cf{region_ids[ths_id], table_id, ColumnFamilyType::Default};
             for (HandleID h = table_limits[ths_id]; h < table_limits[ths_id + 1]; h++)
@@ -942,13 +943,13 @@ try
                     std::nullopt);
             }
         };
-        ths.push_back(std::thread(runId, 0));
+        ths.push_back(std::thread(run_id, 0));
         std::this_thread::sleep_for(std::chrono::milliseconds(300));
 
         ASSERT_EQ(kvs.getOngoingPrehandleTaskCount(), 1);
         for (size_t ths_id = 1; ths_id < region_ids.size(); ths_id++)
         {
-            ths.push_back(std::thread(runId, ths_id));
+            ths.push_back(std::thread(run_id, ths_id));
         }
 
         auto loop = 0;