Merge pull request ClickHouse#63513 from ClickHouse/threadpool-runner-local-prefetch-read-pool

Make `MergeTreePrefetchedReadPool` safer

Author: antonio2368

src/Storages/MergeTree/MergeTreePrefetchedReadPool.cpp

@@ -46,63 +46,38 @@ bool MergeTreePrefetchedReadPool::TaskHolder::operator<(const TaskHolder & other
 }
 
 
-MergeTreePrefetchedReadPool::PrefetchedReaders::~PrefetchedReaders()
-{
-    for (auto & prefetch_future : prefetch_futures)
-        if (prefetch_future.valid())
-            prefetch_future.wait();
-}
-
 MergeTreePrefetchedReadPool::PrefetchedReaders::PrefetchedReaders(
+    ThreadPool & pool,
     MergeTreeReadTask::Readers readers_,
     Priority priority_,
-    MergeTreePrefetchedReadPool & pool_)
+    MergeTreePrefetchedReadPool & read_prefetch)
     : is_valid(true)
     , readers(std::move(readers_))
+    , prefetch_runner(pool, "ReadPrepare")
 {
-    try
-    {
-        prefetch_futures.reserve(1 + readers.prewhere.size());
-
-        prefetch_futures.push_back(pool_.createPrefetchedFuture(readers.main.get(), priority_));
+    prefetch_runner(read_prefetch.createPrefetchedTask(readers.main.get(), priority_));
 
-        for (const auto & reader : readers.prewhere)
-            prefetch_futures.push_back(pool_.createPrefetchedFuture(reader.get(), priority_));
+    for (const auto & reader : readers.prewhere)
+        prefetch_runner(read_prefetch.createPrefetchedTask(reader.get(), priority_));
 
-        fiu_do_on(FailPoints::prefetched_reader_pool_failpoint,
-        {
-            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Failpoint for prefetched reader enabled");
-        });
-    }
-    catch (...) /// in case of memory exceptions we have to wait
+    fiu_do_on(FailPoints::prefetched_reader_pool_failpoint,
     {
-        for (auto & prefetch_future : prefetch_futures)
-            if (prefetch_future.valid())
-                prefetch_future.wait();
-
-        throw;
-    }
+        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Failpoint for prefetched reader enabled");
+    });
 }
 
 void MergeTreePrefetchedReadPool::PrefetchedReaders::wait()
 {
     ProfileEventTimeIncrement<Microseconds> watch(ProfileEvents::WaitPrefetchTaskMicroseconds);
-    for (auto & prefetch_future : prefetch_futures)
-        prefetch_future.wait();
+    prefetch_runner.waitForAllToFinish();
 }
 
 MergeTreeReadTask::Readers MergeTreePrefetchedReadPool::PrefetchedReaders::get()
 {
     SCOPE_EXIT({ is_valid = false; });
     ProfileEventTimeIncrement<Microseconds> watch(ProfileEvents::WaitPrefetchTaskMicroseconds);
 
-    /// First wait for completion of all futures.
-    for (auto & prefetch_future : prefetch_futures)
-        prefetch_future.wait();
-
-    /// Then rethrow first exception if any.
-    for (auto & prefetch_future : prefetch_futures)
-        prefetch_future.get();
+    prefetch_runner.waitForAllToFinishAndRethrowFirstError();
 
     return std::move(readers);
 }
@@ -139,22 +114,20 @@ MergeTreePrefetchedReadPool::MergeTreePrefetchedReadPool(
     fillPerThreadTasks(pool_settings.threads, pool_settings.sum_marks);
 }
 
-std::future<void> MergeTreePrefetchedReadPool::createPrefetchedFuture(IMergeTreeReader * reader, Priority priority)
+std::function<void()> MergeTreePrefetchedReadPool::createPrefetchedTask(IMergeTreeReader * reader, Priority priority)
 {
     /// In order to make a prefetch we need to wait for marks to be loaded. But we just created
     /// a reader (which starts loading marks in its constructor), then if we do prefetch right
     /// after creating a reader, it will be very inefficient. We can do prefetch for all parts
     /// only inside this MergeTreePrefetchedReadPool, where read tasks are created and distributed,
     /// and we cannot block either, therefore make prefetch inside the pool and put the future
     /// into the thread task. When a thread calls getTask(), it will wait for it is not ready yet.
-    auto task = [=, context = getContext()]() mutable
+    return [=, context = getContext()]() mutable
     {
         /// For async read metrics in system.query_log.
         PrefetchIncrement watch(context->getAsyncReadCounters());
         reader->prefetchBeginOfRange(priority);
     };
-
-    return scheduleFromThreadPoolUnsafe<void>(std::move(task), prefetch_threadpool, "ReadPrepare", priority);
 }
 
 void MergeTreePrefetchedReadPool::createPrefetchedReadersForTask(ThreadTask & task)
