[K/N][Runtime] Support thread state switching in spin locks

Issue #KT-49013 Fixed

kotlin-native/runtime/src/legacymm/cpp/Memory.cpp

@@ -227,14 +227,14 @@ class CycleDetector : private kotlin::Pinned, public KonanAllocatorAware {
   }
 
   void insertCandidate(KRef candidate) {
-    std::lock_guard<kotlin::SpinLock> guard(lock_);
+    std::lock_guard guard(lock_);
 
     auto it = candidateList_.insert(candidateList_.begin(), candidate);
     candidateInList_.emplace(candidate, it);
   }
 
   void removeCandidate(KRef candidate) {
-    std::lock_guard<kotlin::SpinLock> guard(lock_);
+    std::lock_guard guard(lock_);
 
     auto it = candidateInList_.find(candidate);
     if (it == candidateInList_.end())
@@ -243,7 +243,7 @@ class CycleDetector : private kotlin::Pinned, public KonanAllocatorAware {
     candidateInList_.erase(it);
   }
 
-  kotlin::SpinLock lock_;
+  kotlin::SpinLock<kotlin::MutexThreadStateHandling::kIgnore> lock_;
   using CandidateList = KStdList<KRef>;
   CandidateList candidateList_;
   KStdUnorderedMap<KRef, CandidateList::iterator> candidateInList_;
@@ -2994,7 +2994,7 @@ ScopedRefHolder::~ScopedRefHolder() {
 CycleDetectorRootset CycleDetector::collectRootset() {
   auto& detector = instance();
   CycleDetectorRootset rootset;
-  std::lock_guard<kotlin::SpinLock> guard(detector.lock_);
+  std::lock_guard guard(detector.lock_);
   for (auto* candidate: detector.candidateList_) {
     // Only frozen candidates are to be analyzed.
     if (!isPermanentOrFrozen(candidate))

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

@@ -47,18 +47,6 @@ void ThrowException(KRef exception) {
 
 namespace {
 
-// This function accesses a TLS variable under the hood, so it must not be called from a thread destructor
-// (see kotlin::mm::IsCurrentThreadRegistered, kotlin::mm::GetMemoryState()).
-[[nodiscard]] kotlin::ThreadStateGuard setNativeStateForRegisteredThread(bool reentrant = true) {
-    if (kotlin::mm::IsCurrentThreadRegistered()) {
-        return kotlin::ThreadStateGuard(kotlin::ThreadState::kNative, reentrant);
-    } else {
-        // The current thread is not registered in the Kotlin runtime,
-        // just return an empty guard which doesn't actually switch the state.
-        return kotlin::ThreadStateGuard();
-    }
-}
-
 class {
     /**
      * Timeout 5 sec for concurrent (second) terminate attempt to give a chance the first one to finish.
@@ -73,7 +61,7 @@ class {
         // block() is supposed to be NORETURN, otherwise go to normal abort()
         konan::abort();
       } else {
-        auto guard = setNativeStateForRegisteredThread();
+        kotlin::NativeOrUnregisteredThreadGuard guard(/* reentrant = */ true);
         sleep(timeoutSec);
         // We come here when another terminate handler hangs for 5 sec, that looks fatally broken. Go to forced exit now.
       }
@@ -147,12 +135,12 @@ class TerminateHandler : private kotlin::Pinned {
           terminateWithUnhandledException(e.GetExceptionObject());
         } catch (...) {
           // Not a Kotlin exception - call default handler
-          auto guard = setNativeStateForRegisteredThread();
+          kotlin::NativeOrUnregisteredThreadGuard guard(/* reentrant = */ true);
           queuedHandler();
         }
       }
       // Come here in case of direct terminate() call or unknown exception - go to default terminate handler.
-      auto guard = setNativeStateForRegisteredThread();
+      kotlin::NativeOrUnregisteredThreadGuard guard(/* reentrant = */ true);
       queuedHandler();
   }
 

kotlin-native/runtime/src/main/cpp/Memory.h

@@ -497,6 +497,20 @@ ALWAYS_INLINE inline R CallWithThreadState(R(*function)(Args...), Args... args)
     return function(std::forward<Args>(args)...);
 }
 
+class NativeOrUnregisteredThreadGuard final : private MoveOnly {
+public:
+    explicit NativeOrUnregisteredThreadGuard(bool reentrant = false) noexcept {
+        // The default ctor of ThreadStateGuard doesn't set the state.
+        // So the actual state switching is performed only if the thread is registered.
+        if (kotlin::mm::IsCurrentThreadRegistered()) {
+            backingGuard_ = kotlin::ThreadStateGuard(kotlin::ThreadState::kNative, reentrant);
+        }
+    }
+
+private:
+    ThreadStateGuard backingGuard_;
+};
+
 extern const bool kSupportsMultipleMutators;
 
 } // namespace kotlin

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

@@ -16,7 +16,7 @@
 namespace kotlin {
 
 // A queue that is constructed by collecting subqueues from several `Producer`s.
-template <typename T>
+template <typename T, typename Mutex>
 class MultiSourceQueue {
 public:
     class Producer;
@@ -66,7 +66,7 @@ public:
             for (auto& node : queue_) {
                 node.owner_ = nullptr;
             }
-            std::lock_guard<SpinLock> guard(owner_.mutex_);
+            std::lock_guard<Mutex> guard(owner_.mutex_);
             owner_.queue_.splice(owner_.queue_.end(), queue_);
             owner_.deletionQueue_.splice(owner_.deletionQueue_.end(), deletionQueue_);
         }
@@ -114,7 +114,7 @@ public:
         explicit Iterable(MultiSourceQueue& owner) noexcept : owner_(owner), guard_(owner_.mutex_) {}
 
         MultiSourceQueue& owner_; // weak
-        std::unique_lock<SpinLock> guard_;
+        std::unique_lock<Mutex> guard_;
     };
 
     // Lock `MultiSourceQueue` for safe iteration. If element was scheduled for deletion,
@@ -123,7 +123,7 @@ public:
 
     // Lock `MultiSourceQueue` and apply deletions. Only deletes elements that were published.
     void ApplyDeletions() noexcept {
-        std::lock_guard<SpinLock> guard(mutex_);
+        std::lock_guard<Mutex> guard(mutex_);
         KStdList<Node*> remainingDeletions;
 
         auto it = deletionQueue_.begin();
@@ -154,7 +154,7 @@ private:
     // which is important for GC mark phase.
     KStdList<Node> queue_;
     KStdList<Node*> deletionQueue_;
-    SpinLock mutex_;
+    Mutex mutex_;
 };
 
 } // namespace kotlin

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

