feat(threadpool): blocking parallelFor with work stealing

Removes the need for a dedicated JXL thread pool

Author: Tom94

include/tev/Task.h

@@ -20,6 +20,7 @@
 
 #include <tev/Common.h>
 
+#include <chrono>
 #include <coroutine>
 #include <future>
 #include <ranges>
@@ -200,6 +201,10 @@ template <typename T> class Task {
         return await_resume();
     }
 
+    bool done() const noexcept { return mState->latch.value() <= (mState->continuation ? 0 : 1); }
+
+    std::future_status wait_for(const std::chrono::microseconds& duration) const noexcept { return mFuture.wait_for(duration); }
+
     bool await_suspend(std::coroutine_handle<> coroutine) noexcept {
         if (!mHandle) {
             tlog::error("Cannot co_await/get() a task multiple times.");

include/tev/ThreadPool.h

@@ -112,6 +112,26 @@ class ThreadPool {
     void shutdownThreads(size_t num);
     void shutdown();
 
+    bool tryRunOneTask();
+
+    // Block the calling thread on `task`, but keep the pool productive by running queued tasks inline instead of sleeping. Useful for
+    // sychronous APIs that need to wait for a task to finish, but don't want to block the threadpool. See `parallelForSync`.
+    template <typename T> T blockAndDrain(Task<T>&& task) {
+        while (!task.done()) {
+            if (!tryRunOneTask()) {
+                // If we've got no tasks to run, back off a bit to avoid busy waiting. No big harm done if tasks are queued for up to 1ms
+                // without being dealt with. If more tasks arrive than can be processed, the threadpool will remain busy once the 1ms are
+                // over. Waiting on the blocked task itself allows early wakeups if the task finishes while we're waiting, so we're not
+                // actually incurring 1ms latency on that front.
+                if (task.wait_for(std::chrono::milliseconds(1)) == std::future_status::ready) {
+                    break;
+                }
+            }
+        }
+
+        return task.get(); // already complete; no real blocking
+    }
+
     size_t numTasksInSystem() const { return mNumTasksInSystem; }
 
     void waitUntilFinished();
@@ -190,6 +210,15 @@ class ThreadPool {
         );
     }
 
+    template <std::integral Int, std::invocable<Int, Int> F>
+    void parallelForSync(Int start, Int end, size_t approxCost, F body, int priority) {
+        blockAndDrain(parallelFor(start, end, approxCost, body, priority));
+    }
+
+    template <std::integral Int, std::invocable<Int> F> void parallelForSync(Int start, Int end, size_t approxCost, F body, int priority) {
+        blockAndDrain(parallelFor(start, end, approxCost, body, priority));
+    }
+
     size_t numThreads() const { return mNumThreads; }
 
 private:

src/ThreadPool.cpp

@@ -92,6 +92,30 @@ void ThreadPool::startThreads(size_t num) {
     }
 }
 
+bool ThreadPool::tryRunOneTask() {
+    if (!mTaskQueueSemaphore.tryWait()) {
+        return false;
+    }
+
+    QueuedTask task;
+    {
+        const std::scoped_lock lock{mTaskQueueMutex};
+        TEV_ASSERT(!mTaskQueue.empty(), "Task queue empty after successful tryWait.");
+
+        if (mTaskQueue.top().stopToken) {
+            // Hand off stop tokens to actual thread pool threads, not draining callers
+            mTaskQueueSemaphore.signal();
+            return false;
+        }
+
+        task = mTaskQueue.pop();
+    }
+
+    task.fun();
+    --mNumTasksInSystem;
+    return true;
+}
+
 void ThreadPool::shutdownThreads(size_t num) {
     mNumThreads -= num;
     for (size_t i = 0; i < num; ++i) {

src/imageio/JxlImageLoader.cpp

@@ -33,7 +33,6 @@
 #include <jxl/thread_parallel_runner.h>
 #include <jxl/thread_parallel_runner_cxx.h>
 
-#include <istream>
 #include <limits>
 #include <span>
 #include <vector>
@@ -172,13 +171,12 @@ Task<vector<ImageData>> JxlImageLoader::load(
             // co-routine. I.e. we need to synchronously wait for the work to finish, which could deadlock the global threadpool. Other mitigation
             // strategies involve temporarily creating and removing extra threads from the global threadpool (which tev previously implemented),
             // but this approach here scales better to huge numbers of images (n-cores extra threads instead of n-images extra threads).
-            static auto jxlPool = ThreadPool();
 
             const auto* runnerDataPtr = static_cast<RunnerData*>(runnerOpaque);
 
             const uint32_t range = endRange - startRange;
             const uint32_t numTasks = std::min(
-                jxlPool.nTasks(
+                ThreadPool::global().nTasks(
                     0u,
                     range,
                     numeric_limits<uint32_t>::max() // Max parallelism up to range tasks & hardware concurrency
@@ -191,26 +189,23 @@ Task<vector<ImageData>> JxlImageLoader::load(
                 return initResult;
             }
 
-            jxlPool
-                .parallelFor(
-                    0u,
-                    numTasks,
-                    numeric_limits<uint32_t>::max(), // Maximum parallelism up to numTasks threads
-                    [&](uint32_t i) {
-                        const uint32_t taskStart = startRange + (range * i / numTasks);
-                        const uint32_t taskEnd = startRange + (range * (i + 1) / numTasks);
-                        TEV_ASSERT(taskStart != taskEnd, "Should not produce tasks with empty range.");
-
-                        for (uint32_t j = taskStart; j < taskEnd; ++j) {
-                            func(jpegxlOpaque, j, (uint32_t)i);
-                        }
-                    },
-                    runnerDataPtr->priority
-                )
-                // The synchronous parallel for loop is janky, because it doesn't follow the coroutine paradigm. But it is the only way to
-                // get the thread pool to cooperate with the JXL API that expects a non-coroutine function here. We will offload the
-                // JxlImageLoader::load() function into a wholly separate thread to avoid blocking the thread pool as a consequence.
-                .get();
+            // Since the JXL API's contract doesn't let us async/await, call the synchronous variant of parallelFor, which will block this
+            // thread by running other thread pool tasks until the parallel for is done.
+            ThreadPool::global().parallelForSync(
+                0u,
+                numTasks,
+                numeric_limits<uint32_t>::max(), // Maximum parallelism up to numTasks threads
+                [&](uint32_t i) {
+                    const uint32_t taskStart = startRange + (range * i / numTasks);
+                    const uint32_t taskEnd = startRange + (range * (i + 1) / numTasks);
+                    TEV_ASSERT(taskStart != taskEnd, "Should not produce tasks with empty range.");
+
+                    for (uint32_t j = taskStart; j < taskEnd; ++j) {
+                        func(jpegxlOpaque, j, (uint32_t)i);
+                    }
+                },
+                runnerDataPtr->priority
+            );
 
             return 0;
         };