spinRWMutex.h

//
// Copyright 2022 Pixar
//
// Licensed under the Apache License, Version 2.0 (the "Apache License")
// with the following modification; you may not use this file except in
// compliance with the Apache License and the following modification to it:
// Section 6. Trademarks. is deleted and replaced with:
//
// 6. Trademarks. This License does not grant permission to use the trade
//    names, trademarks, service marks, or product names of the Licensor
//    and its affiliates, except as required to comply with Section 4(c) of
//    the License and to reproduce the content of the NOTICE file.
//
// You may obtain a copy of the Apache License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the Apache License with the above modification is
// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the Apache License for the specific
// language governing permissions and limitations under the Apache 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();
            }
        }

        void Release() {
            switch (_acqState) {
            default:
            case NotAcquired:
                break;
            case ReadAcquired:
                _ReleaseRead();
                break;
            case WriteAcquired:
                _ReleaseWrite();
                break;
            };
        }

        void AcquireRead() {
            TF_DEV_AXIOM(_acqState == NotAcquired);
            _mutex->AcquireRead();
            _acqState = ReadAcquired;
        }

        void AcquireWrite() {
            TF_DEV_AXIOM(_acqState == NotAcquired);
            _mutex->AcquireWrite();
            _acqState = WriteAcquired;
        }

        bool UpgradeToWriter() {
            TF_DEV_AXIOM(_acqState == ReadAcquired);
            _acqState = WriteAcquired;
            return _mutex->UpgradeToWriter();
        }

        bool DowngradeToReader() {
            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