Program Listing for File rw_spinlock.hpp

↰ Return to documentation for file (/home/runner/work/Legion-Engine/Legion-Engine/legion/engine/core/async/rw_spinlock.hpp)

#pragma once
#include <atomic>
#include <tuple>
#include <unordered_map>
#include <array>
#include <thread>
#include <functional>
#include <core/types/primitives.hpp>
#include <core/platform/platform.hpp>
#include <core/containers/sparse_set.hpp>
#include <core/detail/internals.hpp>
#include <core/async/wait_priority.hpp>

namespace legion::core::async
{
    enum struct lock_state : int { idle = 0, read = 1, write = -1 };

    inline constexpr lock_state lock_state_idle = lock_state::idle;
    inline constexpr lock_state lock_state_read = lock_state::read;
    inline constexpr lock_state lock_state_write = lock_state::write;

    struct rw_spinlock final
    {
    private:
        static bool m_forceRelease;
        static std::atomic_uint m_lastId;

        static thread_local std::unordered_map<uint, int> m_localWriters;
        static thread_local std::unordered_map<uint, int> m_localReaders;
        static thread_local std::unordered_map<uint, lock_state> m_localState;

        uint m_id = m_lastId.fetch_add(1, std::memory_order_relaxed);
        // State of the lock. -1 means that a thread has write permission. 0 means that the lock is unlocked. 1+ means that there are N amount of readers.
        mutable  std::atomic_int m_lockState = { 0 };
        mutable std::thread::id m_writer;

        void read_lock(wait_priority priority = wait_priority::real_time) const;

        bool read_try_lock() const;

        void write_lock(wait_priority priority = wait_priority::real_time) const;

        bool write_try_lock() const;

        void read_unlock() const;

        void write_unlock() const;

    public:
        static void force_release(bool release = true);

        rw_spinlock() = default;

        rw_spinlock(rw_spinlock&& source) noexcept;
        rw_spinlock& operator=(rw_spinlock&& source) noexcept;

        rw_spinlock(const rw_spinlock&) = delete;
        rw_spinlock& operator=(const rw_spinlock&) = delete;

        void lock(lock_state permissionLevel = lock_state::write, wait_priority priority = wait_priority::real_time) const;

        bool try_lock(lock_state permissionLevel = lock_state::write) const;

        void unlock(lock_state permissionLevel = lock_state::write) const;

        void lock_shared() const;

        bool try_lock_shared() const;

        void unlock_shared() const;

        template<typename Guard, typename Func>
        auto critical_section(const Func& func) const -> decltype(auto)
        {
            Guard guard(*this);
            return std::invoke(func);
        }
    };

    class readonly_guard final
    {
    private:
        const rw_spinlock& m_lock;

    public:
        readonly_guard(const rw_spinlock& lock, wait_priority priority = wait_priority::real_time) : m_lock(lock)
        {
            m_lock.lock(lock_state::read, priority);
        }

        readonly_guard(const readonly_guard&) = delete;

        ~readonly_guard()
        {
            m_lock.unlock(lock_state::read);
        }

        readonly_guard& operator=(readonly_guard&&) = delete;
    };

    template<size_type S>
    class readonly_multiguard final
    {
    private:
        std::array<const rw_spinlock*, S> m_locks;

    public:
        template<typename lock_type1 = rw_spinlock, typename lock_type2 = rw_spinlock, typename... lock_typesN>
        readonly_multiguard(const lock_type1& lock1, const lock_type2& lock2, const lock_typesN&... locks) : m_locks{ {&lock1, &lock2, &locks...} }
        {
            int lastLocked = -1; // Index to the last locked lock.

            bool locked = true;
            do
            {
                for (int i = 0; i <= lastLocked; i++) // If we failed to lock all locks we need to unlock the ones we did lock.
                    m_locks[i]->unlock(lock_state::read);

                // Reset variables
                locked = true;
                lastLocked = -1;

                // Try to lock all locks.
                for (int i = 0; i < m_locks.size(); i++)
                {
                    if (m_locks[i]->try_lock(lock_state::read))
                    {
                        lastLocked = i;
                    }
                    else
                    {
                        locked = false;
                        break;
                    }
                }
            } while (!locked);
        }

        readonly_multiguard(const readonly_multiguard&) = delete;

        ~readonly_multiguard()
        {
            for (auto* lock : m_locks)
                lock->unlock(lock_state::read);
        }

