perf(processor): run a one-thread task on the calling thread

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>& 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>& task,
+    processor::ExecutionContext* context, bool launchNewWorkerThread, TaskBeforeRunFn beforeRun) {
+    scheduleTaskAndWaitOrError(task, context->clientContext, launchNewWorkerThread, beforeRun);
 }
 
 void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr<Task>& task,
     main::ClientContext* clientContext, bool launchNewWorkerThread) {
+    scheduleTaskAndWaitOrError(task, clientContext, launchNewWorkerThread, nullptr);
+}
+
+void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr<Task>& 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>& 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>& 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>& 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>& 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>& task,
+    main::ClientContext* clientContext, bool launchNewWorkerThread) {
+    scheduleTaskAndWaitOrError(task, clientContext, launchNewWorkerThread, nullptr);
+}
+
+void TaskScheduler::scheduleTaskAndWaitOrError(const std::shared_ptr<Task>& 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());

src/include/common/task_system/task.h

@@ -15,8 +15,9 @@ using lock_t = std::unique_lock<std::mutex>;
 /**
  * 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

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>& task,
         main::ClientContext* clientContext, bool launchNewWorkerThread = false);
+    void scheduleTaskAndWaitOrError(const std::shared_ptr<Task>& task,
+        processor::ExecutionContext* context, bool launchNewWorkerThread,
+        TaskBeforeRunFn beforeRun);
+    void scheduleTaskAndWaitOrError(const std::shared_ptr<Task>& 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<ScheduledTask> getTaskAndRegister();
     static void runTask(Task* task);
 
+    bool runOnCallingThread(Task* task);
+
 private:
     std::deque<std::shared_ptr<ScheduledTask>> 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>& task,
         main::ClientContext* clientContext, bool launchNewWorkerThread = false);
+    void scheduleTaskAndWaitOrError(const std::shared_ptr<Task>& task,
+        processor::ExecutionContext* context, bool launchNewWorkerThread,
+        TaskBeforeRunFn beforeRun);
+    void scheduleTaskAndWaitOrError(const std::shared_ptr<Task>& task,
+        main::ClientContext* clientContext, bool launchNewWorkerThread, TaskBeforeRunFn beforeRun);
 
     static TaskScheduler* Get(const main::ClientContext& context);
 

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;

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<ProcessorTask*>(task)) {
+        processorTask->narrowToSingleMorselSource();
+    }
+}
+
+} // namespace
+
 std::shared_ptr<ResultSet> ResultSetPool::getOrCreate(ResultSetDescriptor* descriptor,
     MemoryManager* memoryManager) {
     const auto threadID = std::this_thread::get_id();
@@ -75,7 +85,7 @@ std::unique_ptr<main::QueryResult> 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();
 }

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<uint64_t>()},
       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<TableFunctionCall>();
+    if (call->getInfo().function.name != FTableScan::name || call->getInfo().bindData == nullptr) {
+        return;
+    }
+    auto* bindData = call->getInfo().bindData->constPtrCast<FTableScanBindData>();
+    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

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)