RWSpinLock.hpp

Go to the documentation of this file.

#ifndef CPP_UTILITIES_MEMORYSAFETY_RWSPINLOCK_HPP
#define CPP_UTILITIES_MEMORYSAFETY_RWSPINLOCK_HPP
/*
 * Copyright 2011-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <algorithm>
#include <atomic>
#include <thread>
#include "../Common.h"
 
UTILITIES_NAMESPACE_BEGIN
 
namespace Memory {
class RWSpinLock {
    enum : int32_t { READER = 4, UPGRADED = 2, WRITER = 1 };
 
public:
    constexpr RWSpinLock() : bits_(0) {}
 
    RWSpinLock(RWSpinLock const&) = delete;
    RWSpinLock& operator=(RWSpinLock const&) = delete;
 
    void lock() {
        uint_fast32_t count = 0;
        while (!try_lock()) {
            if (++count > 1000) {
                std::this_thread::yield();
            }
        }
    }
 
    void unlock() {
        static_assert(READER > WRITER + UPGRADED, "wrong bits!");
        bits_.fetch_and(~(WRITER | UPGRADED), std::memory_order_release);
    }
 
    void lock_shared() {
        uint_fast32_t count = 0;
        while (!try_lock_shared()) {
            if (++count > 1000) {
                std::this_thread::yield();
            }
        }
    }
 
    void unlock_shared() {
        bits_.fetch_add(-READER, std::memory_order_release);
    }
 
    void unlock_and_lock_shared() {
        bits_.fetch_add(READER, std::memory_order_acquire);
        unlock();
    }
 
    void lock_upgrade() {
        uint_fast32_t count = 0;
        while (!try_lock_upgrade()) {
            if (++count > 1000) {
                std::this_thread::yield();
            }
        }
    }
 
    void unlock_upgrade() {
        bits_.fetch_add(-UPGRADED, std::memory_order_acq_rel);
    }
 
    void unlock_upgrade_and_lock() {
        int64_t count = 0;
        while (!try_unlock_upgrade_and_lock()) {
            if (++count > 1000) {
                std::this_thread::yield();
            }
        }
    }
 
    void unlock_upgrade_and_lock_shared() {
        bits_.fetch_add(READER - UPGRADED, std::memory_order_acq_rel);
    }
 
    void unlock_and_lock_upgrade() {
        // need to do it in two steps here -- as the UPGRADED bit might be OR-ed
        // at the same time when other threads are trying do try_lock_upgrade().
        bits_.fetch_or(UPGRADED, std::memory_order_acquire);
        bits_.fetch_add(-WRITER, std::memory_order_release);
    }
 
    bool try_lock() {
        int32_t expect = 0;
        return bits_.compare_exchange_strong(
                    expect, WRITER, std::memory_order_acq_rel);
    }
 
    bool try_lock_shared() {
        // fetch_add is considerably (100%) faster than compare_exchange,
        // so here we are optimizing for the common (lock success) case.
        int32_t value = bits_.fetch_add(READER, std::memory_order_acquire);
        if (value & (WRITER | UPGRADED)) {
            bits_.fetch_add(-READER, std::memory_order_release);
            return false;
        }
        return true;
    }
 
    bool try_unlock_upgrade_and_lock() {
        int32_t expect = UPGRADED;
        return bits_.compare_exchange_strong(
                    expect, WRITER, std::memory_order_acq_rel);
    }
 
    bool try_lock_upgrade() {
        int32_t value = bits_.fetch_or(UPGRADED, std::memory_order_acquire);
 
        return ((value & (UPGRADED | WRITER)) == 0);
    }
 
    int32_t bits() const {
        return bits_.load(std::memory_order_acquire);
    }
 
    class ReadHolder;
    class UpgradedHolder;
    class WriteHolder;
 
    class ReadHolder {
    public:
        explicit ReadHolder(RWSpinLock* lock) : lock_(lock) {
            if (lock_) {
                lock_->lock_shared();
            }
        }
 
        explicit ReadHolder(RWSpinLock& lock) : lock_(&lock) {
            lock_->lock_shared();
        }
 
        ReadHolder(ReadHolder&& other) noexcept : lock_(other.lock_) {
            other.lock_ = nullptr;
        }
 
        explicit ReadHolder(UpgradedHolder&& upgraded) : lock_(upgraded.lock_) {
            upgraded.lock_ = nullptr;
            if (lock_) {
                lock_->unlock_upgrade_and_lock_shared();
            }
        }
 
        explicit ReadHolder(WriteHolder&& writer) : lock_(writer.lock_) {
            writer.lock_ = nullptr;
            if (lock_) {
                lock_->unlock_and_lock_shared();
            }
        }
 
        ReadHolder& operator=(ReadHolder&& other) {
            using std::swap;
            swap(lock_, other.lock_);
            return *this;
        }
 
        ReadHolder(const ReadHolder& other) = delete;
        ReadHolder& operator=(const ReadHolder& other) = delete;
 
        ~ReadHolder() {
            if (lock_) {
                lock_->unlock_shared();
            }
        }
 
        void reset(RWSpinLock* lock = nullptr) {
            if (lock == lock_) {
                return;
            }
            if (lock_) {
                lock_->unlock_shared();
            }
            lock_ = lock;
            if (lock_) {
                lock_->lock_shared();
            }
        }
 
        void swap(ReadHolder& other) {
            std::swap(lock_, other.lock_);
        }
 
    private:
        friend class UpgradedHolder;
        friend class WriteHolder;
        RWSpinLock* lock_;
    };
 
    class UpgradedHolder {
    public:
        explicit UpgradedHolder(RWSpinLock* lock) : lock_(lock) {
            if (lock_) {
                lock_->lock_upgrade();
            }
        }
 
        explicit UpgradedHolder(RWSpinLock& lock) : lock_(&lock) {
            lock_->lock_upgrade();
        }
 
        explicit UpgradedHolder(WriteHolder&& writer) {
            lock_ = writer.lock_;
            writer.lock_ = nullptr;
            if (lock_) {
                lock_->unlock_and_lock_upgrade();
            }
        }
 
        UpgradedHolder(UpgradedHolder&& other) noexcept : lock_(other.lock_) {
            other.lock_ = nullptr;
        }
 
        UpgradedHolder& operator=(UpgradedHolder&& other) {
            using std::swap;
            swap(lock_, other.lock_);
            return *this;
        }
 
        UpgradedHolder(const UpgradedHolder& other) = delete;
        UpgradedHolder& operator=(const UpgradedHolder& other) = delete;
 
        ~UpgradedHolder() {
            if (lock_) {
                lock_->unlock_upgrade();
            }
        }
 
        void reset(RWSpinLock* lock = nullptr) {
            if (lock == lock_) {
                return;
            }
            if (lock_) {
                lock_->unlock_upgrade();
            }
            lock_ = lock;
            if (lock_) {
                lock_->lock_upgrade();
            }
        }
 
        void swap(UpgradedHolder& other) {
            using std::swap;
            swap(lock_, other.lock_);
        }
 
    private:
        friend class WriteHolder;
        friend class ReadHolder;
        RWSpinLock* lock_;
    };
 
    class WriteHolder {
    public:
        explicit WriteHolder(RWSpinLock* lock) : lock_(lock) {
            if (lock_) {
                lock_->lock();
            }
        }
 
        explicit WriteHolder(RWSpinLock& lock) : lock_(&lock) {
            lock_->lock();
        }
 
        explicit WriteHolder(UpgradedHolder&& upgraded) {
            lock_ = upgraded.lock_;
            upgraded.lock_ = nullptr;
            if (lock_) {
                lock_->unlock_upgrade_and_lock();
            }
        }
 
        WriteHolder(WriteHolder&& other) noexcept : lock_(other.lock_) {
            other.lock_ = nullptr;
        }
 
        WriteHolder& operator=(WriteHolder&& other) {
            using std::swap;
            swap(lock_, other.lock_);
            return *this;
        }
 
        WriteHolder(const WriteHolder& other) = delete;
        WriteHolder& operator=(const WriteHolder& other) = delete;
 
        ~WriteHolder() {
            if (lock_) {
                lock_->unlock();
            }
        }
 
        void reset(RWSpinLock* lock = nullptr) {
            if (lock == lock_) {
                return;
            }
            if (lock_) {
                lock_->unlock();
            }
            lock_ = lock;
            if (lock_) {
                lock_->lock();
            }
        }
 
        void swap(WriteHolder* other) {
            using std::swap;
            swap(lock_, other->lock_);
        }
 
    private:
        friend class ReadHolder;
        friend class UpgradedHolder;
        RWSpinLock* lock_;
    };
 
private:
    std::atomic<int32_t> bits_;
};
} // namespace Memory
UTILITIES_NAMESPACE_END
 
#endif  // CPP_UTILITIES_MEMORYSAFETY_RWSPINLOCK_HPP