Task.h

Go to the documentation of this file.

 
 #pragma once
 
 #include <coroutine>
 #include <exception>
 #include <memory>
 #include <optional>
 #include <vector>
 #include <mutex>
 #include <atomic>
 #include <utility>
 
 namespace swt {
 
 // Forward declaration
 template<typename T>
 class Task;
 
 template<typename T>
 class TaskPromise {
 public:
     using Handle = std::coroutine_handle<TaskPromise<T>>;
 
     Task<T> get_return_object();
 
     std::suspend_always initial_suspend() noexcept {
         return {};
     }
 
     auto final_suspend() noexcept {
         struct FinalAwaiter {
             TaskPromise* promise;
             
             bool await_ready() noexcept { return false; }
             
             std::coroutine_handle<> await_suspend(std::coroutine_handle<TaskPromise<T>> h) noexcept {
                 std::lock_guard<std::mutex> lock(promise->mContinuationMutex);
                 
                 // Resume all waiting continuations
                 for (auto& continuation : promise->mContinuations) {
                     continuation.resume();
                 }
                 promise->mContinuations.clear();
                 
                 return std::noop_coroutine();
             }
             
             void await_resume() noexcept {}
         };
         
         return FinalAwaiter{this};
     }
 
     void return_value(T value) {
         mResult = std::move(value);
         mResultReady.store(true, std::memory_order_release);
     }
 
     void unhandled_exception() {
         mException = std::current_exception();
         mResultReady.store(true, std::memory_order_release);
     }
 
     T getResult() {
         if (mException) {
             std::rethrow_exception(mException);
         }
         if (mResult.has_value()) {
             return mResult.value(); // Don't move, might be called multiple times
         }
         throw std::runtime_error("Task not completed");
     }
 
     bool isReady() const noexcept {
         return mResultReady.load(std::memory_order_acquire);
     }
 
     void addContinuation(std::coroutine_handle<> continuation) {
         std::lock_guard<std::mutex> lock(mContinuationMutex);
         
         // If already completed, resume immediately
         if (mResultReady.load(std::memory_order_acquire)) {
             continuation.resume();
         } else {
             mContinuations.push_back(continuation);
         }
     }
 
 private:
     std::optional<T> mResult;                           
     std::exception_ptr mException;                      
     std::atomic<bool> mResultReady{false};             
     
     std::mutex mContinuationMutex;                      
     std::vector<std::coroutine_handle<>> mContinuations; 
 };
 
 template<>
 class TaskPromise<void> {
 public:
     using Handle = std::coroutine_handle<TaskPromise<void>>;
 
     Task<void> get_return_object();
 
     std::suspend_always initial_suspend() noexcept {
         return {};
     }
 
     auto final_suspend() noexcept {
         struct FinalAwaiter {
             TaskPromise* promise;
             
             bool await_ready() noexcept { return false; }
             
             std::coroutine_handle<> await_suspend(std::coroutine_handle<TaskPromise<void>> h) noexcept {
                 std::lock_guard<std::mutex> lock(promise->mContinuationMutex);
                 
                 for (auto& continuation : promise->mContinuations) {
                     continuation.resume();
                 }
                 promise->mContinuations.clear();
                 
                 return std::noop_coroutine();
             }
             
             void await_resume() noexcept {}
         };
         
         return FinalAwaiter{this};
     }
 
     void return_void() {
         mResultReady.store(true, std::memory_order_release);
     }
 
     void unhandled_exception() {
         mException = std::current_exception();
         mResultReady.store(true, std::memory_order_release);
     }
 
     void getResult() {
         if (mException) {
             std::rethrow_exception(mException);
         }
     }
 
     bool isReady() const noexcept {
         return mResultReady.load(std::memory_order_acquire);
     }
 
     void addContinuation(std::coroutine_handle<> continuation) {
         std::lock_guard<std::mutex> lock(mContinuationMutex);
         
         if (mResultReady.load(std::memory_order_acquire)) {
             continuation.resume();
         } else {
             mContinuations.push_back(continuation);
         }
     }
 
 private:
     std::exception_ptr mException;                      
     std::atomic<bool> mResultReady{false};             
     std::mutex mContinuationMutex;                      
     std::vector<std::coroutine_handle<>> mContinuations; 
 };
 
 template<typename T>
 class Task {
 public:
     using promise_type = TaskPromise<T>;
     using Handle = typename promise_type::Handle;
 
     explicit Task(Handle handle) : mHandle(handle) {}
 
     ~Task() {
         if (mHandle) {
             mHandle.destroy();
         }
     }
 
     // Move-only semantics
     Task(const Task&) = delete;
     Task& operator=(const Task&) = delete;
 
     Task(Task&& other) noexcept : mHandle(std::exchange(other.mHandle, {})) {}
 
     Task& operator=(Task&& other) noexcept {
         if (this != &other) {
             if (mHandle) {
                 mHandle.destroy();
             }
             mHandle = std::exchange(other.mHandle, {});
         }
         return *this;
     }
 
     void start() {
         if (mHandle && !mHandle.done()) {
             mHandle.resume();
         }
     }
 
     bool done() const noexcept {
         return mHandle && mHandle.done();
     }
 
     auto getResult() -> T requires (!std::is_void_v<T>) {
         if (!mHandle) {
             throw std::runtime_error("Invalid task handle");
         }
         
         if (!mHandle.promise().isReady()) {
             throw std::runtime_error("Task not completed - use co_await or check done() first");
         }
         
         return mHandle.promise().getResult();
     }
 
     void getResult() requires std::is_void_v<T> {
         if (!mHandle) {
             throw std::runtime_error("Invalid task handle");
         }
         
         if (!mHandle.promise().isReady()) {
             throw std::runtime_error("Task not completed - use co_await or check done() first");
         }
         
         mHandle.promise().getResult();
     }
 
     bool await_ready() const noexcept {
         return done();
     }
 
     void await_suspend(std::coroutine_handle<> waiter) noexcept {
         if (!mHandle) {
             // Invalid handle, resume immediately
             waiter.resume();
             return;
         }
 
         // Add waiter as continuation
         mHandle.promise().addContinuation(waiter);
         
         // Start this task if not already started
         if (!mHandle.done()) {
             mHandle.resume();
         }
     }
 
     auto await_resume() -> T requires (!std::is_void_v<T>) {
         return getResult();
     }
 
     void await_resume() requires std::is_void_v<T> {
         getResult();
     }
 
     bool isReady() const noexcept {
         return mHandle && mHandle.promise().isReady();
     }
 
 private:
     Handle mHandle; 
 };
 
 // Implementations for get_return_object
 template<typename T>
 inline Task<T> TaskPromise<T>::get_return_object() {
     return Task<T>{Handle::from_promise(*this)};
 }
 
 inline Task<void> TaskPromise<void>::get_return_object() {
     return Task<void>{Handle::from_promise(*this)};
 }
 
 } // namespace swt