From 0fc67ab3c446f30d87a5aabb01b65bd6ceaf1fcc Mon Sep 17 00:00:00 2001 From: Darek Date: Sun, 18 Oct 2026 04:10:08 +0200 Subject: [PATCH] perf(processor): run a one-thread task on the calling thread scheduleTaskAndWaitOrError pushed every task to the queue and woke the whole worker pool, even when only one thread could register. A one-row READ_FTABLE (RETURN 1 and other constant queries) was also marked parallel, so every worker cloned the pipeline before finding no morsel. A task limited to one thread now runs on the caller. Workers and inline callers share numWorkerThreads slots under the scheduler mutex, so max_num_threads still bounds the threads running queries across connections; a caller that finds no free slot uses the queue, and a thread that already holds a slot reuses it. Queries with a timeout and callers that pass launchNewWorkerThread keep the queue path. After a task's children finish, a READ_FTABLE source that fits in one morsel is lowered to one thread. Co-Authored-By: Claude Opus 5.5 --- src/common/task_system/task_scheduler.cpp | 109 +++++++++++- src/include/common/task_system/task.h | 5 +- .../common/task_system/task_scheduler.h | 20 +++ src/include/processor/processor_task.h | 4 + src/processor/processor.cpp | 12 +- src/processor/processor_task.cpp | 20 +++ test/processor/CMakeLists.txt | 2 +- test/processor/task_scheduler_test.cpp | 159 ++++++++++++++++++ 8 files changed, 319 insertions(+), 12 deletions(-) create mode 100644 test/processor/task_scheduler_test.cpp diff --git a/src/common/task_system/task_scheduler.cpp b/src/common/task_system/task_scheduler.cpp index fbb94fcebb..cf04720419 100644 --- a/src/common/task_system/task_scheduler.cpp +++ b/src/common/task_system/task_scheduler.cpp @@ -15,12 +15,21 @@ namespace common { #ifndef __SINGLE_THREADED__ +namespace { + +// The scheduler whose slot this thread occupies: always set on a worker, and on a caller +// while it runs a task inline. A nested inline run on such a thread reuses its slot. +thread_local const TaskScheduler* busySlotOwner = nullptr; + +} // namespace + #if defined(__APPLE__) TaskScheduler::TaskScheduler(uint64_t numWorkerThreads, uint32_t threadQos) #else TaskScheduler::TaskScheduler(uint64_t numWorkerThreads) #endif - : stopWorkerThreads{false}, nextScheduledTaskID{0} { + : stopWorkerThreads{false}, nextScheduledTaskID{0}, maxNumBusyThreads{numWorkerThreads}, + numBusyThreads{0} { #if defined(__APPLE__) this->threadQos = threadQos; #endif @@ -41,17 +50,40 @@ TaskScheduler::~TaskScheduler() { void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, processor::ExecutionContext* context, bool launchNewWorkerThread) { - scheduleTaskAndWaitOrError(task, context->clientContext, launchNewWorkerThread); + scheduleTaskAndWaitOrError(task, context->clientContext, launchNewWorkerThread, nullptr); +} + +void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, + processor::ExecutionContext* context, bool launchNewWorkerThread, TaskBeforeRunFn beforeRun) { + scheduleTaskAndWaitOrError(task, context->clientContext, launchNewWorkerThread, beforeRun); } void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, main::ClientContext* clientContext, bool launchNewWorkerThread) { + scheduleTaskAndWaitOrError(task, clientContext, launchNewWorkerThread, nullptr); +} + +void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, + main::ClientContext* clientContext, bool launchNewWorkerThread, TaskBeforeRunFn beforeRun) { for (auto& dependency : task->children) { - scheduleTaskAndWaitOrError(dependency, clientContext); + scheduleTaskAndWaitOrError(dependency, clientContext, false, beforeRun); if (dependency->terminate()) { return; } } + if (beforeRun != nullptr) { + beforeRun(task.get()); + } + // A one-thread task runs on the caller so the pool stays asleep. A timeout still uses + // the pool: the waiting caller is what notices the clock and interrupts the worker. + if (!launchNewWorkerThread && task->maxNumThreads == 1 && + !(clientContext != nullptr && clientContext->hasTimeout()) && + runOnCallingThread(task.get())) { + if (task->hasException()) { + std::rethrow_exception(task->getExceptionPtr()); + } + return; + } std::thread newWorkerThread; if (launchNewWorkerThread) { // Note that newWorkerThread is not executing yet. However, we still call @@ -60,7 +92,12 @@ void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task // numThreadsRegistered field of the task, tt does not keep track of the thread ids or // anything specific to the thread. task->registerThread(); - newWorkerThread = std::thread(runTask, task.get()); + // The new thread stands in for the blocked caller, so it reuses the caller's slot for + // nested one-thread tasks. + newWorkerThread = std::thread([this, taskPtr = task.get()] { + busySlotOwner = this; + runTask(taskPtr); + }); } auto scheduledTask = pushTaskIntoQueue(task); cv.notify_all(); @@ -103,7 +140,44 @@ void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task } } +bool TaskScheduler::runOnCallingThread(Task* task) { + if (busySlotOwner == this) { + if (!task->registerThread()) { + return false; + } + runTask(task); + return true; + } + { + lock_t lck{taskSchedulerMtx}; + if (numBusyThreads == maxNumBusyThreads || !task->registerThread()) { + return false; + } + numBusyThreads++; + } + // Gives the slot back even if runTask throws. + struct SlotGuard { + TaskScheduler* scheduler; + const TaskScheduler* previousOwner; + ~SlotGuard() { + busySlotOwner = previousOwner; + lock_t lck{scheduler->taskSchedulerMtx}; + scheduler->numBusyThreads--; + // A queued task may be waiting for this slot rather than for a worker. + if (!scheduler->taskQueue.empty()) { + scheduler->cv.notify_one(); + } + } + } guard{this, busySlotOwner}; + busySlotOwner = this; + // Writes made here reach a dependent task: it runs on this thread, or this thread queues it + // under taskSchedulerMtx. Writes of a child run by workers arrive through its taskMtx. + runTask(task); + return true; +} + void TaskScheduler::runWorkerThread() { + busySlotOwner = this; #if defined(__APPLE__) qos_class_t qosClass = (qos_class_t)threadQos; if (qosClass != QOS_CLASS_DEFAULT && qosClass != QOS_CLASS_UNSPECIFIED) { @@ -132,11 +206,17 @@ void TaskScheduler::runWorkerThread() { } scheduledTask->task->deRegisterThreadAndFinalizeTask(); scheduledTask = nullptr; + numBusyThreads--; } cv.wait(lck, [&] { - scheduledTask = getTaskAndRegister(); + if (numBusyThreads < maxNumBusyThreads) { + scheduledTask = getTaskAndRegister(); + } return scheduledTask != nullptr || stopWorkerThreads; }); + if (scheduledTask != nullptr) { + numBusyThreads++; + } lck.unlock(); if (stopWorkerThreads) { return; @@ -158,17 +238,30 @@ TaskScheduler::~TaskScheduler() { void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, processor::ExecutionContext* context, bool launchNewWorkerThread) { - scheduleTaskAndWaitOrError(task, context->clientContext, launchNewWorkerThread); + scheduleTaskAndWaitOrError(task, context->clientContext, launchNewWorkerThread, nullptr); } void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, - main::ClientContext* clientContext, bool) { + processor::ExecutionContext* context, bool launchNewWorkerThread, TaskBeforeRunFn beforeRun) { + scheduleTaskAndWaitOrError(task, context->clientContext, launchNewWorkerThread, beforeRun); +} + +void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, + main::ClientContext* clientContext, bool launchNewWorkerThread) { + scheduleTaskAndWaitOrError(task, clientContext, launchNewWorkerThread, nullptr); +} + +void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr& task, + main::ClientContext* clientContext, bool, TaskBeforeRunFn beforeRun) { for (auto& dependency : task->children) { - scheduleTaskAndWaitOrError(dependency, clientContext); + scheduleTaskAndWaitOrError(dependency, clientContext, false, beforeRun); if (dependency->terminate()) { return; } } + if (beforeRun != nullptr) { + beforeRun(task.get()); + } task->registerThread(); // runTask deregisters, so we don't need to deregister explicitly here runTask(task.get()); diff --git a/src/include/common/task_system/task.h b/src/include/common/task_system/task.h index 7efb1131c1..4bb67fa07e 100644 --- a/src/include/common/task_system/task.h +++ b/src/include/common/task_system/task.h @@ -15,8 +15,9 @@ using lock_t = std::unique_lock; /** * Task represents a task that can be executed by multiple threads in the TaskScheduler. Task is a * virtual class. Users of TaskScheduler need to extend the Task class and implement at - * least a virtual run() function. Users can assume that before T.run() is called, a worker thread W - * has grabbed task T from the TaskScheduler's queue and registered itself to T. They can also + * least a virtual run() function. Users can assume that before T.run() is called, a thread W (a + * worker that grabbed T from the TaskScheduler's queue, or the caller running a one-thread T + * inline) has registered itself to T. They can also * assume that after run() is called W will deregister itself from T. When deregistering, if W is * the last worker to finish on T, i.e., once W finishes, T will be completed, the * finalize() will be called. So the run() and finalize() calls are separate diff --git a/src/include/common/task_system/task_scheduler.h b/src/include/common/task_system/task_scheduler.h index 82230ec255..88cd72d87a 100644 --- a/src/include/common/task_system/task_scheduler.h +++ b/src/include/common/task_system/task_scheduler.h @@ -20,6 +20,10 @@ struct ScheduledTask { uint64_t ID; }; +// Invoked after a task's dependencies finish and before it is launched. Used to +// lower the thread limit once produced input is visible. +using TaskBeforeRunFn = void (*)(Task*); + /** * TaskScheduler is a library that manages a set of worker threads that can execute tasks that are * put into a task queue. Each task accepts a maximum number of threads. Users of TaskScheduler @@ -64,6 +68,11 @@ class LBUG_API TaskScheduler { processor::ExecutionContext* context, bool launchNewWorkerThread = false); void scheduleTaskAndWaitOrError(const std::shared_ptr& task, main::ClientContext* clientContext, bool launchNewWorkerThread = false); + void scheduleTaskAndWaitOrError(const std::shared_ptr& task, + processor::ExecutionContext* context, bool launchNewWorkerThread, + TaskBeforeRunFn beforeRun); + void scheduleTaskAndWaitOrError(const std::shared_ptr& task, + main::ClientContext* clientContext, bool launchNewWorkerThread, TaskBeforeRunFn beforeRun); static TaskScheduler* Get(const main::ClientContext& context); @@ -78,6 +87,8 @@ class LBUG_API TaskScheduler { std::shared_ptr getTaskAndRegister(); static void runTask(Task* task); + bool runOnCallingThread(Task* task); + private: std::deque> taskQueue; bool stopWorkerThreads; @@ -88,6 +99,10 @@ class LBUG_API TaskScheduler { #if defined(__APPLE__) uint32_t threadQos; // Thread quality of service for worker threads. #endif + // Workers and callers running a task inline share numWorkerThreads slots, so inline + // execution does not raise the database-wide thread cap. Guarded by taskSchedulerMtx. + uint64_t maxNumBusyThreads; + uint64_t numBusyThreads; }; #else // Single-threaded version of TaskScheduler @@ -102,6 +117,11 @@ class TaskScheduler { processor::ExecutionContext* context, bool launchNewWorkerThread = false); void scheduleTaskAndWaitOrError(const std::shared_ptr& task, main::ClientContext* clientContext, bool launchNewWorkerThread = false); + void scheduleTaskAndWaitOrError(const std::shared_ptr& task, + processor::ExecutionContext* context, bool launchNewWorkerThread, + TaskBeforeRunFn beforeRun); + void scheduleTaskAndWaitOrError(const std::shared_ptr& task, + main::ClientContext* clientContext, bool launchNewWorkerThread, TaskBeforeRunFn beforeRun); static TaskScheduler* Get(const main::ClientContext& context); diff --git a/src/include/processor/processor_task.h b/src/include/processor/processor_task.h index 5d2247df62..3187969756 100644 --- a/src/include/processor/processor_task.h +++ b/src/include/processor/processor_task.h @@ -18,6 +18,10 @@ class ProcessorTask : public common::Task { bool terminate() override; + // A READ_FTABLE source that already fits in one morsel has no work for a second thread. + // Called after child tasks, so an accumulate scan sees the filled table. + void narrowToSingleMorselSource(); + private: bool sharedStateInitialized; Sink* sink; diff --git a/src/processor/processor.cpp b/src/processor/processor.cpp index 6029643faf..c4a6bd8f1c 100644 --- a/src/processor/processor.cpp +++ b/src/processor/processor.cpp @@ -14,6 +14,16 @@ using namespace lbug::storage; namespace lbug { namespace processor { +namespace { + +void narrowSingleMorselScan(Task* task) { + if (auto* processorTask = dynamic_cast(task)) { + processorTask->narrowToSingleMorselSource(); + } +} + +} // namespace + std::shared_ptr ResultSetPool::getOrCreate(ResultSetDescriptor* descriptor, MemoryManager* memoryManager) { const auto threadID = std::this_thread::get_id(); @@ -75,7 +85,7 @@ std::unique_ptr QueryProcessor::execute(PhysicalPlan* physica initTask(task.get()); auto progressBar = ProgressBar::Get(*context->clientContext); progressBar->startProgress(context->queryID); - taskScheduler->scheduleTaskAndWaitOrError(task, context); + taskScheduler->scheduleTaskAndWaitOrError(task, context, false, narrowSingleMorselScan); progressBar->endProgress(context->queryID); return sink->getQueryResult(); } diff --git a/src/processor/processor_task.cpp b/src/processor/processor_task.cpp index 2ae71809d7..5e7261e283 100644 --- a/src/processor/processor_task.cpp +++ b/src/processor/processor_task.cpp @@ -5,6 +5,8 @@ #include "main/database.h" #include "main/settings.h" #include "processor/execution_context.h" +#include "processor/operator/table_function_call.h" +#include "processor/operator/table_scan/ftable_scan_function.h" #include "processor/processor.h" #include "processor/result/result_set.h" #include "storage/buffer_manager/memory_manager.h" @@ -19,6 +21,24 @@ ProcessorTask::ProcessorTask(Sink* sink, ExecutionContext* executionContext) .getValue()}, sharedStateInitialized{false}, sink{sink}, executionContext{executionContext} {} +void ProcessorTask::narrowToSingleMorselSource() { + PhysicalOperator* op = sink; + while (!op->isSource()) { + op = op->getChild(0); + } + if (op->getOperatorType() != PhysicalOperatorType::TABLE_FUNCTION_CALL) { + return; + } + auto* call = op->ptrCast(); + if (call->getInfo().function.name != FTableScan::name || call->getInfo().bindData == nullptr) { + return; + } + auto* bindData = call->getInfo().bindData->constPtrCast(); + if (bindData->table->getNumTuples() <= bindData->morselSize) { + setSingleThreadedTask(); + } +} + void ProcessorTask::run() { // We need the lock when cloning because multiple threads can be accessing to clone, // which is not thread safe diff --git a/test/processor/CMakeLists.txt b/test/processor/CMakeLists.txt index 57a6659b2d..bbae316da6 100644 --- a/test/processor/CMakeLists.txt +++ b/test/processor/CMakeLists.txt @@ -1 +1 @@ -add_lbug_test(processor_test intersect_kernels_test.cpp) +add_lbug_test(processor_test intersect_kernels_test.cpp task_scheduler_test.cpp) diff --git a/test/processor/task_scheduler_test.cpp b/test/processor/task_scheduler_test.cpp new file mode 100644 index 0000000000..ec12b83e96 --- /dev/null +++ b/test/processor/task_scheduler_test.cpp @@ -0,0 +1,159 @@ +#include +#include +#include +#include +#include + +#include "common/task_system/task_scheduler.h" +#include "gtest/gtest.h" +#include "main/client_context.h" +#include "main/connection.h" +#include "main/database.h" + +using namespace lbug; + +namespace { + +void updatePeak(std::atomic& peak, uint64_t now) { + auto prev = peak.load(); + while (now > prev && !peak.compare_exchange_weak(prev, now)) {} +} + +class CountingTask final : public common::Task { +public: + CountingTask(std::atomic& active, std::atomic& peak) + : Task{1}, active{active}, peak{peak} {} + + void run() override { + threadID = std::this_thread::get_id(); + updatePeak(peak, ++active); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + --active; + } + + std::thread::id threadID; + +private: + std::atomic& active; + std::atomic& peak; +}; + +class BlockingTask final : public common::Task { +public: + BlockingTask(std::atomic& running, std::atomic& release) + : Task{1}, running{running}, release{release} {} + + void run() override { + ++running; + // Spin so that every blocked caller gives its slot back at nearly the same moment. + while (!release.load()) {} + } + +private: + std::atomic& running; + std::atomic& release; +}; + +class NestingTask final : public common::Task { +public: + NestingTask(main::ClientContext* context, std::shared_ptr child) + : Task{1}, context{context}, child{std::move(child)} {} + + void run() override { + threadID = std::this_thread::get_id(); + common::TaskScheduler::Get(*context)->scheduleTaskAndWaitOrError(child, context); + } + + std::thread::id threadID; + +private: + main::ClientContext* context; + std::shared_ptr child; +}; + +void schedule(main::Connection& conn, const std::shared_ptr& task) { + auto* context = conn.getClientContext(); + common::TaskScheduler::Get(*context)->scheduleTaskAndWaitOrError(task, context); +} + +} // namespace + +TEST(TaskSchedulerTest, OneThreadTaskRunsOnCallingThread) { + main::Database db(":memory:", main::SystemConfig(-1u, 4)); + main::Connection conn(&db); + std::atomic active{0}, peak{0}; + auto task = std::make_shared(active, peak); + schedule(conn, task); + EXPECT_EQ(task->threadID, std::this_thread::get_id()); +} + +TEST(TaskSchedulerTest, OneThreadTasksStayWithinThreadLimit) { + constexpr uint64_t maxNumThreads = 2; + constexpr uint64_t numCallers = 8; + main::Database db(":memory:", main::SystemConfig(-1u, maxNumThreads)); + std::vector> conns; + for (auto i = 0u; i < numCallers; i++) { + conns.push_back(std::make_unique(&db)); + } + std::atomic active{0}, peak{0}; + std::vector callers; + for (auto& conn : conns) { + callers.emplace_back([&, conn = conn.get()] { + for (auto i = 0; i < 3; i++) { + schedule(*conn, std::make_shared(active, peak)); + } + }); + } + for (auto& caller : callers) { + caller.join(); + } + EXPECT_GE(peak.load(), 1u); + EXPECT_LE(peak.load(), maxNumThreads); +} + +// Callers that find every slot taken queue their tasks. When the inline callers finish, each +// freed slot must wake a worker, so the queued tasks run side by side. +TEST(TaskSchedulerTest, QueuedTasksRunInParallelOnceInlineSlotsFree) { + constexpr uint64_t maxNumThreads = 4; + main::Database db(":memory:", main::SystemConfig(-1u, maxNumThreads)); + std::vector> conns; + for (auto i = 0u; i < 2 * maxNumThreads; i++) { + conns.push_back(std::make_unique(&db)); + } + std::atomic running{0}; + std::atomic release{false}; + std::vector callers; + for (auto i = 0u; i < maxNumThreads; i++) { + callers.emplace_back([&, conn = conns[i].get()] { + schedule(*conn, std::make_shared(running, release)); + }); + } + while (running.load() < maxNumThreads) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + std::atomic active{0}, peak{0}; + for (auto i = maxNumThreads; i < 2 * maxNumThreads; i++) { + callers.emplace_back([&, conn = conns[i].get()] { + schedule(*conn, std::make_shared(active, peak)); + }); + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + EXPECT_EQ(peak.load(), 0u); + release = true; + for (auto& caller : callers) { + caller.join(); + } + EXPECT_EQ(peak.load(), maxNumThreads); +} + +// With one slot held by the outer task, a nested one-thread task can only finish by reusing it. +TEST(TaskSchedulerTest, NestedOneThreadTaskReusesCallerSlot) { + main::Database db(":memory:", main::SystemConfig(-1u, 1)); + main::Connection conn(&db); + std::atomic active{0}, peak{0}; + auto child = std::make_shared(active, peak); + auto outer = std::make_shared(conn.getClientContext(), child); + schedule(conn, outer); + EXPECT_EQ(outer->threadID, std::this_thread::get_id()); + EXPECT_EQ(child->threadID, std::this_thread::get_id()); +}