pingcap · ti-chi-bot · Dec 14, 2023 · Dec 8, 2023 · Dec 14, 2023 · Dec 14, 2023
diff --git a/dbms/src/Flash/Pipeline/Schedule/Tasks/RFWaitTask.h b/dbms/src/Flash/Pipeline/Schedule/Tasks/RFWaitTask.h
@@ -64,14 +64,15 @@ class RFWaitTask : public Task
 
     static void submitReadyRfsAndSegmentTaskPool(
         const RuntimeFilteList & ready_rf_list,
-        const DM::SegmentReadTaskPoolPtr & task_pool)
+        const DM::SegmentReadTaskPoolPtr & task_pool,
+        const LoggerPtr & log)
     {
         for (const RuntimeFilterPtr & rf : ready_rf_list)
         {
             auto rs_operator = rf->parseToRSOperator(task_pool->getColumnToRead());
             task_pool->appendRSOperator(rs_operator);
         }
-        DM::SegmentReadTaskScheduler::instance().add(task_pool);
+        DM::SegmentReadTaskScheduler::instance().add(task_pool, log);
     }
 
 private:
@@ -82,7 +83,7 @@ class RFWaitTask : public Task
         filterAndMoveReadyRfs(waiting_rf_list, ready_rf_list);
         if (waiting_rf_list.empty() || stopwatch.elapsed() >= max_wait_time_ns)
         {
-            submitReadyRfsAndSegmentTaskPool(ready_rf_list, task_pool);
+            submitReadyRfsAndSegmentTaskPool(ready_rf_list, task_pool, log);
             return ExecTaskStatus::FINISHED;
         }
         return ExecTaskStatus::WAITING;

