spinRWMutex.h

//
// Copyright 2022 Pixar
//
// Licensed under the terms set forth in the LICENSE.txt file available at
// https://openusd.org/license.
//
#ifndef PXR_BASE_TF_SPIN_RW_MUTEX_H
#define PXR_BASE_TF_SPIN_RW_MUTEX_H
 
#include "pxr/pxr.h"
#include "pxr/base/tf/api.h"
 
#include "pxr/base/arch/hints.h"
#include "pxr/base/tf/diagnosticLite.h"
 
#include <atomic>
#include <utility>
 
PXR_NAMESPACE_OPEN_SCOPE
 
class TfSpinRWMutex
{
    static constexpr int OneReader = 2;
    static constexpr int WriterFlag = 1;
    
public:
 
    TfSpinRWMutex() : _lockState(0) {}
 
    struct DeferAcquire {};
 
    static constexpr DeferAcquire deferAcquire {};
 
    struct ScopedLock {
 
        explicit ScopedLock(TfSpinRWMutex &m, bool write=true)
            : _mutex(&m)
            , _acqState(_NotAcquired) {
            Acquire(write);
        }
 
        ScopedLock(TfSpinRWMutex &m, TfSpinRWMutex::DeferAcquire)
            : _mutex(&m), _acqState(_NotAcquired) {}
 
        ScopedLock() : _mutex(nullptr), _acqState(_NotAcquired) {}
 
        ScopedLock(ScopedLock &&other) noexcept
            : _mutex(std::exchange(other._mutex, nullptr))
            , _acqState(std::exchange(other._acqState, _NotAcquired)) {}
 
        ScopedLock &operator=(ScopedLock &&other) noexcept {
            if (this != &other) {
                Release();
                _mutex = std::exchange(other._mutex, nullptr);
                _acqState = std::exchange(other._acqState, _NotAcquired);
            }
            return *this;
        }        
 
        ~ScopedLock() {
            Release();
        }
 
        void Acquire(TfSpinRWMutex &m, bool write=true) {
            Release();
            _mutex = &m;
            Acquire(write);
        }            
 
        void Acquire(bool write=true) {
            if (write) {
                AcquireWrite();
            }
            else {
                AcquireRead();
            }
        }
 
        bool TryAcquire(TfSpinRWMutex &m, bool write=true) {
            Release();
            _mutex = &m;
            return TryAcquire(write);
        }
 
        bool TryAcquire(bool write=true) {
            return write ? TryAcquireWrite() : TryAcquireRead();
        }
 
        void Release() {
            switch (_acqState) {
            default:
            case _NotAcquired:
                break;
            case _ReadAcquired:
                _ReleaseRead();
                break;
            case _WriteAcquired:
                _ReleaseWrite();
                break;
            };
        }
 
