Resolve a number of thread-safety/scheduling issues

Author: jeremyong

ARCHITECTURE.md

@@ -19,7 +19,11 @@ That is, it doesn't get requeued on the thread pool for later execution.
 
 The concurrent queue used to push work to worker threads is provided by [`moodycamel::ConcurrentQueue`](https://github.com/cameron314/concurrentqueue).
 Under the hood, the queue provides multiple-consumer multiple-producer usage, although in this case, only a single producer per queue
-exists.
+exists. The thread pool worker threads currently do *not* support work stealing, which is a slightly more complicated endeavor
+for job schedulers that support task affinity.
+
+The granularity of your jobs shouldn't be too fine - maybe having jobs that are at least 100 us or more is a good idea, or you'll
+end up paying disproportionately for scheduling costs.
 
 The Win32 event awaiter works by having a single IO thread which blocks in a single `WaitForMultipleObjects` call. One of the
 events it waits on is used to signal the available of more events to wait on. All the other events waited on are user awaited.

include/coop/detail/tracer.hpp

@@ -1,5 +1,19 @@
 #pragma once
 
+#include <cstdint>
+#include <thread>
+
+namespace coop
+{
+namespace detail
+{
+    inline size_t thread_id() noexcept
+    {
+        return std::hash<std::thread::id>{}(std::this_thread::get_id());
+    }
+} // namespace detail
+} // namespace coop
+
 #if defined(COOP_TRACE) && !defined(NDEBUG)
 #    include <cstdio>
 

include/coop/task.hpp

@@ -15,8 +15,7 @@ using experimental::suspend_never;
 #endif
 #include <cstdlib>
 #include <limits>
-#include <semaphore>
-#include <unordered_map>
+#include <mutex>
 
 namespace coop
 {
@@ -56,6 +55,9 @@ struct promise_base_t
     // resume point, which immediately following the suspend point.
     std::coroutine_handle<> continuation = nullptr;
 
+    std::mutex mutex;
+    bool flag = false;
+
     // Do not suspend immediately on entry of a coroutine
     std::suspend_never initial_suspend() const noexcept
     {
@@ -80,21 +82,84 @@ struct final_awaiter_t
     {
     }
 
-    std::coroutine_handle<> await_suspend(std::coroutine_handle<P> coroutine) const noexcept
+    std::coroutine_handle<>
+    await_suspend(std::coroutine_handle<P> coroutine) const noexcept
     {
         // Check if this coroutine is being finalized from the
         // middle of a "continuation" coroutine and hop back there to
         // continue execution while *this* coroutine is suspended.
 
-        auto continuation = coroutine.promise().continuation;
-        if (continuation)
+        if constexpr (P::joinable_v)
+        {
+            // Joinable tasks are never awaited and so cannot have a
+            // continuation by definition
+            return std::noop_coroutine();
+        }
+        else
         {
-            return continuation;
+            COOP_LOG("Final await for coroutine %p on thread %zu\n",
+                     coroutine.address(),
+                     detail::thread_id());
+            std::scoped_lock lock{coroutine.promise().mutex};
+            if (coroutine.promise().flag)
+            {
+                // We're not the first to reach here
+                auto continuation = coroutine.promise().continuation;
+                if (continuation)
+                {
+                    COOP_LOG("Resuming continuation %p on %p on thread %zu\n",
+                             continuation.address(),
+                             coroutine.address(),
+                             detail::thread_id());
+                    return continuation;
+                }
+                else
+                {
+                    COOP_LOG(
+                        "Coroutine %p on thread %zu missing continuation\n",
+                        coroutine.address(),
+                        detail::thread_id());
+                }
+            }
+            coroutine.promise().flag = true;
+            return std::noop_coroutine();
         }
-        return std::noop_coroutine();
     }
 };
 
+namespace detail
+{
+    // Helper function for awaiting on a task. The next resume point is
+    // installed as a continuation of the task being awaited.
+    template <typename P>
+    std::coroutine_handle<>
+    await_suspend(std::coroutine_handle<P> base, std::coroutine_handle<> next)
+    {
+        if constexpr (P::joinable_v)
+        {
+            // Joinable tasks are never awaited and so cannot have a
+            // continuation by definition
+            return std::noop_coroutine();
+        }
+        else
+        {
+            std::scoped_lock lock{base.promise().mutex};
+            if (!base.promise().flag)
+            {
+                // We're the first to reach here
+                base.promise().flag = true;
+                COOP_LOG("Installing continuation %p for %p on thread %zu\n",
+                         next.address(),
+                         base.address(),
+                         detail::thread_id());
+                base.promise().continuation = next;
+                return std::noop_coroutine();
+            }
+            return base;
+        }
+    }
+} // namespace detail
+
 template <bool Joinable>
 class task_base_t
 {
@@ -108,6 +173,8 @@ template <>
 class task_base_t<true>
 {
 public:
+    constexpr static bool joinable_v = true;
+
     void join()
     {
         join_sem_->acquire();
@@ -133,6 +200,8 @@ class task_t final : public task_base_t<Joinable>
 public:
     struct promise_type : public task_base_t<Joinable>::promise_type
     {
+        constexpr static bool joinable_v = Joinable;
+
         T data;
 
         static void* operator new(size_t size)
@@ -190,7 +259,6 @@ class task_t final : public task_base_t<Joinable>
     task_t(std::coroutine_handle<promise_type> coroutine) noexcept
         : coroutine_{coroutine}
     {
-        COOP_LOG("task %p born\n", coroutine_.address());
     }
     task_t(task_t const&) = delete;
     task_t& operator=(task_t const&) = delete;
@@ -238,7 +306,6 @@ class task_t final : public task_base_t<Joinable>
     {
         if (coroutine_)
         {
-            COOP_LOG("task %p dying\n", coroutine_.address());
             coroutine_.destroy();
         }
     }
@@ -275,10 +342,9 @@ class task_t final : public task_base_t<Joinable>
 
     // When suspending from a coroutine *within* a task's coroutine, save the
     // resume point (to be resumed when the inner coroutine finalizes)
-    void
-    await_suspend(std::coroutine_handle<> coroutine) noexcept
+    std::coroutine_handle<> await_suspend(std::coroutine_handle<> coroutine) noexcept
     {
-        coroutine_.promise().continuation = coroutine;
+        return detail::await_suspend(coroutine_, coroutine);
     }
 
 private:
@@ -293,6 +359,8 @@ class task_t<void, Joinable, C> : public task_base_t<Joinable>
 public:
     struct promise_type : public task_base_t<Joinable>::promise_type
     {
+        constexpr static bool joinable_v = Joinable;
+
         static void* operator new(size_t size)
         {
             return C::alloc(size);
@@ -334,7 +402,6 @@ class task_t<void, Joinable, C> : public task_base_t<Joinable>
     task_t(std::coroutine_handle<promise_type> coroutine) noexcept
         : coroutine_{coroutine}
     {
-        COOP_LOG("task %p born\n", coroutine_.address());
     }
     task_t(task_t const&) = delete;
     task_t& operator=(task_t const&) = delete;
@@ -382,7 +449,6 @@ class task_t<void, Joinable, C> : public task_base_t<Joinable>
     {
         if (coroutine_)
         {
-            COOP_LOG("task %p dying\n", coroutine_.address());
             coroutine_.destroy();
         }
     }
@@ -404,10 +470,9 @@ class task_t<void, Joinable, C> : public task_base_t<Joinable>
 
     // When suspending from a coroutine *within* this task's coroutine, save the
     // resume point (to be resumed when the inner coroutine finalizes)
-    void
-    await_suspend(std::coroutine_handle<> coroutine) noexcept
+    std::coroutine_handle<> await_suspend(std::coroutine_handle<> coroutine) noexcept
     {
-        coroutine_.promise().continuation = coroutine;
+        return detail::await_suspend(coroutine_, coroutine);
     }
 
 private:
@@ -452,9 +517,6 @@ inline auto suspend(S& scheduler                             = S::instance(),
         }
     };
 
-    COOP_LOG("Suspending coroutine from thread %zu\n",
-             std::hash<std::thread::id>{}(std::this_thread::get_id()));
-
     return awaiter_t{scheduler, cpu_mask, priority, source_location};
 }
 

src/scheduler.cpp

@@ -181,7 +181,16 @@ void scheduler_t::schedule(std::coroutine_handle<> coroutine,
     // discrepancy (Kronecker recurrence sequence)
     uint32_t index = static_cast<uint32_t>(update_++ * std::numbers::phi_v<float>)
                      % std::popcount(cpu_affinity);
-    queues_[index].enqueue(coroutine, priority, source_location);
+
+    // Iteratively unset bits to determine the nth set bit
+    for (uint32_t i = 0; i != index; ++i)
+    {
+        cpu_affinity &= ~(1 << (std::countr_zero(cpu_affinity) + 1));
+    }
+    uint32_t queue = std::countr_zero(cpu_affinity);
+    COOP_LOG("Work queue %i identified\n", queue);
+
+    queues_[queue].enqueue(coroutine, priority, source_location);
 }
 
 void scheduler_t::schedule(std::coroutine_handle<> coroutine,

src/work_queue.cpp

@@ -26,7 +26,7 @@ work_queue_t::work_queue_t(scheduler_t& scheduler, uint32_t id)
     thread_ = std::thread([this] {
 #if defined(_WIN32)
         SetThreadAffinityMask(
-            thread_.native_handle(), static_cast<uint32_t>(1ull << id_));
+            GetCurrentThread(), static_cast<uint32_t>(1ull << id_));
 #elif defined(__linux__)
         // TODO: Android implementation
         pthread_t thread = pthread_self();
@@ -64,12 +64,10 @@ work_queue_t::work_queue_t(scheduler_t& scheduler, uint32_t id)
                     std::coroutine_handle<> coroutine;
                     if (queues_[i].try_dequeue(coroutine))
                     {
-                        COOP_LOG(
-                            "Dequeueing coroutine on CPU %i thread %zu %i\n",
-                            id_,
-                            std::hash<std::thread::id>{}(
-                                std::this_thread::get_id()),
-                            coroutine.done());
+                        COOP_LOG("Dequeueing coroutine %p on thread %zu (%i)\n",
+                                 coroutine.address(),
+                                 detail::thread_id(),
+                                 id_);
                         did_dequeue = true;
                         coroutine.resume();
                         break;
@@ -93,9 +91,10 @@ void work_queue_t::enqueue(std::coroutine_handle<> coroutine,
                            uint32_t priority,
                            source_location_t source_location)
 {
-    priority = std::clamp<uint32_t>(priority, 0, COOP_PRIORITY_COUNT);
-    COOP_LOG("Enqueueing coroutine on CPU %i (%s:%zu)\n",
-             id_,
+    priority = std::clamp<uint32_t>(priority, 0, COOP_PRIORITY_COUNT - 1);
+    COOP_LOG("Enqueueing coroutine %p on thread %zu (%s:%zu)\n",
+             coroutine.address(),
+             detail::thread_id(),
              source_location.file,
              source_location.line);
     queues_[priority].enqueue(coroutine);

test/test.cpp

@@ -5,16 +5,13 @@
 #include <coop/task.hpp>
 #include <thread>
 
-coop::task_t<> suspend_test1()
-{
-    COOP_SUSPEND();
-}
-
-coop::task_t<void, true> suspend_test2()
+coop::task_t<void, true> suspend_time()
 {
+    // std::printf("%zu start thread\n", coop::detail::thread_id());
     auto t1 = std::chrono::system_clock::now();
     COOP_SUSPEND();
     auto t2 = std::chrono::system_clock::now();
+    // std::printf("%zu end thread\n", coop::detail::thread_id());
     size_t us
         = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();
     std::printf("Duration for suspend test: %zu us\n", us);
@@ -23,7 +20,9 @@ coop::task_t<void, true> suspend_test2()
 TEST_CASE("suspend overhead")
 {
     std::printf("Calling suspend_test2 coroutine\n");
-    suspend_test2().join();
+    // auto task = suspend_time();
+    // task.join();
+    suspend_time().join();
     std::printf("suspend_test2 joined\n");
 }
 
@@ -46,25 +45,27 @@ TEST_CASE("test suspend")
     CHECK(id != next);
 }
 
-coop::task_t<int> chain1()
+coop::task_t<int> chain1(int core)
 {
     std::printf("chain1 suspending\n");
-    COOP_SUSPEND();
+    COOP_SUSPEND4(1 << core);
     std::printf("chain1 resumed\n");
     co_return 1;
 }
 
 coop::task_t<int> chain2()
 {
     std::printf("chain2\n");
-    COOP_SUSPEND();
-    co_return co_await chain1();
+    COOP_SUSPEND4(1 << 3);
+    auto t1 = chain1(5);
+    auto t2 = chain1(6);
+    co_return co_await t1 + co_await t2;
 }
 
 coop::task_t<void, true> chain3(int& result)
 {
     std::printf("chain3 suspending\n");
-    co_await coop::suspend(coop::scheduler_t::instance(), 0x2);
+    COOP_SUSPEND4(1 << 4);
     std::printf("chain3 resumed\n");
     result = co_await chain2();
 }
@@ -76,7 +77,7 @@ TEST_CASE("chained continuation")
     std::printf("Joining chained continuation task\n");
     task.join();
     std::printf("Task chained continuation joined\n");
-    CHECK(x == 1);
+    CHECK(x == 2);
 }
 
 coop::task_t<> in_flight1()