storage: Refactor disaggregated read flow

dbms/src/Common/ThreadedWorker.h

@@ -0,0 +1,204 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Common/Logger.h>
+#include <Common/MPMCQueue.h>
+#include <Common/Stopwatch.h>
+#include <Common/ThreadManager.h>
+#include <common/logger_useful.h>
+
+#include <ext/scope_guard.h>
+#include <magic_enum.hpp>
+#include <mutex>
+
+namespace DB
+{
+
+/// A ThreadedWorker spawns n (n=concurrency) threads to process tasks.
+/// Each thread takes a task from the source queue, do some work, and put the result into the
+/// result queue.
+/// ThreadedWorker manages the state of its result queue, for example, when all tasks are finished
+/// or when there are errors, the result queue will be finished or cancelled.
+template <typename Src, typename Dest>
+class ThreadedWorker
+{
+public:
+    virtual ~ThreadedWorker()
+    {
+        wait();
+    }
+
+    void startInBackground() noexcept
+    {
+        std::call_once(start_flag, [this] {
+            watch_start.restart();
+
+            LOG_DEBUG(log, "Starting {} workers, concurrency={}", getName(), concurrency);
+            active_workers = concurrency;
+
+            for (size_t index = 0; index < concurrency; ++index)
+            {
+                thread_manager->schedule(true, getName(), [this, index] {
+                    total_wait_schedule_ms += watch_start.elapsedMilliseconds();
+
+                    workerLoop(index);
+                    handleWorkerFinished();
+                });
+            }
+        });
+    }
+
+    void wait() noexcept
+    {
+        std::call_once(wait_flag, [this] {
+            try
+            {
+                // thread_manager->wait can be only called once.
+                thread_manager->wait();
+            }
+            catch (...)
+            {
+                // This should not occur, as we should have caught all exceptions in workerLoop.
+                auto error = getCurrentExceptionMessage(false);
+                LOG_WARNING(log, "{} meet unexepcted error: {}", getName(), error);
+            }
+        });
+    }
+
+public:
+    const size_t concurrency;
+    const std::shared_ptr<MPMCQueue<Src>> source_queue;
+    const std::shared_ptr<MPMCQueue<Dest>> result_queue;
+
+protected:
+    const LoggerPtr log;
+    std::shared_ptr<ThreadManager> thread_manager;
+    std::atomic<Int64> active_workers = 0;
+
+    ThreadedWorker(
+        std::shared_ptr<MPMCQueue<Src>> source_queue_,
+        std::shared_ptr<MPMCQueue<Dest>> result_queue_,
+        LoggerPtr log_,
+        size_t concurrency_)
+        : concurrency(concurrency_)
+        , source_queue(source_queue_)
+        , result_queue(result_queue_)
+        , log(log_)
+        , thread_manager(newThreadManager())
+    {
+        RUNTIME_CHECK(concurrency > 0);
+    }
+
+    virtual String getName() const noexcept = 0;
+
+    virtual Dest doWork(const Src & task) = 0;
+
+private:
+    void handleWorkerFinished()
+    {
+        active_workers--;
+        if (active_workers == 0)
+        {
+            std::call_once(finish_flag, [this] {
+                LOG_DEBUG(
+                    log,
+                    "{} workers finished, total_processed_tasks={} concurrency={} elapsed={:.3f}s total_wait_schedule={:.3f}s total_wait_upstream={:.3f}s",
+                    getName(),
+                    total_processed_tasks,
+                    concurrency,
+                    watch_start.elapsedSeconds(),
+                    total_wait_schedule_ms / 1000.0,
+                    total_wait_upstream_ms / 1000.0);
+                // Note: the result queue may be already cancelled, but it is fine.
+                result_queue->finish();
+            });
+        }
+    }
+
+    void workerLoop(size_t thread_idx)
+    {
+        try
+        {
+            while (true)
+            {
+                Src task;
+
+                Stopwatch w;
+                auto pop_result = source_queue->pop(task);
+                total_wait_upstream_ms += w.elapsedMilliseconds();
+
+                if (pop_result != MPMCQueueResult::OK)
+                {
+                    if (pop_result == MPMCQueueResult::FINISHED)
+                    {
+                        // No more work to do, just exit.
+                        // The FINISH signal will be passed to downstreams when
+                        // all workers are exited.
+                        break;
+                    }
+                    else if (pop_result == MPMCQueueResult::CANCELLED)
+                    {
+                        // There are errors, populate the error to downstreams
+                        // immediately.
+                        auto cancel_reason = source_queue->getCancelReason();
+                        LOG_WARNING(log, "{}#{} meeting error from upstream: {}", getName(), thread_idx, cancel_reason);
+                        result_queue->cancelWith(cancel_reason);
+                        break;
+                    }
+                    else
+                    {
+                        RUNTIME_CHECK_MSG(false, "Unexpected pop MPMCQueueResult: {}", magic_enum::enum_name(pop_result));
+                    }
+                }
+
+                auto work_result = doWork(task);
+                auto push_result = result_queue->push(work_result);
+                if (push_result != MPMCQueueResult::OK)
+                {
+                    if (push_result == MPMCQueueResult::CANCELLED)
+                    {
+                        // The result_queue has been closed by other workers, so we
+                        // no longer process any tasks from the source.
+                        break;
+                    }
+                    else
+                    {
+                        RUNTIME_CHECK_MSG(false, "Unexpected push MPMCQueueResult: {}", magic_enum::enum_name(push_result));
+                    }
+                }
+
+                total_processed_tasks++;
+            }
+        }
+        catch (...)
+        {
+            auto error = getCurrentExceptionMessage(false);
+            LOG_ERROR(log, "{}#{} meet error: {}", getName(), thread_idx, error);
+            result_queue->cancelWith(fmt::format("{} failed: {}", getName(), error));
+        }
+    }
+
+private:
+    Stopwatch watch_start;
+    std::once_flag start_flag;
+    std::once_flag wait_flag;
+    std::once_flag finish_flag;
+    std::atomic<Int64> total_processed_tasks = 0;
+    std::atomic<Int64> total_wait_schedule_ms = 0;
+    std::atomic<Int64> total_wait_upstream_ms = 0;
+};
+
+} // namespace DB

