[K/N] Add kotlin::RWSpinLock ^KT-56233

Trivially constructible and destructible std::shared_mutex
implementation with additional guarantee that try_* methods fail only if
the mutex is locked.

kotlin-native/runtime/src/main/cpp/Mutex.hpp

@@ -18,6 +18,7 @@
 #define RUNTIME_MUTEX_H
 
 #include <atomic>
+#include <cinttypes>
 #include <thread>
 
 #include "KAssert.h"
@@ -78,6 +79,87 @@ private:
     std::atomic_flag flag_ = ATOMIC_FLAG_INIT;
 };
 
+template <MutexThreadStateHandling threadStateHandling>
+class RWSpinLock;
+
+template <>
+class RWSpinLock<MutexThreadStateHandling::kIgnore> : private Pinned {
+    using State = uint64_t;
+
+public:
+    void lock() noexcept {
+        auto state = state_.load(std::memory_order_relaxed);
+        while (!try_lock_impl(state)) {
+            if (state % 2 == 0) {
+                // There're some readers and no writers are pending. Try announcing the writer.
+                state_.compare_exchange_strong(state, state + 1, std::memory_order_relaxed, std::memory_order_relaxed);
+                continue;
+            }
+            // Lock is busy, wait.
+            yield();
+            state = state_.load(std::memory_order_relaxed);
+        }
+    }
+
+    bool try_lock() noexcept {
+        auto state = state_.load(std::memory_order_relaxed);
+        return try_lock_impl(state);
+    }
+
+    void unlock() noexcept {
+        auto actual = state_.exchange(0, std::memory_order_release);
+        RuntimeAssert(
+                actual == kLocked, "Broken RWSpinLock@%p::unlock. Expected state to be %" PRIu64 " actual %" PRIu64, this, kLocked, actual);
+    }
+
+    void lock_shared() noexcept {
+        auto state = state_.load(std::memory_order_relaxed);
+        while (!try_lock_shared_impl(state)) {
+            // Lock is busy, wait.
+            yield();
+            state = state_.load(std::memory_order_relaxed);
+        }
+    }
+
+    // Fails only if there're pending writers.
+    bool try_lock_shared() noexcept {
+        auto state = state_.load(std::memory_order_relaxed);
+        return try_lock_shared_impl(state);
+    }
+
+    void unlock_shared() noexcept { state_.fetch_sub(2, std::memory_order_release); }
+
+private:
+    bool try_lock_impl(State& state) noexcept {
+        do {
+            if (state > 1) {
+                // Readers are pending, or it's blocked by another writer.
+                return false;
+            }
+            // Only writers are pending. Let's try locking.
+        } while (!state_.compare_exchange_weak(state, kLocked, std::memory_order_acquire, std::memory_order_relaxed));
+        return true;
+    }
+
+    bool try_lock_shared_impl(State& state) noexcept {
+        do {
+            if (state % 2 != 0) {
+                // There's a pending or active writer. Fail.
+                return false;
+            }
+            // Only readers
+        } while (!state_.compare_exchange_weak(state, state + 2, std::memory_order_acquire, std::memory_order_relaxed));
+        return true;
+    }
+
+    // No need to check for external calls, because we explicitly ignore thread state.
+    static NO_EXTERNAL_CALLS_CHECK NO_INLINE void yield() { std::this_thread::yield(); }
+
+    static inline constexpr State kLocked = std::numeric_limits<State>::max();
+    static_assert(kLocked % 2 == 1, "Must be odd");
+    std::atomic<State> state_{0};
+};
+
 } // namespace kotlin
 
 #endif // RUNTIME_MUTEX_H

kotlin-native/runtime/src/main/cpp/MutexTest.cpp

@@ -6,7 +6,9 @@
 #include "Mutex.hpp"
 
 #include <mutex>
+#include <shared_mutex>
 
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 #include "ScopedThread.hpp"
@@ -47,4 +49,179 @@ TYPED_TEST(MutexTest, SmokeDetachedThread) {
     EXPECT_EQ(protectedCounter, 2);
 }
 
+TEST(RWSpinLockTest, Lock) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    mutex.lock();
+    mutex.unlock();
+    // Check that unlock unlocked.
+    mutex.lock();
+    mutex.unlock();
+}
 
