Storages: Refine SegmentReadTaskScheduler::add to reduce lock content…

…ion (#9027) (#9029)

close #9024

dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.cpp

@@ -38,25 +38,15 @@ SegmentReadTaskScheduler::~SegmentReadTaskScheduler()
 }
 
 void SegmentReadTaskScheduler::add(const SegmentReadTaskPoolPtr & pool)
+{
+    // To avoid schedule from always failing to acquire the pending_mtx.
+    std::lock_guard lock(add_mtx);
+    submitPendingPool(pool);
+}
+
+void SegmentReadTaskScheduler::addPool(const SegmentReadTaskPoolPtr & pool)
 {
     assert(pool != nullptr);
-    if (pool->getPendingSegmentCount() <= 0)
-    {
-        LOG_INFO(
-            pool->getLogger(),
-            "Ignored for no segment to read, pool_id={} table_id={}",
-            pool->pool_id,
-            pool->physical_table_id);
-        return;
-    }
-    Stopwatch sw_add;
-    // `add_lock` is only used in this function to make all threads calling `add` to execute serially.
-    std::lock_guard add_lock(add_mtx);
-    add_waittings.fetch_add(1, std::memory_order_relaxed);
-    // `lock` is used to protect data.
-    std::lock_guard lock(mtx);
-    add_waittings.fetch_sub(1, std::memory_order_relaxed);
-    Stopwatch sw_do_add;
     read_pools.emplace(pool->pool_id, pool);
 
     const auto & tasks = pool->getTasks();
@@ -65,15 +55,43 @@ void SegmentReadTaskScheduler::add(const SegmentReadTaskPoolPtr & pool)
         auto seg_id = pa.first;
         merging_segments[pool->physical_table_id][seg_id].push_back(pool->pool_id);
     }
+}
+
+void SegmentReadTaskScheduler::submitPendingPool(SegmentReadTaskPoolPtr pool)
+{
+    assert(pool != nullptr);
+    if (pool->getPendingSegmentCount() <= 0)
+    {
+        LOG_INFO(pool->getLogger(), "Ignored for no segment to read, pool_id={}", pool->pool_id);
+        return;
+    }
+    Stopwatch sw;
+    std::lock_guard lock(pending_mtx);
+    pending_pools.push_back(pool);
     LOG_INFO(
         pool->getLogger(),
-        "Added, pool_id={} block_slots={} segment_count={} pool_count={} cost={:.3f}us do_add_cost={:.3f}us", //
+        "Submitted, pool_id={} segment_count={} pending_pools={} cost={}ns",
         pool->pool_id,
-        pool->getFreeBlockSlots(),
-        tasks.size(),
-        read_pools.size(),
-        sw_add.elapsed() / 1000.0,
-        sw_do_add.elapsed() / 1000.0);
+        pool->getPendingSegmentCount(),
+        pending_pools.size(),
+        sw.elapsed());
+}
+
+void SegmentReadTaskScheduler::reapPendingPools()
+{
+    SegmentReadTaskPools pools;
+    {
+        std::lock_guard lock(pending_mtx);
+        pools.swap(pending_pools);
+    }
+    if (!pools.empty())
+    {
+        for (const auto & pool : pools)
+        {
+            addPool(pool);
+        }
+        LOG_INFO(log, "Added, pool_ids={}, pool_count={}", pools, read_pools.size());
+    }
 }
 
 MergedTaskPtr SegmentReadTaskScheduler::scheduleMergedTask(SegmentReadTaskPoolPtr & pool)
@@ -263,49 +281,43 @@ std::tuple<UInt64, UInt64, UInt64> SegmentReadTaskScheduler::scheduleOneRound()
 
 bool SegmentReadTaskScheduler::schedule()
 {
-    Stopwatch sw_sched_total;
-    std::lock_guard lock(mtx);
-    Stopwatch sw_do_sched;
-
-    auto pool_count = read_pools.size();
+    Stopwatch sw_sched;
     UInt64 erased_pool_count = 0;
     UInt64 sched_null_count = 0;
     UInt64 sched_succ_count = 0;
     UInt64 sched_round = 0;
     bool can_sched_more_tasks = false;
+    UInt64 reap_pending_pools_ns = 0;
     do
     {
         ++sched_round;
+        Stopwatch sw;
+        reapPendingPools();
+        reap_pending_pools_ns += sw.elapsed();
         auto [erase, null, succ] = scheduleOneRound();
         erased_pool_count += erase;
         sched_null_count += null;
         sched_succ_count += succ;
 
         can_sched_more_tasks = succ > 0 && !read_pools.empty();
-        // If no thread is waitting to add tasks and there are some tasks to be scheduled, run scheduling again.
-        // Avoid releasing and acquiring `mtx` repeatly.
-        // This is common when query concurrency is low, but individual queries are heavy.
-    } while (add_waittings.load(std::memory_order_relaxed) <= 0 && can_sched_more_tasks);
+    } while (can_sched_more_tasks);
 
     if (read_pools.empty())
     {
         GET_METRIC(tiflash_storage_read_thread_counter, type_sche_no_pool).Increment();
     }
 