dbms/src/Common/TiFlashMetrics.h

@@ -232,7 +232,8 @@ namespace DB
         F(type_cache_occupy, {{"type", "cache_occupy"}}, ExpBuckets{0.01, 2, 20}),                                                                  \
         F(type_build_read_task, {{"type", "build_read_task"}}, ExpBuckets{0.01, 2, 20}),                                                            \
         F(type_seg_next_task, {{"type", "seg_next_task"}}, ExpBuckets{0.01, 2, 20}),                                                                \
-        F(type_seg_build_stream, {{"type", "seg_build_stream"}}, ExpBuckets{0.01, 2, 20}))                                                          \
+        F(type_seg_build_stream, {{"type", "seg_build_stream"}}, ExpBuckets{0.01, 2, 20}), \
+        F(type_seg_read_prefix, {{"type", "seg_read_prefix"}}, ExpBuckets{0.01, 2, 20}))                                                          \
     M(tiflash_disaggregated_details, "", Counter,                                                                                                   \
         F(type_cftiny_read, {{"type", "cftiny_read"}}),                                                                                             \
         F(type_cftiny_fetch, {{"type", "cftiny_fetch"}}))                                                                                           \

dbms/src/DataStreams/AddExtraTableIDColumnInputStream.cpp

@@ -20,8 +20,9 @@ namespace DB
 AddExtraTableIDColumnInputStream::AddExtraTableIDColumnInputStream(
     BlockInputStreamPtr input,
     int extra_table_id_index,
-    TableID physical_table_id)
-    : action(input->getHeader(), extra_table_id_index, physical_table_id)
+    TableID physical_table_id_)
+    : physical_table_id(physical_table_id_)
+    , action(input->getHeader(), extra_table_id_index)
 {
     children.push_back(input);
 }
@@ -32,7 +33,7 @@ Block AddExtraTableIDColumnInputStream::readImpl()
     if (!res)
         return res;
 
-    auto ok = action.transform(res);
+    auto ok = action.transform(res, physical_table_id);
     if (!ok)
         return {};
 

dbms/src/DataStreams/AddExtraTableIDColumnInputStream.h

