AsyncMutex.hpp

Go to the documentation of this file.

//
//  CeresEngine - A game development framework
//
//  Created by Rogiel Sulzbach.
//  Copyright (c) 2018-2022 Rogiel Sulzbach. All rights reserved.
//
 
#pragma once
 
#include "ExecutionContext.hpp"
#include "Threading.hpp"
 
#include "CeresEngine/DataTypes.hpp"
#include "CeresEngine/Macros.hpp"
 
#include "CeresEngine/Foundation/Coroutine.hpp"
#include "CeresEngine/Foundation/Container/Atomic.hpp"
 
namespace CeresEngine {
 
    template<typename ExecutorType = InlineExecutor> class TAsyncMutex {
    public:
        class Lock;
        class LockOperation;
        class ScopedLockOperation;
 
    private:
        friend class LockOperation;
 
        ExecutorType mExecutor;
 
        static constexpr std::uintptr_t kNotLocked = 1;
        static constexpr std::uintptr_t kLockedNoWaiters = 0;
 
        // NOTE: assumes == reinterpret_cast<std::uintptr_t>(static_cast<void*>(nullptr))
 
        Atomic<std::uintptr_t> mState = kNotLocked;
 
        LockOperation* mWaiters = nullptr;
 
    public:
        TAsyncMutex() noexcept = default;
 
        explicit TAsyncMutex(ExecutorType&& executor) noexcept : mExecutor(std::forward<ExecutorType>(executor)) {}
 
        ~TAsyncMutex() noexcept {
            [[maybe_unused]] const std::uintptr_t state = mState.load(std::memory_order_relaxed);
            CE_ASSERT(state == kNotLocked || state == kLockedNoWaiters);
            CE_ASSERT(mWaiters == nullptr);
        }
 
        [[nodiscard]] bool tryLock() noexcept {
            // Try to atomically transition from nullptr (not-locked) -> this (locked-no-waiters).
            std::uintptr_t oldState = kNotLocked;
            return mState.compare_exchange_strong(oldState, kLockedNoWaiters, std::memory_order_acquire, std::memory_order_relaxed);
        }
 
        [[nodiscard]] bool try_lock() noexcept { return tryLock(); }
 
        [[nodiscard]] LockOperation lock() noexcept { return ScopedLockOperation{*this}; }
 
        [[nodiscard]] ScopedLockOperation scopedLock() noexcept { return ScopedLockOperation{*this}; }
 
        void unlock() {
            CE_ASSERT(mState.load(std::memory_order_relaxed) != kNotLocked);
 
            LockOperation* waitersHead = mWaiters;
            if(waitersHead == nullptr) {
                std::uintptr_t oldState = kLockedNoWaiters;
                const bool releasedLock = mState.compare_exchange_strong(oldState, kNotLocked, std::memory_order_release, std::memory_order_relaxed);
                if(releasedLock) {
                    return;
                }
 
                // At least one new waiter.
                // Acquire the list of new waiter operations atomically.
                oldState = mState.exchange(kLockedNoWaiters, std::memory_order_acquire);
                CE_ASSERT(oldState != kLockedNoWaiters && oldState != kNotLocked);
 
                // Transfer the list to mWaiters, reversing the list in the process so
                // that the head of the list is the first to be resumed.
                LockOperation* next = reinterpret_cast<LockOperation*>(oldState);
                do {
                    LockOperation* const temporaryNext = next->mNext;
                    next->mNext = waitersHead;
                    waitersHead = next;
                    next = temporaryNext;
                } while(next != nullptr);
            }
            CE_ASSERT(waitersHead != nullptr);
 
            mWaiters = waitersHead->mNext;
 
            // Resume the waiter.
            // This will pass the ownership of the lock on to that operation/coroutine.
            dispatch(mExecutor, [awaiter = waitersHead->mAwaiter]() mutable { awaiter.resume(); });
        }
 
    public:
        class Lock {
        public:
            explicit Lock(TAsyncMutex& mutex, std::adopt_lock_t) noexcept : mMutex(&mutex) {}
 
            Lock(Lock&& other) noexcept : mMutex(std::exchange(other.mMutex, nullptr)) {}
            Lock& operator=(Lock&& other) noexcept {
                release();
                mMutex = std::exchange(mMutex, nullptr);
                return *this;
            }
 
            Lock(const Lock& other) = delete;
            Lock& operator=(const Lock& other) = delete;
 
            // Releases the lock.
            ~Lock() { release(); }
 
        public:
            inline void unlock() { release(); }
 
        private:
            void release() {
                if(mMutex != nullptr) {
                    mMutex->unlock();
                }
            }
 
            TAsyncMutex* mMutex;
        };
 
        class LockOperation {
            friend class TAsyncMutex;
 
        private:
            LockOperation* mNext;
 
            CoroutineHandle<> mAwaiter;
 
        protected:
            TAsyncMutex& mMutex;
 
        public:
            explicit LockOperation(TAsyncMutex& mutex) noexcept : mMutex(mutex) {}
 
        public: // Coroutine interface
            [[nodiscard]] bool await_ready() const noexcept { return false; }
            [[nodiscard]] bool await_suspend(const CoroutineHandle<> awaiter) noexcept {
                mAwaiter = awaiter;
 
                std::uintptr_t oldState = mMutex.mState.load(std::memory_order_acquire);
                while(true) {
                    if(oldState == TAsyncMutex::kNotLocked) {
                        if(mMutex.mState.compare_exchange_weak(oldState, TAsyncMutex::kLockedNoWaiters, std::memory_order_acquire, std::memory_order_relaxed)) {
                            // Acquired lock, don't suspend.
                            return false;
                        }
                    } else {
                        // Try to push this operation onto the head of the waiter stack.
                        mNext = reinterpret_cast<LockOperation*>(oldState);
                        if(mMutex.mState.compare_exchange_weak(oldState, reinterpret_cast<std::uintptr_t>(this), std::memory_order_release, std::memory_order_relaxed)) {
                            // Queued operation to waiters list, suspend now.
                            return true;
                        }
                    }
                }
            }
            void await_resume() const noexcept {}
        };
 
        class ScopedLockOperation : public LockOperation {
        public:
            using LockOperation::LockOperation;
 
            [[nodiscard]] Lock await_resume() const noexcept { return Lock{this->mMutex, std::adopt_lock}; }
        };
    };
 
