Semaphore.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <atomic>
#include <coroutine>
#include <cstddef>
 
namespace corium::async {
 
class AsyncSemaphore {
public:
    explicit constexpr AsyncSemaphore(ptrdiff_t initialCount = 1) noexcept
        : m_count(initialCount)
    {}
 
    ~AsyncSemaphore() = default;
    AsyncSemaphore(const AsyncSemaphore&) = delete;
    AsyncSemaphore& operator=(const AsyncSemaphore&) = delete;
 
    bool tryAcquire() noexcept {
        ptrdiff_t current = m_count.load(std::memory_order_relaxed);
        while (current > 0) {
            if (m_count.compare_exchange_weak(current, current - 1,
                                               std::memory_order_acquire,
                                               std::memory_order_relaxed)) {
                return true;
            }
        }
        return false;
    }
 
    void release(ptrdiff_t update = 1) noexcept {
        m_count.fetch_add(update, std::memory_order_release);
        auto h = m_waiter.exchange(nullptr, std::memory_order_acq_rel);
        if (h && !h.done()) {
            h.resume();
        }
    }
 
    [[nodiscard]] ptrdiff_t available() const noexcept {
        return m_count.load(std::memory_order_relaxed);
    }
 
    struct AcquireAwaiter {
        AsyncSemaphore& sem;
 
        [[nodiscard]] bool await_ready() const noexcept {
            return sem.tryAcquire();
        }
 
        bool await_suspend(std::coroutine_handle<> h) noexcept {
            if (sem.tryAcquire()) {
                return false;
            }
            sem.m_waiter.store(h, std::memory_order_release);
            return !sem.tryAcquire();
        }
 
        constexpr void await_resume() const noexcept {}
    };
 
    [[nodiscard]] AcquireAwaiter acquire() noexcept {
        return AcquireAwaiter{*this};
    }
 
private:
    std::atomic<ptrdiff_t> m_count{1};
    std::atomic<std::coroutine_handle<>> m_waiter{nullptr};
};
 
} // namespace corium::async