@@ -42,6 +42,7 @@ class AddExtraTableIDColumnInputStream : public IProfilingBlockInputStream
     Block readImpl() override;
 
 private:
+    const TableID physical_table_id;
     AddExtraTableIDColumnTransformAction action;
 };
 

dbms/src/DataStreams/AddExtraTableIDColumnTransformAction.cpp

@@ -47,21 +47,17 @@ Block AddExtraTableIDColumnTransformAction::buildHeader(
 
 AddExtraTableIDColumnTransformAction::AddExtraTableIDColumnTransformAction(
     const Block & inner_header_,
-    int extra_table_id_index_,
-    TableID physical_table_id_)
+    int extra_table_id_index_)
     : header(buildHeader(inner_header_, extra_table_id_index_))
     , extra_table_id_index(extra_table_id_index_)
-    , physical_table_id(physical_table_id_)
 {
 }
 
 AddExtraTableIDColumnTransformAction::AddExtraTableIDColumnTransformAction(
     const DM::ColumnDefines & columns_to_read_,
-    int extra_table_id_index_,
-    TableID physical_table_id_)
+    int extra_table_id_index_)
     : header(buildHeader(columns_to_read_, extra_table_id_index_))
     , extra_table_id_index(extra_table_id_index_)
-    , physical_table_id(physical_table_id_)
 {
 }
 
@@ -70,7 +66,7 @@ Block AddExtraTableIDColumnTransformAction::getHeader() const
     return header;
 }
 
-bool AddExtraTableIDColumnTransformAction::transform(Block & block)
+bool AddExtraTableIDColumnTransformAction::transform(Block & block, TableID physical_table_id)
 {
     if (unlikely(!block))
         return true;

dbms/src/DataStreams/AddExtraTableIDColumnTransformAction.h

@@ -31,15 +31,13 @@ struct AddExtraTableIDColumnTransformAction
 
     AddExtraTableIDColumnTransformAction(
         const Block & inner_header_,
-        int extra_table_id_index_,
-        TableID physical_table_id_);
+        int extra_table_id_index_);
 
     AddExtraTableIDColumnTransformAction(
         const DM::ColumnDefines & columns_to_read_,
-        int extra_table_id_index_,
-        TableID physical_table_id_);
+        int extra_table_id_index_);
 
-    bool transform(Block & block);
+    bool transform(Block & block, TableID physical_table_id);
 
     Block getHeader() const;
 
@@ -48,12 +46,10 @@ struct AddExtraTableIDColumnTransformAction
         return total_rows;
     }
 
-
 private:
     Block header;
     // position of the ExtraPhysTblID column in column_names parameter in the StorageDeltaMerge::read function.
     const int extra_table_id_index;
-    const TableID physical_table_id;
 
     size_t total_rows = 0;
 };

dbms/src/Flash/Coprocessor/GenSchemaAndColumn.cpp

@@ -83,10 +83,10 @@ NamesAndTypes genNamesAndTypes(const TiDBTableScan & table_scan, const StringRef
     return genNamesAndTypes(table_scan.getColumns(), column_prefix);
 }
 
-std::tuple<DM::ColumnDefinesPtr, size_t> genColumnDefinesForDisaggregatedRead(const TiDBTableScan & table_scan)
+std::tuple<DM::ColumnDefinesPtr, int> genColumnDefinesForDisaggregatedRead(const TiDBTableScan & table_scan)
 {
     auto column_defines = std::make_shared<DM::ColumnDefines>();
-    size_t extra_table_id_index = InvalidColumnID;
+    int extra_table_id_index = InvalidColumnID;
     column_defines->reserve(table_scan.getColumnSize());
     for (Int32 i = 0; i < table_scan.getColumnSize(); ++i)
     {

dbms/src/Flash/Coprocessor/GenSchemaAndColumn.h

@@ -39,7 +39,7 @@ using ColumnDefinesPtr = std::shared_ptr<std::vector<ColumnDefine>>;
 } // namespace DM
 
 // The column defines and `extra table id index`
-std::tuple<DM::ColumnDefinesPtr, size_t>
+std::tuple<DM::ColumnDefinesPtr, int>
 genColumnDefinesForDisaggregatedRead(const TiDBTableScan & table_scan);
 
 } // namespace DB