+TEST(RWSpinLockTest, TryLock) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    EXPECT_TRUE(mutex.try_lock());
+    EXPECT_FALSE(mutex.try_lock());
+    mutex.unlock();
+    mutex.lock();
+    EXPECT_FALSE(mutex.try_lock());
+    mutex.unlock();
+}
+
+TEST(RWSpinLockTest, LockShared) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    mutex.lock_shared();
+    mutex.lock_shared();
+    mutex.unlock_shared();
+    mutex.lock_shared();
+    mutex.unlock_shared();
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, TryLockShared) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    EXPECT_TRUE(mutex.try_lock_shared());
+    EXPECT_TRUE(mutex.try_lock_shared());
+    mutex.unlock_shared();
+    EXPECT_TRUE(mutex.try_lock_shared());
+    mutex.unlock_shared();
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, LockSharedWhileLocked) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    bool done = false;
+    mutex.lock();
+    ScopedThread thread([&] {
+        mutex.lock_shared();
+        EXPECT_TRUE(done);
+        mutex.unlock_shared();
+    });
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    done = true;
+    mutex.unlock();
+}
+
+TEST(RWSpinLockTest, TryLockSharedWhileLocked) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    mutex.lock();
+    EXPECT_FALSE(mutex.try_lock_shared());
+    mutex.unlock();
+}
+
+TEST(RWSpinLockTest, LockWhileLockShared) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    bool done = false;
+    mutex.lock_shared();
+    ScopedThread thread([&] {
+        mutex.lock();
+        EXPECT_TRUE(done);
+        mutex.unlock();
+    });
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    // Shared lock is not where one does mutability. But it's okay here.
+    done = true;
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, TryLockWhileLockShared) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    mutex.lock_shared();
+    EXPECT_FALSE(mutex.try_lock());
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, LockSharedWhileLockSharedWithPendingLock) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    int state = 0;
+    mutex.lock_shared();
+    ScopedThread thread_lock([&] {
+        mutex.lock();
+        EXPECT_THAT(state, 1);
+        state = 2;
+        mutex.unlock();
+    });
+    // Wait for thread_lock to say that it wants to lock the mutex.
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    ScopedThread thread_lock_shared([&] {
+        mutex.lock_shared();
+        EXPECT_THAT(state, 2);
+        mutex.unlock_shared();
+    });
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    // Shared lock is not where one does mutability. But it's okay here.
+    state = 1;
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, TryLockSharedWhileLockSharedWithPendingLock) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    bool done = false;
+    mutex.lock_shared();
+    ScopedThread thread_lock([&] {
+        mutex.lock();
+        EXPECT_TRUE(done);
+        mutex.unlock();
+    });
+    // Wait for thread_lock to say that it wants to lock the mutex.
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    EXPECT_FALSE(mutex.try_lock_shared());
+    // Shared lock is not where one does mutability. But it's okay here.
+    done = true;
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, LockSharedWhileLockSharedWithFailedPendingLock) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    mutex.lock_shared();
+    EXPECT_FALSE(mutex.try_lock());
+    mutex.lock_shared();
+    mutex.unlock_shared();
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, TryLockSharedWhileLockSharedWithFailedPendingLock) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    mutex.lock_shared();
+    EXPECT_FALSE(mutex.try_lock());
+    EXPECT_TRUE(mutex.try_lock_shared());
+    mutex.unlock_shared();
+    mutex.unlock_shared();
+}
+
+TEST(RWSpinLockTest, Counter) {
+    RWSpinLock<MutexThreadStateHandling::kIgnore> mutex;
+    // tsan will help catch if the lock breaks memory ordering.
+    uint64_t counter = 0;
+    constexpr uint64_t boundary = 10000;
+    std::atomic<bool> canStart = false;
+    std::vector<ScopedThread> writers;
+    for (int i = 0; i < 2; ++i) {
+        writers.emplace_back([&] {
+            while (!canStart.load(std::memory_order_acquire)) {
+            }
+            while (true) {
+                std::unique_lock guard(mutex);
+                if (counter == boundary) return;
+                counter += 1;
+            }
+        });
+    }
+    std::vector<ScopedThread> readers;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        readers.emplace_back([&] {
+            while (!canStart.load(std::memory_order_acquire)) {
+            }
+            while (true) {
+                std::shared_lock guard(mutex);
+                EXPECT_THAT(counter, testing::Ge(uint64_t(0)));
+                EXPECT_THAT(counter, testing::Le(boundary));
+                if (counter == boundary) return;
+            }
+        });
+    }
+    canStart.store(true, std::memory_order_release);
+    readers.clear();
+    writers.clear();
+    EXPECT_THAT(counter, boundary);
+}