@@ -164,7 +137,7 @@ void MergeTreePrefetchedReadPool::createPrefetchedReadersForTask(ThreadTask & ta
 
     auto extras = getExtras();
     auto readers = MergeTreeReadTask::createReaders(task.read_info, extras, task.ranges);
-    task.readers_future = std::make_unique<PrefetchedReaders>(std::move(readers), task.priority, *this);
+    task.readers_future = std::make_unique<PrefetchedReaders>(prefetch_threadpool, std::move(readers), task.priority, *this);
 }
 
 void MergeTreePrefetchedReadPool::startPrefetches()

src/Storages/MergeTree/MergeTreePrefetchedReadPool.h

@@ -1,5 +1,6 @@
 #pragma once
 #include <Storages/MergeTree/MergeTreeReadPoolBase.h>
+#include <Common/threadPoolCallbackRunner.h>
 #include <Common/ThreadPool_fwd.h>
 #include <IO/AsyncReadCounters.h>
 #include <boost/heap/priority_queue.hpp>
@@ -51,18 +52,18 @@ class MergeTreePrefetchedReadPool : public MergeTreeReadPoolBase, private WithCo
     class PrefetchedReaders
     {
     public:
-        PrefetchedReaders() = default;
-        PrefetchedReaders(MergeTreeReadTask::Readers readers_, Priority priority_, MergeTreePrefetchedReadPool & pool_);
+        PrefetchedReaders(
+            ThreadPool & pool, MergeTreeReadTask::Readers readers_, Priority priority_, MergeTreePrefetchedReadPool & read_prefetch);
 
         void wait();
         MergeTreeReadTask::Readers get();
         bool valid() const { return is_valid; }
-        ~PrefetchedReaders();
 
     private:
         bool is_valid = false;
         MergeTreeReadTask::Readers readers;
-        std::vector<std::future<void>> prefetch_futures;
+
+        ThreadPoolCallbackRunnerLocal<void> prefetch_runner;
     };
 
     struct ThreadTask
@@ -108,7 +109,7 @@ class MergeTreePrefetchedReadPool : public MergeTreeReadPoolBase, private WithCo
 
     void startPrefetches();
     void createPrefetchedReadersForTask(ThreadTask & task);
-    std::future<void> createPrefetchedFuture(IMergeTreeReader * reader, Priority priority);
+    std::function<void()> createPrefetchedTask(IMergeTreeReader * reader, Priority priority);
 
     MergeTreeReadTaskPtr stealTask(size_t thread, MergeTreeReadTask * previous_task);
     MergeTreeReadTaskPtr createTask(ThreadTask & thread_task, MergeTreeReadTask * previous_task);

src/Storages/MergeTree/MergeTreeReadPoolBase.cpp

@@ -53,6 +53,19 @@ void MergeTreeReadPoolBase::fillPerPartInfos()
         MergeTreeReadTaskInfo read_task_info;
 
         read_task_info.data_part = part_with_ranges.data_part;
+
+        const auto & data_part = read_task_info.data_part;
+        if (data_part->isProjectionPart())
+        {
+            read_task_info.parent_part = data_part->storage.getPartIfExists(
+                data_part->getParentPartName(),
+                {MergeTreeDataPartState::PreActive, MergeTreeDataPartState::Active, MergeTreeDataPartState::Outdated});
+
+            if (!read_task_info.parent_part)
+                throw Exception(ErrorCodes::LOGICAL_ERROR, "Did not find parent part {} for projection part {}",
+                            data_part->getParentPartName(), data_part->getDataPartStorage().getFullPath());
+        }
+
         read_task_info.part_index_in_query = part_with_ranges.part_index_in_query;
         read_task_info.alter_conversions = part_with_ranges.alter_conversions;
 

src/Storages/MergeTree/MergeTreeReadTask.h

@@ -56,6 +56,8 @@ struct MergeTreeReadTaskInfo
 {
     /// Data part which should be read while performing this task
     DataPartPtr data_part;
+    /// Parent part of the projection part
+    DataPartPtr parent_part;
     /// For `part_index` virtual column
     size_t part_index_in_query;
     /// Alter converversionss that should be applied on-fly for part.