-    auto total_ms = sw_sched_total.elapsedMilliseconds();
-    if (total_ms >= 100)
+    if (auto total_ms = sw_sched.elapsedMilliseconds(); total_ms >= 50)
     {
         LOG_INFO(
             log,
-            "schedule sched_round={} pool_count={} erased_pool_count={} sched_null_count={} sched_succ_count={} "
-            "cost={}ms do_sched_cost={}ms",
+            "schedule sched_round={} erased_pool_count={} sched_null_count={} sched_succ_count={} reap={}ms cost={}ms",
             sched_round,
-            pool_count,
             erased_pool_count,
             sched_null_count,
             sched_succ_count,
-            total_ms,
-            sw_do_sched.elapsedMilliseconds());
+            reap_pending_pools_ns / 1000'000,
+            total_ms);
     }
     return can_sched_more_tasks;
 }

dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.h

@@ -18,15 +18,18 @@
 
 namespace DB::DM
 {
-
-// `SegmentReadTaskScheduler` is a global singleton. All `SegmentReadTaskPool` will be added to it and be scheduled by it.
-// 1. `UnorderedInputStream`/`UnorderedSourceOps` will call `SegmentReadTaskScheduler::add` to add a `SegmentReadTaskPool`
-// object to the `read_pools` list and index segments information into `merging_segments`.
-// 2. A schedule-thread will scheduling read tasks:
-//   a. It scans the `read_pools` list and check if `SegmentReadTaskPool` need be scheduled.
-//   b. Chooses a `SegmentReadTask` of the `SegmentReadTaskPool`, if other `SegmentReadTaskPool` will read the same
-//      `SegmentReadTask`, pop them, and build a `MergedTask`.
-//   c. Sends the MergedTask to read threads(SegmentReader).
+namespace tests
+{
+class SegmentReadTasksPoolTest;
+}
+// SegmentReadTaskScheduler is a global singleton. All SegmentReadTaskPool objects will be added to it and be scheduled by it.
+// - Threads of computational layer will call SegmentReadTaskScheduler::add to add a SegmentReadTaskPool object to the `pending_pools`.
+// - Call path: UnorderedInputStream/UnorderedSourceOps -> SegmentReadTaskScheduler::add -> SegmentReadTaskScheduler::submitPendingPool
+//
+// - `sched_thread` will scheduling read tasks.
+// - Call path: schedLoop -> schedule -> reapPendingPools -> scheduleOneRound
+// - reapPeningPools will swap the `pending_pools` and add these pools to `read_pools` and `merging_segments`.
+// - scheduleOneRound will scan `read_pools` and choose segments to read.
 class SegmentReadTaskScheduler
 {
 public:
@@ -39,7 +42,7 @@ class SegmentReadTaskScheduler
     ~SegmentReadTaskScheduler();
     DISALLOW_COPY_AND_MOVE(SegmentReadTaskScheduler);
 
-    // Add SegmentReadTaskPool to `read_pools` and index segments into merging_segments.
+    // Add `pool` to `pending_pools`.
     void add(const SegmentReadTaskPoolPtr & pool);
 
     void pushMergedTask(const MergedTaskPtr & p) { merged_task_pool.push(p); }
@@ -72,11 +75,15 @@ class SegmentReadTaskScheduler
     // `erased_pool_count` - how many stale pools have beed erased.
     // `sched_null_count` - how many pools do not require scheduling.
     // `sched_succ_count` - how many pools is scheduled.
-    std::tuple<UInt64, UInt64, UInt64> scheduleOneRound() EXCLUSIVE_LOCKS_REQUIRED(mtx);
+    std::tuple<UInt64, UInt64, UInt64> scheduleOneRound();
+
+    void submitPendingPool(SegmentReadTaskPoolPtr pool);
+    void reapPendingPools();
+    void addPool(const SegmentReadTaskPoolPtr & pool);
 
     // To restrict the instantaneous concurrency of `add` and avoid `schedule` from always failing to acquire the lock.
     std::mutex add_mtx;
-    std::mutex mtx;
+    // `read_pools` and `merging_segment` are only accessed by `sched_thread`.
     // pool_id -> pool
     std::unordered_map<UInt64, SegmentReadTaskPoolPtr> read_pools;
     // table_id -> {seg_id -> pool_ids, seg_id -> pool_ids, ...}
@@ -89,7 +96,9 @@ class SegmentReadTaskScheduler
 
     LoggerPtr log;
 
-    // To count how many threads are waitting to add tasks.
-    std::atomic<Int64> add_waittings{0};
+    std::mutex pending_mtx;
+    SegmentReadTaskPools pending_pools GUARDED_BY(pending_mtx);
+
+    friend class tests::SegmentReadTasksPoolTest;
 };
 } // namespace DB::DM

dbms/src/Storages/DeltaMerge/SegmentReadTaskPool.h

@@ -304,3 +304,15 @@ struct fmt::formatter<DB::DM::SegmentReadTaskPtr>
             t->read_snapshot->delta->getDeltaIndexEpoch());
     }
 };
+
+template <>
+struct fmt::formatter<DB::DM::SegmentReadTaskPoolPtr>
+{
+    static constexpr auto parse(format_parse_context & ctx) { return ctx.begin(); }
+
+    template <typename FormatContext>
+    auto format(const DB::DM::SegmentReadTaskPoolPtr & pool, FormatContext & ctx) const
+    {
+        return fmt::format_to(ctx.out(), "{}", pool->pool_id);
+    }
+};