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>
 
PXR_NAMESPACE_OPEN_SCOPE
 
class TfSpinRWMutex
{
    static constexpr int OneReader = 2;
    static constexpr int WriterFlag = 1;
    
public:
 
    TfSpinRWMutex() : _lockState(0) {}
 
    struct ScopedLock {
 
        // Acquisition states.
        static constexpr int NotAcquired = 0;
        static constexpr int ReadAcquired = 1;
        static constexpr int WriteAcquired = 2;
 
        explicit ScopedLock(TfSpinRWMutex &m, bool write=true)
            : _mutex(&m)
            , _acqState(NotAcquired) {
            Acquire(write);
        }
 
        ScopedLock() : _mutex(nullptr), _acqState(NotAcquired) {}
 
        ~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 UpgradeToWriter() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == ReadAcquired);
            _acqState = WriteAcquired;
            return _mutex->UpgradeToWriter();
        }
 
        bool DowngradeToReader() {
            TF_DEV_AXIOM(_mutex);
            TF_DEV_AXIOM(_acqState == WriteAcquired);
            _acqState = ReadAcquired;
            return _mutex->DowngradeToReader();
        }
        
    private:
 
        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) & 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.
        _lockState -= OneReader;
        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() {
        // Just decrement the count.
        _lockState -= OneReader;
    }
 
    inline bool TryAcquireWrite() {
        int state = _lockState.fetch_or(WriterFlag);
        if (!(state & WriterFlag)) {
            // We set the flag, wait for readers.
            if (state != 0) {
                // Wait for pending readers.
                _WaitForReaders();
            }
            return true;
        }
        return false;
    }
 
    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 &= ~WriterFlag;
    }
 
    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.
        bool atomic = true;
        while (true) {
            int state = _lockState.fetch_or(WriterFlag);
            if (!(state & WriterFlag)) {
                // We set the flag, release our reader count and wait for any
                // other pending readers.
                if (_lockState.fetch_sub(
                        OneReader) != (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 by adding
        // (OneReader-1).
        _lockState += (OneReader-1);
        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);
            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