        void AcquireRead() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == _NotAcquired);
            _mutex->AcquireRead();
            _acqState = _ReadAcquired;
        }
 
        bool TryAcquireRead() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == _NotAcquired);
            if (_mutex->TryAcquireRead()) {
                _acqState = _ReadAcquired;
                return true;
            }
            return false;
        }
 
        void AcquireWrite() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == _NotAcquired);
            _mutex->AcquireWrite();
            _acqState = _WriteAcquired;
        }
 
        bool TryAcquireWrite() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == _NotAcquired);
            if (_mutex->TryAcquireWrite()) {
                _acqState = _WriteAcquired;
                return true;
            }
            return false;
        }
 
        bool TryAcquireWriteIfReleased() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == _NotAcquired);
            if (_mutex->TryAcquireWriteIfReleased()) {
                _acqState = _WriteAcquired;
                return true;
            }
            return false;
        }
 
        bool UpgradeToWriter() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == _ReadAcquired);
            bool result = _mutex->UpgradeToWriter();
            _acqState = _WriteAcquired;
            return result;
        }
 
        bool DowngradeToReader() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == _WriteAcquired);
            _acqState = _ReadAcquired;
            return _mutex->DowngradeToReader();
        }
        
    private:
 
        // Acquisition states.
        static constexpr int _NotAcquired = 0;
        static constexpr int _ReadAcquired = 1;
        static constexpr int _WriteAcquired = 2;
 
        void _ReleaseRead() {
            TF_DEV_AXIOM(_acqState == _ReadAcquired);
            _mutex->ReleaseRead();
            _acqState = _NotAcquired;
        }
 
        void _ReleaseWrite() {
            TF_DEV_AXIOM(_acqState == _WriteAcquired);
            _mutex->ReleaseWrite();
            _acqState = _NotAcquired;
        }
 
        TfSpinRWMutex *_mutex;
        int _acqState; // _NotAcquired   (0),
                       // _ReadAcquired  (1),
                       // _WriteAcquired (2)
    };
 
    inline bool TryAcquireRead() {
        // Optimistically increment the reader count.
        if (ARCH_LIKELY(!(_lockState.fetch_add(
                              OneReader, std::memory_order_acquire) &
                          WriterFlag))) {
            // We incremented the reader count and observed no writer activity,
            // we have a read lock.
            return true;
        }
        // Otherwise there's writer activity.  Undo the increment and return
        // false.  Release ordering ensures a waiting writer can observe this
        // reader count drop to zero.
        _lockState.fetch_sub(OneReader, std::memory_order_release);
        return false;
    }
    
    inline void AcquireRead() {
        while (true) {
            if (TryAcquireRead()) {
                return;
            }
            // There's writer activity.  Wait to see no writer activity and
            // retry.
            _WaitForWriter();
        }
    }
 
    inline void ReleaseRead() {
        // Release ordering ensures the writer waiting in _WaitForReaders can
        // observe this reader count drop to zero.
        _lockState.fetch_sub(OneReader, std::memory_order_release);
    }
 
    inline bool TryAcquireWrite() {
        int state = _lockState.fetch_or(WriterFlag, std::memory_order_acquire);
        if (!(state & WriterFlag)) {
            // We set the flag, wait for readers.
            if (state != 0) {
                // Wait for pending readers.
                _WaitForReaders();
            }
            return true;
        }
        return false;
    }
 
    inline bool TryAcquireWriteIfReleased() {
        int expected = 0;
        return _lockState.compare_exchange_strong(
            expected, WriterFlag,
            std::memory_order_acquire,
            std::memory_order_relaxed);
    }
 
    void AcquireWrite() {
        // Attempt to acquire -- if we fail then wait to see no other writer and
        // retry.
        while (true) {
            if (TryAcquireWrite()) {
                return;
            }
            _WaitForWriter();
        }
    }
 
    inline void ReleaseWrite() {
        _lockState.fetch_and(~WriterFlag, std::memory_order_release);
    }
 
    bool UpgradeToWriter() {
        // This thread owns a read lock, attempt to upgrade to write lock.  If
        // we do so without an intervening writer, return true, otherwise return
        // false.
        const auto acquire = std::memory_order_acquire;
        const auto release = std::memory_order_release;
        bool atomic = true;
        while (true) {
            int state = _lockState.fetch_or(WriterFlag, acquire);
            if (!(state & WriterFlag)) {
                // We set the flag, release our reader count and wait for any
                // other pending readers.  Release ordering on the fetch_sub
                // pairs with the acquire in _WaitForReaders.
                if (_lockState.fetch_sub(
                        OneReader, release) != (OneReader | WriterFlag)) {
                    _WaitForReaders();
                }
                return atomic;
            }
            // There was other writer activity -- wait for it to clear, then
            // retry.
            atomic = false;
            _WaitForWriter();
        }
    }
 
    bool DowngradeToReader() {
        // Simultaneously add a reader count and clear the writer bit.  Since we
        // own the WriterFlag (=1) and we want to bump the reader count by
        // OneReader (=2) we can achieve this by adding OneReader and then
        // clearing WriterFlag.  But that's the same as adding 2 and subtracting
        // 1, so we can do this in a single step by just adding 1.  We write
        // this as 'OneReader - 1' so it's clear this isn't a simple increment.
        //
        // acq_rel: release for the write side (make protected writes visible to
        // subsequent readers), acquire for the read side (establish
        // happens-before for subsequent reads of protected data).
        _lockState.fetch_add(OneReader - 1, std::memory_order_acq_rel);
        return true;
    }
 
private:
    friend class TfBigRWMutex;
    
    // Helpers for staged-acquire-write that BigRWMutex uses.
    enum _StagedAcquireWriteState {
        _StageNotAcquired,
        _StageAcquiring,
        _StageAcquired
    };
 
    // This API lets TfBigRWMutex acquire a write lock step-by-step so that it
    // can begin acquiring write locks on several mutexes without waiting
    // serially for pending readers to complete.  Call _StagedAcquireWriteStep
    // with _StageNotAcquired initially, and save the returned value.  Continue
    // repeatedly calling _StagedAcquireWriteStep, passing the previously
    // returned value until this function returns _StageAcquired.  At this
    // point the write lock is acquired.
    _StagedAcquireWriteState
    _StagedAcquireWriteStep(_StagedAcquireWriteState curState) {
        int state;
        switch (curState) {
        case _StageNotAcquired:
            state = _lockState.fetch_or(WriterFlag, std::memory_order_acquire);
            if (!(state & WriterFlag)) {
                // We set the flag. If there were no readers we're done,
                // otherwise we'll have to wait for them, next step.
                return state == 0 ? _StageAcquired : _StageAcquiring;
            }
            // Other writer activity, must retry next step.
            return _StageNotAcquired;
        case _StageAcquiring:
            // We have set the writer flag but must wait to see no readers.
            _WaitForReaders();
            return _StageAcquired;
        case _StageAcquired:
        default:
            return _StageAcquired;
        };
    }
    
    TF_API void _WaitForReaders() const;
    TF_API void _WaitForWriter() const;
    
    std::atomic<int> _lockState;
};
 
PXR_NAMESPACE_CLOSE_SCOPE
 
#endif // PXR_BASE_TF_SPIN_RW_MUTEX_H