    extern template class TAsyncMutex<>;
 
    using AsyncMutex = TAsyncMutex<>;
 
    template<typename ExecutorType> class TAsyncSharedMutex;
 
    template<typename ExecutorType = InlineExecutor> class TAsyncSharedMutex {
    public:
        class Lock;
        class LockOperation;
 
    private:
        friend class LockOperation;
 
        enum class State {
            Unlocked,
            LockedShared,
            LockedExclusive
        };
 
        ExecutorType mExecutor;
 
        Mutex mMutex;
        State mState = State::Unlocked;
 
        uint64_t mSharedUsers = 0;
 
        uint64_t mExclusiveWaiters = 0;
 
        LockOperation* mHeadWaiter = nullptr;
        LockOperation* mTailWaiter = nullptr;
 
    public:
        explicit TAsyncSharedMutex() {}
 
        explicit TAsyncSharedMutex(ExecutorType&& executor) : mExecutor(std::forward<ExecutorType>(executor)) {}
        ~TAsyncSharedMutex() = default;
 
        TAsyncSharedMutex(const TAsyncSharedMutex&) = delete;
        TAsyncSharedMutex(TAsyncSharedMutex&&) = delete;
        TAsyncSharedMutex& operator=(const TAsyncSharedMutex&) = delete;
        TAsyncSharedMutex& operator=(TAsyncSharedMutex&&) = delete;
 
        [[nodiscard]] LockOperation lock_shared() { return LockOperation{*this, false}; }
 
        [[nodiscard]] LockOperation lock() { return LockOperation{*this, true}; }
 
        auto try_lock_shared() -> bool {
            // To acquire the shared lock the state must be one of two states:
            //   1) unlocked
            //   2) shared locked with zero exclusive waiters
            //          Zero exclusive waiters prevents exclusive starvation if shared locks are
            //          always continuously happening.
 
            UniqueLock<Mutex> lk{mMutex};
            return try_lock_shared_locked(lk);
        }
 
        [[nodiscard]] bool try_lock() {
            // To acquire the exclusive lock the state must be unlocked.
            UniqueLock<Mutex> lk{mMutex};
            return try_lock_locked(lk);
        }
 
        void unlock_shared() {
            UniqueLock<Mutex> lk{mMutex};
            --mSharedUsers;
 
            // Only wake waiters from shared state if all shared users have completed.
            if(mSharedUsers == 0) {
                if(mHeadWaiter != nullptr) {
                    wakeWaiters(lk);
                } else {
                    mState = State::Unlocked;
                }
            }
        }
 
        void unlock() {
            UniqueLock<Mutex> lk{mMutex};
            if(mHeadWaiter != nullptr) {
                wakeWaiters(lk);
            } else {
                mState = State::Unlocked;
            }
        }
 
    public:
        class Lock {
        public:
            Lock(TAsyncSharedMutex& sm, const bool exclusive) : mMutex(&sm), mExclusive(exclusive) {}
 
            ~Lock() { unlock(); }
 
            Lock(const Lock&) = delete;
            Lock& operator=(const Lock&) = delete;
 
            Lock(Lock&& other) noexcept : mMutex(std::exchange(other.mMutex, nullptr)), mExclusive(other.mExclusive) {}
            Lock& operator=(Lock&& other) noexcept {
                if(std::addressof(other) != this) {
                    mMutex = std::exchange(other.mMutex, nullptr);
                    mExclusive = other.mExclusive;
                }
                return *this;
            }
 
            void unlock() {
                if(mMutex != nullptr) {
                    if(mExclusive) {
                        mMutex->unlock();
                    } else {
                        mMutex->unlock_shared();
                    }
 
                    mMutex = nullptr;
                }
            }
 
            [[nodiscard]] bool isExclusive() const noexcept { return mExclusive; }
 
        private:
            TAsyncSharedMutex* mMutex{nullptr};
            bool mExclusive{false};
        };
 
        class LockOperation {
            friend class TAsyncSharedMutex;
 
        private:
            TAsyncSharedMutex& mMutex;
            bool mExclusive = false;
            CoroutineHandle<> mAwaitingCoroutine;
            LockOperation* mNext = nullptr;
 
        public:
            LockOperation(TAsyncSharedMutex& sm, const bool exclusive) : mMutex(sm), mExclusive(exclusive) {}
 
            [[nodiscard]] bool await_ready() const noexcept {
                if(mExclusive) {
                    return mMutex.try_lock();
                } else {
                    return mMutex.try_lock_shared();
                }
            }
 
            [[nodiscard]] bool await_suspend(const CoroutineHandle<> awaitingCoroutine) noexcept {
                UniqueLock<Mutex> lk{mMutex.mMutex};
                // Its possible the lock has been released between await_ready() and await_suspend(), double
                // check and make sure we are not going to suspend when nobody holds the lock.
                if(mExclusive) {
                    if(mMutex.try_lock_locked(lk)) {
                        return false;
                    }
                } else {
                    if(mMutex.try_lock_shared_locked(lk)) {
                        return false;
                    }
                }
 
                // For sure the lock is currently held in a manner that it cannot be acquired, suspend ourself
                // at the end of the waiter list.
 
                if(mMutex.mTailWaiter == nullptr) {
                    mMutex.mHeadWaiter = this;
                    mMutex.mTailWaiter = this;
                } else {
                    mMutex.mTailWaiter->mNext = this;
                    mMutex.mTailWaiter = this;
                }
 
                // If this is an exclusive lock acquire then mark it as so so that shared locks after this
                // exclusive one will also suspend so this exclusive lock doens't get starved.
                if(mExclusive) {
                    ++mMutex.mExclusiveWaiters;
                }
 
                mAwaitingCoroutine = awaitingCoroutine;
                return true;
            }
            [[nodiscard]] Lock await_resume() noexcept { return Lock{mMutex, mExclusive}; }
        };
 
    private:
        [[nodiscard]] bool try_lock_shared_locked(UniqueLock<Mutex>& lk) {
            if(mState == State::Unlocked) {
                // If the shared mutex is unlocked put it into shared mode and add ourself as using the lock.
                mState = State::LockedShared;
                ++mSharedUsers;
                lk.unlock();
                return true;
            } else if(mState == State::LockedShared && mExclusiveWaiters == 0) {
                // If the shared mutex is in a shared locked state and there are no exclusive waiters
                // the add ourself as using the lock.
                ++mSharedUsers;
                lk.unlock();
                return true;
            }
 
            // If the lock is in shared mode but there are exclusive waiters then we will also wait so
            // the writers are not starved.
            // If the lock is in exclusive mode already then we need to wait.
 
            return false;
        }
 
        [[nodiscard]] bool try_lock_locked(UniqueLock<Mutex>& lk) {
            if(mState == State::Unlocked) {
                mState = State::LockedExclusive;
                lk.unlock();
                return true;
            }
            return false;
        }
 
        void wakeWaiters(UniqueLock<Mutex>& lk) {
            // First determine what the next lock state will be based on the first waiter.
            if(mHeadWaiter->mExclusive) {
                // If its exclusive then only this waiter can be woken up.
                mState = State::LockedExclusive;
                LockOperation* waiterToResume = mHeadWaiter;
                mHeadWaiter = mHeadWaiter->mNext;
                --mExclusiveWaiters;
                if(mHeadWaiter == nullptr) {
                    mTailWaiter = nullptr;
                }
 
                // Since this is an exclusive lock waiting we can resume it directly.
                lk.unlock();
                dispatch(mExecutor, [handle = waiterToResume->mAwaitingCoroutine]() mutable { handle.resume(); });
            } else {
                // If its shared then we will scan forward and awake all shared waiters onto the given
                // thread pool so they can run in parallel.
                mState = State::LockedShared;
                do {
                    LockOperation* waiterToResume = mHeadWaiter;
                    mHeadWaiter = mHeadWaiter->mNext;
                    if(mHeadWaiter == nullptr) {
                        mTailWaiter = nullptr;
                    }
                    ++mSharedUsers;
 
                    dispatch(mExecutor, [handle = waiterToResume->mAwaitingCoroutine]() mutable { handle.resume(); });
                } while(mHeadWaiter != nullptr && !mHeadWaiter->mExclusive);
 
                // Cannot unlock until the entire set of shared waiters has
                // been traversed.  I think this makes more sense than
                // allocating space for all the shared waiters, unlocking,
                // and then resuming in a batch?
                lk.unlock();
            }
        }
    };
 
    extern template class TAsyncSharedMutex<>;
 
    using AsyncSharedMutex = TAsyncSharedMutex<>;
 
} // namespace CeresEngine