        readonly_multiguard& operator=(readonly_multiguard&&) = delete;
    };

#if !defined(DOXY_EXCLUDE)
    template<typename... types>
    readonly_multiguard(types...)->readonly_multiguard<sizeof...(types)>;
#endif

    class readwrite_guard final
    {
    private:
        const rw_spinlock& m_lock;

    public:
        readwrite_guard(const rw_spinlock& lock, wait_priority priority = wait_priority::real_time) : m_lock(lock)
        {
            m_lock.lock(lock_state::write, priority);
        }

        readwrite_guard(const readwrite_guard&) = delete;

        ~readwrite_guard()
        {
            m_lock.unlock(lock_state::write);
        }

        readwrite_guard& operator=(readwrite_guard&&) = delete;
    };


    template<size_type S>
    class readwrite_multiguard final
    {
    private:
        std::array<const rw_spinlock*, S> m_locks;

    public:
        template<typename lock_type1, typename lock_type2, typename... lock_typesN>
        readwrite_multiguard(const lock_type1& lock1, const lock_type2& lock2, const lock_typesN&... locks) : m_locks{ {&lock1, &lock2, &locks...} }
        {
            int lastLocked = -1; // Index to the last locked lock.

            bool locked = true;
            do
            {
                for (int i = 0; i <= lastLocked; i++) // If we failed to lock all locks we need to unlock the ones we did lock.
                    m_locks[i]->unlock(lock_state::write);

                // Reset variables
                locked = true;
                lastLocked = -1;

                // Try to lock all locks.
                for (int i = 0; i < m_locks.size(); i++)
                {
                    if (m_locks[i]->try_lock(lock_state::write))
                    {
                        lastLocked = i;
                    }
                    else
                    {
                        locked = false;
                        break;
                    }
                }
            } while (!locked);
        }

        readwrite_multiguard(const readwrite_multiguard&) = delete;

        ~readwrite_multiguard()
        {
            for (auto* lock : m_locks)
                lock->unlock(lock_state::write);
        }

        readwrite_multiguard& operator=(readwrite_multiguard&&) = delete;
    };

#if !defined(DOXY_EXCLUDE)
    template<typename... types>
    readwrite_multiguard(types...)->readwrite_multiguard<sizeof...(types)>;
#endif

    template<size_type S>
    class mixed_multiguard final
    {
    private:
        std::array<const rw_spinlock*, S / 2> m_locks;
        std::array<lock_state, S / 2> m_states;

        // Recursive function for filling the arrays with the neccessary data from the template arguments.
        template<size_type I, typename... types>
        void fill(const rw_spinlock& lock, lock_state state, types&&... args)
        {
            if constexpr (I > 2)
            {
                fill<I - 2>(args...);
            }

            m_locks[(I / 2) - 1] = &lock;
            m_states[(I / 2) - 1] = state;
        }

    public:
        template<typename... types>
        explicit mixed_multiguard(types&&... arguments)
        {
            static_assert(sizeof...(types) % 2 == 0, "Argument order should be (lock, lock-state, lock, lock-state). Argument count should thus be even.");

            fill<sizeof...(types)>(arguments...);

            int lastLocked = -1; // Index to the last locked lock.

            bool locked = true;
            do
            {
                for (int i = 0; i <= lastLocked; i++) // If we failed to lock all locks we need to unlock the ones we did lock.
                    m_locks[i]->unlock(m_states[i]);

                // Reset variables
                locked = true;
                lastLocked = -1;

                // Try to lock all locks.
                for (int i = 0; i < m_locks.size(); i++)
                {
                    if (m_locks[i]->try_lock(m_states[i]))
                    {
                        lastLocked = i;
                    }
                    else
                    {
                        locked = false;
                        break;
                    }
                }
            } while (!locked);
        }

        mixed_multiguard(const mixed_multiguard&) = delete;
        mixed_multiguard(mixed_multiguard&&) = delete;

        ~mixed_multiguard()
        {
            for (int i = 0; i < m_locks.size(); i++)
                m_locks[i]->unlock(m_states[i]);
        }

        mixed_multiguard& operator=(mixed_multiguard&&) = delete;
        mixed_multiguard& operator=(const mixed_multiguard&) = delete;
    };

#if !defined(DOXY_EXCLUDE)
    // CTAD so you don't need to input the size of the guard parameters.
    template<typename... types>
    mixed_multiguard(types...)->mixed_multiguard<sizeof...(types)>;
#endif
}