diff --git a/dbms/src/Operators/UnorderedSourceOp.cpp b/dbms/src/Operators/UnorderedSourceOp.cpp
@@ -64,7 +64,7 @@ void UnorderedSourceOp::operatePrefixImpl()
     std::call_once(task_pool->addToSchedulerFlag(), [&]() {
         if (waiting_rf_list.empty())
         {
-            DM::SegmentReadTaskScheduler::instance().add(task_pool);
+            DM::SegmentReadTaskScheduler::instance().add(task_pool, log);
         }
         else
         {
@@ -74,7 +74,7 @@ void UnorderedSourceOp::operatePrefixImpl()
 
             if (max_wait_time_ms <= 0 || waiting_rf_list.empty())
             {
-                RFWaitTask::submitReadyRfsAndSegmentTaskPool(ready_rf_list, task_pool);
+                RFWaitTask::submitReadyRfsAndSegmentTaskPool(ready_rf_list, task_pool, log);
             }
             else
             {

diff --git a/dbms/src/Storages/DeltaMerge/ReadThread/MergedTask.cpp b/dbms/src/Storages/DeltaMerge/ReadThread/MergedTask.cpp
@@ -44,7 +44,7 @@ void MergedTask::initOnce()
         auto & [pool, task, stream] = units[cur_idx];
         if (!pool->valid())
         {
-            setStreamFinished(cur_idx);
+            setUnitFinish(cur_idx);
             continue;
         }
         if (pool->isRUExhausted())
@@ -65,7 +65,7 @@ int MergedTask::readOneBlock()
     int read_block_count = 0;
     for (cur_idx = 0; cur_idx < static_cast<int>(units.size()); cur_idx++)
     {
-        if (isStreamFinished(cur_idx))
+        if (units[cur_idx].isFinished())
         {
             continue;
         }
@@ -74,7 +74,7 @@ int MergedTask::readOneBlock()
 
         if (!pool->valid())
         {
-            setStreamFinished(cur_idx);
+            setUnitFinish(cur_idx);
             continue;
         }
 
@@ -94,37 +94,33 @@ int MergedTask::readOneBlock()
         }
         else
         {
-            setStreamFinished(cur_idx);
+            setUnitFinish(cur_idx);
         }
     }
     return read_block_count;
 }
 
 void MergedTask::setException(const DB::Exception & e)
 {
-    for (auto & unit : units)
-    {
-        if (unit.pool != nullptr)
-        {
-            unit.pool->setException(e);
-        }
-    }
+    std::for_each(units.begin(), units.end(), [&e](auto & u) {
+        if (u.pool != nullptr)
+            u.pool->setException(e);
+    });
 }
 
 MergedTaskPtr MergedTaskPool::pop(uint64_t pool_id)
 {
     std::lock_guard lock(mtx);
-    MergedTaskPtr target;
-    for (auto itr = merged_task_pool.begin(); itr != merged_task_pool.end(); ++itr)
+    auto itr = std::find_if(merged_task_pool.begin(), merged_task_pool.end(), [pool_id](const auto & merged_task) {
+        return merged_task->containPool(pool_id);
+    });
+    if (itr != merged_task_pool.end())
     {
-        if ((*itr)->containPool(pool_id))
-        {
-            target = *itr;
-            merged_task_pool.erase(itr);
-            break;
-        }
+        auto target = *itr;
+        merged_task_pool.erase(itr);
+        return target;
     }
-    return target;
+    return nullptr; // Not Found.
 }
 
 void MergedTaskPool::push(const MergedTaskPtr & t)
@@ -136,13 +132,8 @@ void MergedTaskPool::push(const MergedTaskPtr & t)
 bool MergedTaskPool::has(UInt64 pool_id)
 {
     std::lock_guard lock(mtx);
-    for (const auto & t : merged_task_pool)
-    {
-        if (t->containPool(pool_id))
-        {
-            return true;
-        }
-    }
-    return false;
+    return std::any_of(merged_task_pool.begin(), merged_task_pool.end(), [pool_id](const auto & merged_task) {
+        return merged_task->containPool(pool_id);
+    });
 }
 } // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/ReadThread/MergedTask.h b/dbms/src/Storages/DeltaMerge/ReadThread/MergedTask.h
@@ -25,6 +25,29 @@ struct MergedUnit
         , task(task_)
     {}
 
+    ~MergedUnit()
+    {
+        // Calling `setFinish()` for updating memory statistics of `MemoryTracker`.
+        [[maybe_unused]] auto res = setFinish();
+    }
+
+    [[nodiscard]] bool isFinished() const { return pool == nullptr && task == nullptr && stream == nullptr; }
+
+    // If setted return true else return false.
+    [[nodiscard]] bool setFinish()
+    {
+        if (!isFinished())
+        {
+            // For updating memory statistics of `MemoryTracker`.
+            MemoryTrackerSetter setter(true, pool->mem_tracker.get());
+            task = nullptr;
+            stream = nullptr;
+            pool = nullptr;
+            return true;
+        }
+        return false;
+    }
+
     SegmentReadTaskPoolPtr pool; // The information of a read request.
     SegmentReadTaskPtr task; // The information of a segment that want to read.
     BlockInputStreamPtr stream; // BlockInputStream of a segment, will be created by read threads.
@@ -44,7 +67,6 @@ class MergedTask
         , inited(false)
         , cur_idx(-1)
         , finished_count(0)
-        , log(Logger::get())
     {
         passive_merged_segments.fetch_add(units.size() - 1, std::memory_order_relaxed);
         GET_METRIC(tiflash_storage_read_thread_gauge, type_merged_task).Increment();
@@ -54,8 +76,6 @@ class MergedTask
         passive_merged_segments.fetch_sub(units.size() - 1, std::memory_order_relaxed);
         GET_METRIC(tiflash_storage_read_thread_gauge, type_merged_task).Decrement();
         GET_METRIC(tiflash_storage_read_thread_seconds, type_merged_task).Observe(sw.elapsedSeconds());
-        // `setAllStreamFinished` must be called to explicitly releasing all streams for updating memory statistics of `MemoryTracker`.
-        setAllStreamsFinished();
     }
 
     int readBlock();
@@ -64,36 +84,27 @@ class MergedTask
 
     const GlobalSegmentID & getSegmentId() const { return seg_id; }
 
-    size_t getPoolCount() const { return units.size(); }
-
-    std::vector<uint64_t> getPoolIds() const
+    bool containPool(uint64_t pool_id) const
     {
-        std::vector<uint64_t> ids;
-        ids.reserve(units.size());
-        for (const auto & unit : units)
-        {
-            if (unit.pool != nullptr)
-            {
-                ids.push_back(unit.pool->pool_id);
-            }
-        }
-        return ids;
+        return std::any_of(units.begin(), units.end(), [pool_id](const auto & u) {
+            return u.pool != nullptr && u.pool->pool_id == pool_id;
+        });
     }
 
-    bool containPool(uint64_t pool_id) const
+    void setException(const DB::Exception & e);
+
+    String toString() const
     {
-        for (const auto & unit : units)
-        {
-            if (unit.pool != nullptr && unit.pool->pool_id == pool_id)
-            {
-                return true;
-            }
-        }
-        return false;
+        std::vector<UInt64> ids;
+        ids.reserve(units.size());
+        std::for_each(units.begin(), units.end(), [&ids](const auto & u) {
+            if (u.pool != nullptr)
+                ids.push_back(u.pool->pool_id);
+        });
+        return fmt::format("seg_id:{} pool_id:{}", seg_id, ids);
     }
-    void setException(const DB::Exception & e);
 
-    const LoggerPtr getCurrentLogger() const
+    LoggerPtr getCurrentLogger() const
     {
         // `std::cmp_*` is safety to compare negative signed integers and unsigned integers.
         if (likely(
@@ -115,41 +126,13 @@ class MergedTask
 private:
     void initOnce();
     int readOneBlock();
+    void setUnitFinish(int i) { finished_count += units[i].setFinish(); }
 
-    bool isStreamFinished(size_t i) const
-    {
-        return units[i].pool == nullptr && units[i].task == nullptr && units[i].stream == nullptr;
-    }
-
-    void setStreamFinished(size_t i)
-    {
-        if (!isStreamFinished(i))
-        {
-            // `MergedUnit.stream` must be released explicitly for updating memory statistics of `MemoryTracker`.
-            auto & [pool, task, stream] = units[i];
-            {
-                MemoryTrackerSetter setter(true, pool->mem_tracker.get());
-                task = nullptr;
-                stream = nullptr;
-            }
-            pool = nullptr;
-            finished_count++;
-        }
-    }
-
-    void setAllStreamsFinished()
-    {
-        for (size_t i = 0; i < units.size(); ++i)
-        {
-            setStreamFinished(i);
-        }
-    }
     GlobalSegmentID seg_id;
     std::vector<MergedUnit> units;
     bool inited;
     int cur_idx;
     size_t finished_count;
-    LoggerPtr log;
     Stopwatch sw;
     inline static std::atomic<int64_t> passive_merged_segments{0};
 };
@@ -162,17 +145,12 @@ using MergedTaskPtr = std::shared_ptr<MergedTask>;
 class MergedTaskPool
 {
 public:
-    MergedTaskPool()
-        : log(Logger::get())
-    {}
-
     MergedTaskPtr pop(uint64_t pool_id);
     void push(const MergedTaskPtr & t);
     bool has(UInt64 pool_id);
 
 private:
     std::mutex mtx;
     std::list<MergedTaskPtr> merged_task_pool GUARDED_BY(mtx);
-    LoggerPtr log;
 };
 } // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.cpp b/dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.cpp
@@ -30,10 +30,12 @@ SegmentReadTaskScheduler::~SegmentReadTaskScheduler()
     sched_thread.join();
 }
 
-void SegmentReadTaskScheduler::add(const SegmentReadTaskPoolPtr & pool)
+void SegmentReadTaskScheduler::add(const SegmentReadTaskPoolPtr & pool, const LoggerPtr & req_log)
 {
     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);
+    // `lock` is used to protect data.
     std::lock_guard lock(mtx);
     Stopwatch sw_do_add;
     read_pools.add(pool);
@@ -43,12 +45,11 @@ void SegmentReadTaskScheduler::add(const SegmentReadTaskPoolPtr & pool)
     {
         merging_segments[seg_id].push_back(pool->pool_id);
     }
-    auto block_slots = pool->getFreeBlockSlots();
     LOG_DEBUG(
-        log,
+        req_log,
         "Added, pool_id={} block_slots={} segment_count={} pool_count={} cost={:.3f}us do_add_cost={:.3f}us", //
         pool->pool_id,
-        block_slots,
+        pool->getFreeBlockSlots(),
         tasks.size(),
         read_pools.size(),
         sw_add.elapsed() / 1000.0,
@@ -235,9 +236,8 @@ bool SegmentReadTaskScheduler::schedule()
             {
                 LOG_DEBUG(
                     log,
-                    "scheduleMergedTask segment_id={} pool_ids={} cost={}ms pool_count={}",
-                    merged_task->getSegmentId(),
-                    merged_task->getPoolIds(),
+                    "scheduleMergedTask merged_task=<{}> cost={}ms pool_count={}",
+                    merged_task->toString(),
                     elapsed_ms,
                     pool_count);
             }

diff --git a/dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.h b/dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.h
@@ -17,7 +17,6 @@
 #include <Storages/DeltaMerge/ReadThread/MergedTask.h>
 #include <Storages/DeltaMerge/SegmentReadTaskPool.h>
 
-#include <memory>
 namespace DB::DM
 {
 using SegmentReadTaskPoolList = CircularScanList<SegmentReadTaskPool>;
@@ -43,7 +42,7 @@ class SegmentReadTaskScheduler
     DISALLOW_COPY_AND_MOVE(SegmentReadTaskScheduler);
 
     // Add SegmentReadTaskPool to `read_pools` and index segments into merging_segments.
-    void add(const SegmentReadTaskPoolPtr & pool) LOCKS_EXCLUDED(add_mtx, mtx);
+    void add(const SegmentReadTaskPoolPtr & pool, const LoggerPtr & req_log) LOCKS_EXCLUDED(add_mtx, mtx);
 
     void pushMergedTask(const MergedTaskPtr & p) { merged_task_pool.push(p); }
 

diff --git a/dbms/src/Storages/DeltaMerge/ReadThread/SegmentReader.cpp b/dbms/src/Storages/DeltaMerge/ReadThread/SegmentReader.cpp
@@ -99,12 +99,7 @@ class SegmentReader
             }
             if (read_count <= 0)
             {
-                LOG_DEBUG(
-                    log,
-                    "All finished, pool_ids={} segment_id={} read_count={}",
-                    merged_task->getPoolIds(),
-                    merged_task->getSegmentId(),
-                    read_count);
+                LOG_DEBUG(log, "All finished, merged_task=<{}> read_count={}", merged_task->toString(), read_count);
             }
             // If `merged_task` is pushed back to `MergedTaskPool`, it can be accessed by another read thread if it is scheduled.
             // So do not push back to `MergedTaskPool` when exception happened since current read thread can still access to this `merged_task` object and set exception message to it.

diff --git a/dbms/src/Storages/DeltaMerge/ReadThread/UnorderedInputStream.h b/dbms/src/Storages/DeltaMerge/ReadThread/UnorderedInputStream.h
@@ -131,7 +131,7 @@ class UnorderedInputStream : public IProfilingBlockInputStream
         }
         std::call_once(task_pool->addToSchedulerFlag(), [&]() {
             prepareRuntimeFilter();
-            SegmentReadTaskScheduler::instance().add(task_pool);
+            SegmentReadTaskScheduler::instance().add(task_pool, log);
         });
         task_pool_added = true;
     }