@@ -18,8 +18,8 @@ using namespace kotlin;
 
 namespace {
 
-template <typename T>
+template <typename T, typename Mutex>
-KStdVector<T> Collect(MultiSourceQueue<T>& queue) {
+KStdVector<T> Collect(MultiSourceQueue<T, Mutex>& queue) {
     KStdVector<T> result;
     for (const auto& element : queue.LockForIter()) {
         result.push_back(element);
@@ -29,7 +29,7 @@ KStdVector<T> Collect(MultiSourceQueue<T>& queue) {
 
 } // namespace
 
-using IntQueue = MultiSourceQueue<int>;
+using IntQueue = MultiSourceQueue<int, SpinLock<MutexThreadStateHandling::kIgnore>>;
 
 TEST(MultiSourceQueueTest, Insert) {
     IntQueue queue;

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

@@ -20,11 +20,20 @@
 #include <cstdint>
 
 #include "KAssert.h"
+#include "Memory.h"
 #include "Utils.hpp"
 
 namespace kotlin {
 
-class SpinLock : private Pinned {
+enum class MutexThreadStateHandling {
+    kIgnore, kSwitchIfRegistered
+};
+
+template <MutexThreadStateHandling threadStateHandling>
+class SpinLock;
+
+template <>
+class SpinLock<MutexThreadStateHandling::kIgnore> : private Pinned {
 public:
     void lock() noexcept {
         while (!__sync_bool_compare_and_swap(&atomicInt, 0, 1)) {
@@ -43,6 +52,32 @@ private:
     int32_t atomicInt = 0;
 };
 
+template <>
+class SpinLock<MutexThreadStateHandling::kSwitchIfRegistered> : private Pinned {
+public:
+    void lock() noexcept {
+        // Fast path without thread state switching.
+        if (__sync_bool_compare_and_swap(&atomicInt, 0, 1)) {
+            return;
+        }
+
+        kotlin::NativeOrUnregisteredThreadGuard guard(/* reentrant = */ true);
+        while (!__sync_bool_compare_and_swap(&atomicInt, 0, 1)) {
+            // TODO: yield.
+        }
+    }
+
+    void unlock() noexcept {
+        if (!__sync_bool_compare_and_swap(&atomicInt, 1, 0)) {
+            RuntimeAssert(false, "Unable to unlock");
+        }
+    }
+
+private:
+    // TODO: Consider using `std::atomic_flag`.
+    int32_t atomicInt = 0;
+};
+
 } // namespace kotlin
 
 #endif // RUNTIME_MUTEX_H

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "Mutex.hpp"
+
+#include <mutex>
+#include <thread>
+
+#include "gtest/gtest.h"
+#include "TestSupport.hpp"
+
+using namespace kotlin;
+
+template <typename T>
+class MutexTest : public testing::Test {};
+
+using LockTypes = testing::Types<SpinLock<MutexThreadStateHandling::kIgnore>, SpinLock<MutexThreadStateHandling::kSwitchIfRegistered>>;
+TYPED_TEST_SUITE(MutexTest, LockTypes);
+
+TYPED_TEST(MutexTest, SmokeDetachedThread) {
+    using LockUnderTest = TypeParam;
+
+    LockUnderTest lock;
+    std::thread secondThread;
+    std::atomic<bool> started = false;
+    std::atomic<int32_t> protectedCounter = 0;
+
+    {
+        std::unique_lock guard1(lock);
+        secondThread = std::thread([&lock, &started, &protectedCounter]() {
+            started = true;
+            std::unique_lock guard2(lock);
+            protectedCounter++;
+        });
+
+        protectedCounter++;
+        while (!started) { std::this_thread::yield(); }
+
+        // Wait to give the second thread a chance to try to acquire the lock.
+        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        EXPECT_EQ(protectedCounter, 1);
+    }
+    secondThread.join();
+    EXPECT_EQ(protectedCounter, 2);
+}
+
+

kotlin-native/runtime/src/main/cpp/ObjCExport.mm

@@ -982,7 +982,7 @@ static const TypeInfo* createTypeInfo(Class clazz, const TypeInfo* superType, co
   return result;
 }
 
-static kotlin::SpinLock typeInfoCreationMutex;
+static kotlin::SpinLock<kotlin::MutexThreadStateHandling::kSwitchIfRegistered> typeInfoCreationMutex;
 
 static const TypeInfo* getOrCreateTypeInfo(Class clazz) {
   const TypeInfo* result = Kotlin_ObjCExport_getAssociatedTypeInfo(clazz);
@@ -996,7 +996,7 @@ static const TypeInfo* getOrCreateTypeInfo(Class clazz) {
     theForeignObjCObjectTypeInfo :
     getOrCreateTypeInfo(superClass);
 
-  std::lock_guard<kotlin::SpinLock> lockGuard(typeInfoCreationMutex);
+  std::lock_guard lockGuard(typeInfoCreationMutex);
 
   result = Kotlin_ObjCExport_getAssociatedTypeInfo(clazz); // double-checking.
   if (result == nullptr) {
@@ -1016,7 +1016,7 @@ const TypeInfo* Kotlin_ObjCExport_createTypeInfoWithKotlinFieldsFrom(Class clazz
   return createTypeInfo(clazz, superType, fieldsInfo);
 }
 
-static kotlin::SpinLock classCreationMutex;
+static kotlin::SpinLock<kotlin::MutexThreadStateHandling::kSwitchIfRegistered> classCreationMutex;
 static int anonymousClassNextId = 0;
 
 static void addVirtualAdapters(Class clazz, const ObjCTypeAdapter* typeAdapter) {
@@ -1092,7 +1092,7 @@ static Class getOrCreateClass(const TypeInfo* typeInfo) {
   } else {
     Class superClass = getOrCreateClass(typeInfo->superType_);
 
-    std::lock_guard<kotlin::SpinLock> lockGuard(classCreationMutex); // Note: non-recursive
+    std::lock_guard lockGuard(classCreationMutex); // Note: non-recursive
 
     result = typeInfo->writableInfo_->objCExport.objCClass; // double-checking.
     if (result == nullptr) {

kotlin-native/runtime/src/main/cpp/ObjCInterop.mm

@@ -253,7 +253,7 @@ static void AddMethods(Class clazz, const struct ObjCMethodDescription* methods,
   }
 }
 
-static kotlin::SpinLock classCreationMutex;
+static kotlin::SpinLock<kotlin::MutexThreadStateHandling::kSwitchIfRegistered> classCreationMutex;
 static int anonymousClassNextId = 0;
 
 NO_EXTERNAL_CALLS_CHECK static Class allocateClass(const KotlinObjCClassInfo* info) {
@@ -284,7 +284,7 @@ NO_EXTERNAL_CALLS_CHECK static Class allocateClass(const KotlinObjCClassInfo* in
 }
 
 void* CreateKotlinObjCClass(const KotlinObjCClassInfo* info) {
-  std::lock_guard<kotlin::SpinLock> lockGuard(classCreationMutex);
+  std::lock_guard lockGuard(classCreationMutex);
 
   void* createdClass = *info->createdClass;
   if (createdClass != nullptr) {

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

@@ -20,7 +20,7 @@
 namespace kotlin {
 
 // TODO: Consider different locking mechanisms.
-template <typename Value, typename Mutex = SpinLock>
+template <typename Value, typename Mutex>
 class SingleLockList : private Pinned {
 public:
     class Node;

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

@@ -19,7 +19,7 @@ using namespace kotlin;
 
 namespace {
 
-using IntList = SingleLockList<int>;
+using IntList = SingleLockList<int, SpinLock<MutexThreadStateHandling::kIgnore>>;
 
 } // namespace
 
@@ -376,7 +376,7 @@ private:
 } // namespace
 
 TEST(SingleLockListTest, PinnedType) {
-    SingleLockList<PinnedType> list;
+    SingleLockList<PinnedType, SpinLock<MutexThreadStateHandling::kIgnore>> list;
     constexpr int kFirst = 1;
 
     auto* itemNode = list.Emplace(kFirst);
@@ -414,7 +414,7 @@ private:
 TEST(SingleLockListTest, Destructor) {
     testing::StrictMock<testing::MockFunction<DestructorHook>> hook;
     {
-        SingleLockList<WithDestructorHook> list;
+        SingleLockList<WithDestructorHook, SpinLock<MutexThreadStateHandling::kIgnore>> list;
         auto* first = list.Emplace(hook.AsStdFunction())->Get();
         auto* second = list.Emplace(hook.AsStdFunction())->Get();
         auto* third = list.Emplace(hook.AsStdFunction())->Get();

kotlin-native/runtime/src/mm/cpp/GlobalsRegistry.hpp

@@ -15,16 +15,18 @@ namespace kotlin {
 namespace mm {
 
 class GlobalsRegistry : Pinned {
+    using Mutex = SpinLock<MutexThreadStateHandling::kIgnore>;
+
 public:
-    class ThreadQueue : public MultiSourceQueue<ObjHeader**>::Producer {
+    class ThreadQueue : public MultiSourceQueue<ObjHeader**, Mutex>::Producer {
     public:
         explicit ThreadQueue(GlobalsRegistry& registry) : Producer(registry.globals_) {}
         // Do not add fields as this is just a wrapper and Producer does not have virtual destructor.
     };
 
-    using Iterable = MultiSourceQueue<ObjHeader**>::Iterable;
+    using Iterable = MultiSourceQueue<ObjHeader**, Mutex>::Iterable;
 
-    using Iterator = MultiSourceQueue<ObjHeader**>::Iterator;
+    using Iterator = MultiSourceQueue<ObjHeader**, Mutex>::Iterator;
 
     GlobalsRegistry();
     ~GlobalsRegistry();
@@ -47,7 +49,7 @@ public:
 
 private:
     // TODO: Add-only MultiSourceQueue can be made more efficient. Measure, if it's a problem.
-    MultiSourceQueue<ObjHeader**> globals_;
+    MultiSourceQueue<ObjHeader**, Mutex> globals_;
 };
 
 } // namespace mm

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

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "Mutex.hpp"
+
+#include <thread>
+
+#include "gtest/gtest.h"
+#include "TestSupport.hpp"
+#include "ThreadState.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+template <MutexThreadStateHandling Handling, ThreadState State>
+class Mode {
+public:
+    using Lock = SpinLock<Handling>;
+    static constexpr MutexThreadStateHandling threadStateHandling = Handling;
+    static constexpr ThreadState initialState = State;
+};
+
+} // namespace
+
+template <typename T>
+class MutexTestNewMM : public testing::Test {};
+
+using TestModes = testing::Types<
+        Mode<MutexThreadStateHandling::kIgnore, ThreadState::kRunnable>,
+        Mode<MutexThreadStateHandling::kIgnore, ThreadState::kNative>,
+        Mode<MutexThreadStateHandling::kSwitchIfRegistered, ThreadState::kRunnable>,
+        Mode<MutexThreadStateHandling::kSwitchIfRegistered, ThreadState::kNative>>;
+
+TYPED_TEST_SUITE(MutexTestNewMM, TestModes);
+
+TYPED_TEST(MutexTestNewMM, SmokeAttachedThread) {
+    RunInNewThread([](){
+        using LockUnderTest = typename TypeParam::Lock;
+        ThreadState initialThreadState = TypeParam::initialState;
+
+        LockUnderTest lock;
+        std::thread secondThread;
+        std::atomic<bool> started = false;
+        std::atomic<int32_t> protectedCounter = 0;
+        mm::ThreadData* secondThreadData = nullptr;
+
+        {
+            std::unique_lock guard1(lock);
+            secondThread = std::thread([&lock, &started, &protectedCounter, &secondThreadData, &initialThreadState]() {
+                ScopedMemoryInit init;
+                SwitchThreadState(init.memoryState(), initialThreadState, /* reentrant = */ true);
+
+                secondThreadData = init.memoryState()->GetThreadData();
+                ASSERT_EQ(secondThreadData->state(), initialThreadState);
+
+                started = true;
+                std::unique_lock guard2(lock);
+                EXPECT_EQ(secondThreadData->state(), initialThreadState);
+                protectedCounter++;
+            });
+
+            protectedCounter++;
+            while (!started) {
+                std::this_thread::yield();
+            }
+
+            // Wait to give the second thread a chance to try to acquire the lock.
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+            EXPECT_EQ(protectedCounter, 1);
+
+            auto expectedSecondThreadState =
+                    (TypeParam::threadStateHandling == MutexThreadStateHandling::kSwitchIfRegistered)
+                    ? ThreadState::kNative
+                    : initialThreadState;
+            EXPECT_EQ(secondThreadData->state(), expectedSecondThreadState);
+        }
+        secondThread.join();
+        EXPECT_EQ(protectedCounter, 2);
+    });
+}

kotlin-native/runtime/src/mm/cpp/ObjectFactory.hpp

@@ -163,7 +163,7 @@ public:
                 return;
             }
 
-            std::lock_guard<SpinLock> guard(owner_.mutex_);
+            std::lock_guard guard(owner_.mutex_);
 
             owner_.AssertCorrectUnsafe();
 
@@ -328,7 +328,7 @@ public:
 
     private:
         ObjectFactoryStorage& owner_; // weak
-        std::unique_lock<SpinLock> guard_;
+        std::unique_lock<SpinLock<MutexThreadStateHandling::kIgnore>> guard_;
     };
 
     ~ObjectFactoryStorage() {
@@ -389,7 +389,7 @@ private:
     unique_ptr<Node> root_;
     Node* last_ = nullptr;
     size_t size_ = 0;
-    SpinLock mutex_;
+    SpinLock<MutexThreadStateHandling::kIgnore> mutex_;
 };
 
 class SimpleAllocator {

kotlin-native/runtime/src/mm/cpp/StableRefRegistry.hpp

@@ -15,16 +15,18 @@ namespace mm {
 
 // Registry for all objects that have references outside of Kotlin.
 class StableRefRegistry : Pinned {
+    using Mutex = SpinLock<MutexThreadStateHandling::kIgnore>;
+
 public:
-    class ThreadQueue : public MultiSourceQueue<ObjHeader*>::Producer {
+    class ThreadQueue : public MultiSourceQueue<ObjHeader*, Mutex>::Producer {
     public:
         explicit ThreadQueue(StableRefRegistry& registry) : Producer(registry.stableRefs_) {}
         // Do not add fields as this is just a wrapper and Producer does not have virtual destructor.
     };
 
-    using Iterable = MultiSourceQueue<ObjHeader*>::Iterable;
+    using Iterable = MultiSourceQueue<ObjHeader*, Mutex>::Iterable;
-    using Iterator = MultiSourceQueue<ObjHeader*>::Iterator;
+    using Iterator = MultiSourceQueue<ObjHeader*, Mutex>::Iterator;
-    using Node = MultiSourceQueue<ObjHeader*>::Node;
+    using Node = MultiSourceQueue<ObjHeader*, Mutex>::Node;
 
     StableRefRegistry();
     ~StableRefRegistry();
@@ -64,7 +66,7 @@ private:
     //   before posting queue to the global registry)
     //
     // TODO: Measure to understand, if this approach is problematic.
-    MultiSourceQueue<ObjHeader*> stableRefs_;
+    MultiSourceQueue<ObjHeader*, Mutex> stableRefs_;
 };
 
 } // namespace mm