atomic_wait.h

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#pragma once
 
#include <atomic>
#include <mutex>
#include <condition_variable>
#include <thread>
#include <chrono>
#include <type_traits>
 
#ifdef USE_STD_CONCEPTS
#include <concepts>
#endif
 
#ifdef HAS_STD_ATOMIC_WAIT
// C++20: std::atomic already has wait/notify
// No additional implementation needed
#else
 
namespace kcenon::thread {
 
template<typename T>
class atomic_wait_helper {
    static_assert(std::is_trivially_copyable<T>::value,
                  "T must be trivially copyable for atomic operations");
 
public:
    void wait(std::atomic<T>& atomic, T old) noexcept {
        // Phase 1: Short spin-wait
        // Rationale: Many atomic operations complete quickly, avoiding syscall overhead
        constexpr int SPIN_COUNT = 40;
        for (int i = 0; i < SPIN_COUNT; ++i) {
            if (atomic.load(std::memory_order_acquire) != old) {
                return;
            }
            // Hint to CPU: this is a spin-wait loop
            #if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)
                __builtin_ia32_pause();
            #elif defined(__aarch64__) || defined(_M_ARM64)
                __asm__ __volatile__("yield" ::: "memory");
            #else
                std::this_thread::yield();
            #endif
        }
 
        // Phase 2: Exponential backoff with brief sleeps
        // Rationale: Longer operations benefit from yielding CPU
        constexpr int BACKOFF_COUNT = 5;
        auto backoff_time = std::chrono::microseconds(1);
 
        for (int i = 0; i < BACKOFF_COUNT; ++i) {
            if (atomic.load(std::memory_order_acquire) != old) {
                return;
            }
            std::this_thread::sleep_for(backoff_time);
            backoff_time *= 2; // Exponential backoff
        }
 
        // Phase 3: Blocking wait
        // Rationale: Value hasn't changed after spin+backoff, use efficient blocking
        std::unique_lock<std::mutex> lock(mutex_);
        cv_.wait(lock, [&] {
            return atomic.load(std::memory_order_acquire) != old;
        });
    }
 
    void notify_one() noexcept {
        // Lock-free fast path: if no waiters, skip mutex
        // Note: This is an optimization, correctness doesn't depend on it
        std::lock_guard<std::mutex> lock(mutex_);
        cv_.notify_one();
    }
 
    void notify_all() noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        cv_.notify_all();
    }
 
private:
    mutable std::mutex mutex_;
    mutable std::condition_variable cv_;
};
 
template<typename T>
class atomic_with_wait {
public:
    atomic_with_wait() noexcept : value_(T{}) {}
    explicit atomic_with_wait(T initial) noexcept : value_(initial) {}
 
    // Disable copy/move (consistent with std::atomic)
    atomic_with_wait(const atomic_with_wait&) = delete;
    atomic_with_wait& operator=(const atomic_with_wait&) = delete;
 
    // Basic atomic operations
    T load(std::memory_order order = std::memory_order_seq_cst) const noexcept {
        return value_.load(order);
    }
 
    void store(T desired, std::memory_order order = std::memory_order_seq_cst) noexcept {
        value_.store(desired, order);
    }
 
    T exchange(T desired, std::memory_order order = std::memory_order_seq_cst) noexcept {
        return value_.exchange(desired, order);
    }
 
    bool compare_exchange_weak(T& expected, T desired,
                               std::memory_order order = std::memory_order_seq_cst) noexcept {
        return value_.compare_exchange_weak(expected, desired, order);
    }
 
    bool compare_exchange_strong(T& expected, T desired,
                                 std::memory_order order = std::memory_order_seq_cst) noexcept {
        return value_.compare_exchange_strong(expected, desired, order);
    }
 
    // Wait/notify operations
    void wait(T old, std::memory_order order = std::memory_order_seq_cst) noexcept {
        waiter_.wait(value_, old);
    }
 
    void notify_one() noexcept {
        waiter_.notify_one();
    }
 
    void notify_all() noexcept {
        waiter_.notify_all();
    }
 
    // Conversion operator
    operator T() const noexcept {
        return load();
    }
 
    // Assignment operator
    T operator=(T desired) noexcept {
        store(desired);
        return desired;
    }
 
    // Atomic arithmetic operations (for integral types)
#ifdef USE_STD_CONCEPTS
    T fetch_add(T arg, std::memory_order order = std::memory_order_seq_cst) noexcept
        requires std::integral<T>
    {
        return value_.fetch_add(arg, order);
    }
 
    T fetch_sub(T arg, std::memory_order order = std::memory_order_seq_cst) noexcept
        requires std::integral<T>
    {
        return value_.fetch_sub(arg, order);
    }
 
    T operator++() noexcept
        requires std::integral<T>
    {
        return fetch_add(1) + 1;
    }
 
    T operator++(int) noexcept
        requires std::integral<T>
    {
        return fetch_add(1);
    }
 
    T operator--() noexcept
        requires std::integral<T>
    {
        return fetch_sub(1) - 1;
    }
 
    T operator--(int) noexcept
        requires std::integral<T>
    {
        return fetch_sub(1);
    }
#else
    // C++17 fallback using SFINAE
    template<typename U = T>
    typename std::enable_if<std::is_integral<U>::value, U>::type
    fetch_add(T arg, std::memory_order order = std::memory_order_seq_cst) noexcept {
        return value_.fetch_add(arg, order);
    }
 
    template<typename U = T>
    typename std::enable_if<std::is_integral<U>::value, U>::type
    fetch_sub(T arg, std::memory_order order = std::memory_order_seq_cst) noexcept {
        return value_.fetch_sub(arg, order);
    }
 
    template<typename U = T>
    typename std::enable_if<std::is_integral<U>::value, U>::type
    operator++() noexcept {
        return fetch_add(1) + 1;
    }
 
    template<typename U = T>
    typename std::enable_if<std::is_integral<U>::value, U>::type
    operator++(int) noexcept {
        return fetch_add(1);
    }
 
    template<typename U = T>
    typename std::enable_if<std::is_integral<U>::value, U>::type
    operator--() noexcept {
        return fetch_sub(1) - 1;
    }
 
    template<typename U = T>
    typename std::enable_if<std::is_integral<U>::value, U>::type
    operator--(int) noexcept {
        return fetch_sub(1);
    }
#endif
 
private:
    std::atomic<T> value_;
    atomic_wait_helper<T> waiter_;
};
 
} // namespace kcenon::thread
 
#endif // HAS_STD_ATOMIC_WAIT