synchronization.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 <mutex>
#include <condition_variable>
#include <atomic>
#include <functional>
 
#ifdef HAS_STD_LATCH
#include <latch>
#include <barrier>
#endif
 
namespace kcenon::thread {
 
#ifdef HAS_STD_LATCH
    // C++20: Use standard library implementations
    using std::latch;
    using std::barrier;
#else
    class latch {
    public:
        explicit latch(std::ptrdiff_t count)
            : count_(count) {}
 
        void count_down(std::ptrdiff_t n = 1) {
            std::unique_lock<std::mutex> lock(mutex_);
            count_ -= n;
            if (count_ <= 0) {
                cv_.notify_all();
            }
        }
 
        bool try_wait() const noexcept {
            std::unique_lock<std::mutex> lock(mutex_);
            return count_ <= 0;
        }
 
        void wait() const {
            std::unique_lock<std::mutex> lock(mutex_);
            cv_.wait(lock, [this] { return count_ <= 0; });
        }
 
        void arrive_and_wait(std::ptrdiff_t n = 1) {
            count_down(n);
            wait();
        }
 
    private:
        mutable std::mutex mutex_;
        mutable std::condition_variable cv_;
        std::ptrdiff_t count_;
    };
 
    template<typename CompletionFunction = std::function<void()>>
    class barrier {
    public:
        explicit barrier(std::ptrdiff_t count,
                        CompletionFunction completion = [](){})
            : threshold_(count)
            , count_(count)
            , generation_(0)
            , completion_(std::move(completion)) {}
 
        void arrive_and_wait() {
            std::unique_lock<std::mutex> lock(mutex_);
            auto gen = generation_;
 
            if (--count_ == 0) {
                // Last thread to arrive
                generation_++;
                count_ = threshold_;
 
                // Call completion function (if provided)
                if constexpr (!std::is_same_v<CompletionFunction, std::function<void()>>) {
                    completion_();
                } else if (completion_) {
                    completion_();
                }
 
                cv_.notify_all();
            } else {
                // Wait for all threads to arrive
                cv_.wait(lock, [this, gen] { return gen != generation_; });
            }
        }
 
        void arrive_and_drop() {
            std::unique_lock<std::mutex> lock(mutex_);
            --threshold_;
            arrive_and_wait();
        }
 
    private:
        mutable std::mutex mutex_;
        mutable std::condition_variable cv_;
        std::ptrdiff_t threshold_;  // Total number of threads
        std::ptrdiff_t count_;       // Remaining threads in current phase
        std::size_t generation_;     // Current phase number
        CompletionFunction completion_;
    };
 
#endif // HAS_STD_LATCH
 
} // namespace kcenon::thread