[K/N][Runtime] Separate thread registering and making it 'Runnable'

The new GC will require threads to suspend right after registering
if Stop-The-World is requested. This patch changes the initial
thread state to kNative and adds a separate state switch right
after thread registering. This switch suspends if it is necessary.

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

@@ -243,6 +243,10 @@ class {
 
 //! Process exception hook (if any) or just printStackTrace + write crash log
 void processUnhandledKotlinException(KRef throwable) {
+  // Use the reentrant switch because both states are possible here:
+  //  - runnable, if the exception occured in a pure Kotlin thread (except initialization of globals).
+  //  - native, if the throwing code was called from ObjC/Swift or if the exception occured during initialization of globals.
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kRunnable, /* reentrant = */ true);
   OnUnhandledException(throwable);
 #if KONAN_REPORT_BACKTRACE_TO_IOS_CRASH_LOG
   ReportBacktraceToIosCrashLog(throwable);

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

@@ -422,8 +422,11 @@ ALWAYS_INLINE inline void AssertThreadState(std::initializer_list<ThreadState> e
 }
 
 // Scopely sets the given thread state for the given thread.
-class ThreadStateGuard final : private Pinned {
+class ThreadStateGuard final : private MoveOnly {
 public:
+    // Do not set any state. Useful to create a variable to move another guard into.
+    ThreadStateGuard() : thread_(nullptr), oldState_(ThreadState::kNative), reentrant_(false) {}
+
     // Set the state for the given thread.
     ThreadStateGuard(MemoryState* thread, ThreadState state, bool reentrant = false) noexcept : thread_(thread), reentrant_(reentrant) {
         oldState_ = SwitchThreadState(thread_, state, reentrant_);
@@ -433,9 +436,25 @@ public:
     explicit ThreadStateGuard(ThreadState state, bool reentrant = false) noexcept
         : ThreadStateGuard(mm::GetMemoryState(), state, reentrant) {};
 
-    ~ThreadStateGuard() noexcept {
-        SwitchThreadState(thread_, oldState_, reentrant_);
+    ThreadStateGuard(ThreadStateGuard&& other) noexcept
+        : thread_(other.thread_), oldState_(other.oldState_), reentrant_(other.reentrant_) {
+        other.thread_ = nullptr;
     }
+
+    ~ThreadStateGuard() noexcept {
+        if (thread_ != nullptr) {
+            SwitchThreadState(thread_, oldState_, reentrant_);
+        }
+    }
+
+    ThreadStateGuard& operator=(ThreadStateGuard&& other) noexcept {
+        thread_ = other.thread_;
+        oldState_ = other.oldState_;
+        reentrant_ = other.reentrant_;
+        other.thread_ = nullptr;
+        return *this;
+    }
+
 private:
     MemoryState* thread_;
     ThreadState oldState_;

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

@@ -5,12 +5,11 @@
 
 #include "ObjectTestSupport.hpp"
 
-#include <thread>
-
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 #include "Natives.h"
+#include "TestSupport.hpp"
 
 using namespace kotlin;
 
@@ -68,17 +67,6 @@ class ObjectTestSupportObjectTest : public testing::Test {};
 using ObjectTestCases = testing::Types<RegularObjectTestCase, IrregularObjectTestCase>;
 TYPED_TEST_SUITE(ObjectTestSupportObjectTest, ObjectTestCases, ObjectTestCaseNames);
 
-// TODO: Replace with common implementation.
-template <typename F>
-void RunInNewThread(F f) {
-    std::thread([&f]() {
-        auto* memory = InitMemory(false);
-        f();
-        ClearMemoryForTests(memory);
-        DeinitMemory(memory, false);
-    }).join();
-}
-
 template <typename Payload>
 KStdVector<ObjHeader**> Collect(test_support::Object<Payload>& object) {
     KStdVector<ObjHeader**> result;

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

@@ -106,10 +106,15 @@ RuntimeState* initRuntime() {
   ::runtimeState = result;
 
   bool firstRuntime = false;
+  // We set this guard in the `switch` below, after memory initialization.
+  kotlin::ThreadStateGuard stateGuard;
   switch (Kotlin_getDestroyRuntimeMode()) {
       case DESTROY_RUNTIME_LEGACY:
           compareAndSwap(&globalRuntimeStatus, kGlobalRuntimeUninitialized, kGlobalRuntimeRunning);
           result->memoryState = InitMemory(false); // The argument will be ignored for legacy DestroyRuntimeMode
+          // Switch thread state because worker and globals inits require the runnable state.
+          // This call may block if GC requested suspending threads.
+          stateGuard = kotlin::ThreadStateGuard(result->memoryState, kotlin::ThreadState::kRunnable);
           result->worker = WorkerInit(result->memoryState, true);
           firstRuntime = atomicAdd(&aliveRuntimesCount, 1) == 1;
           if (!kotlin::kSupportsMultipleMutators && !firstRuntime) {
@@ -131,6 +136,9 @@ RuntimeState* initRuntime() {
               konan::abort();
           }
           result->memoryState = InitMemory(firstRuntime);
+          // Switch thread state because worker and globals inits require the runnable state.
+          // This call may block if GC requested suspending threads.
+          stateGuard = kotlin::ThreadStateGuard(result->memoryState, kotlin::ThreadState::kRunnable);
           result->worker = WorkerInit(result->memoryState, true);
   }
 
@@ -148,8 +156,6 @@ RuntimeState* initRuntime() {
   RuntimeAssert(result->status == RuntimeStatus::kUninitialized, "Runtime must still be in the uninitialized state");
   result->status = RuntimeStatus::kRunning;
 
-  kotlin::SwitchThreadState(result->memoryState, kotlin::ThreadState::kNative);
-
   return result;
 }
 
@@ -173,8 +179,14 @@ void deinitRuntime(RuntimeState* state, bool destroyRuntime) {
   ClearTLS(state->memoryState);
   if (destroyRuntime)
     InitOrDeinitGlobalVariables(DEINIT_GLOBALS, state->memoryState);
+
+  // Worker deinit must be performed in the runnable state because
+  // Worker's destructor unregisters stable refs.
   auto workerId = GetWorkerId(state->worker);
   WorkerDeinit(state->worker);
+
+  // Do not use ThreadStateGuard because memoryState will be destroyed during DeinitMemory.
+  kotlin::SwitchThreadState(state->memoryState, kotlin::ThreadState::kNative);
   DeinitMemory(state->memoryState, destroyRuntime);
   konanDestructInstance(state);
   WorkerDestroyThreadDataIfNeeded(workerId);

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

@@ -27,9 +27,13 @@ void DeinitMemoryForTests(MemoryState* memoryState);
 // Scopely initializes the memory subsystem of the current thread for tests.
 class ScopedMemoryInit : private kotlin::Pinned {
 public:
-    ScopedMemoryInit() : memoryState_(InitMemoryForTests()) {}
+    ScopedMemoryInit() : memoryState_(InitMemoryForTests()) {
+        kotlin::SwitchThreadState(memoryState(), ThreadState::kRunnable);
+    }
     ~ScopedMemoryInit() {
         ClearMemoryForTests(memoryState());
+        // Ensure that memory deinit is performed in the native state.
+        SwitchThreadState(memoryState(), ThreadState::kNative, /* reentrant = */ true);
         DeinitMemoryForTests(memoryState());
     }
 

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

@@ -60,7 +60,7 @@ protected:
 // the variable it works with to avoid invalid memory access.
 template<typename T1, typename T2>
 class AutoReset final : private Pinned {
-    static_assert(std::is_assignable<T1, T2>::value);
+    static_assert(std::is_assignable<T1&, T2>::value);
 
 public:
     AutoReset(T1* variable, T2 value) : variable_(variable), oldValue_(*variable) {

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

@@ -95,8 +95,13 @@ extern "C" MemoryState* InitMemory(bool firstRuntime) {
 }
 
 extern "C" void DeinitMemory(MemoryState* state, bool destroyRuntime) {
+    // We need the native state to avoid a deadlock on unregistering the thread.
+    // The deadlock is possible if we are in the runnable state and the GC already locked
+    // the thread registery and waits for threads to suspend or go to the native state.
+    AssertThreadState(state, ThreadState::kNative);
     auto* node = mm::FromMemoryState(state);
     if (destroyRuntime) {
+        ThreadStateGuard guard(state, ThreadState::kRunnable);
         node->Get()->gc().PerformFullGC();
         // TODO: Also make sure that finalizers are run.
     }

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

@@ -35,7 +35,7 @@ public:
         stableRefThreadQueue_(StableRefRegistry::Instance()),
         gc_(GlobalData::Instance().gc(), *this),
         objectFactoryThreadQueue_(GlobalData::Instance().objectFactory(), gc_),
-        suspensionData_(ThreadState::kRunnable) {}
+        suspensionData_(ThreadState::kNative) {}
 
     ~ThreadData() = default;
 

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

@@ -7,6 +7,7 @@
 
 #include "GlobalData.hpp"
 #include "ThreadData.hpp"
+#include "ThreadState.hpp"
 
 using namespace kotlin;
 
@@ -17,6 +18,7 @@ mm::ThreadRegistry& mm::ThreadRegistry::Instance() noexcept {
 
 mm::ThreadRegistry::Node* mm::ThreadRegistry::RegisterCurrentThread() noexcept {
     auto* threadDataNode = list_.Emplace(pthread_self());
+    AssertThreadState(threadDataNode->Get(), ThreadState::kNative);
     Node*& currentDataNode = currentThreadDataNode_;
     RuntimeAssert(currentDataNode == nullptr, "This thread already had some data assigned to it.");
     currentDataNode = threadDataNode;
@@ -24,6 +26,7 @@ mm::ThreadRegistry::Node* mm::ThreadRegistry::RegisterCurrentThread() noexcept {
 }
 
 void mm::ThreadRegistry::Unregister(Node* threadDataNode) noexcept {
+    AssertThreadState(threadDataNode->Get(), ThreadState::kNative);
     list_.Erase(threadDataNode);
     // Do not touch `currentThreadData_` as TLS may already have been deallocated.
 }

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

@@ -139,6 +139,24 @@ TEST_F(ThreadStateTest, CallWithRunnableState) {
     });
 }
 
+TEST_F(ThreadStateTest, MovingGuard) {
+    RunInNewThread([](MemoryState* memoryState) {
+        auto& threadData = *memoryState->GetThreadData();
+        ASSERT_EQ(threadData.state(), ThreadState::kRunnable);
+        {
+            ThreadStateGuard outerGuard;
+            EXPECT_EQ(threadData.state(), ThreadState::kRunnable);
+            {
+                ThreadStateGuard innerGuard(memoryState, ThreadState::kNative);
+                EXPECT_EQ(threadData.state(), ThreadState::kNative);
+                outerGuard = std::move(innerGuard);
+            }
+            EXPECT_EQ(threadData.state(), ThreadState::kNative);
+        }
+        EXPECT_EQ(threadData.state(), ThreadState::kRunnable);
+    });
+}
+
 TEST(ThreadStateDeathTest, StateAsserts) {
     RunInNewThread([](MemoryState* memoryState) {
         mm::ThreadData* threadData = memoryState->GetThreadData();
@@ -192,4 +210,27 @@ TEST(ThreadStateDeathTest, ReentrantStateSwitch) {
                     testing::ExitedWithCode(0),
                     testing::Not(testing::ContainsRegex("runtime assert: Illegal thread state switch.")));
     });
+}
+
+TEST(ThreadStateDeathTest, MovingReentrantGuard) {
+    RunInNewThread([](MemoryState* memoryState) {
+
+        auto blockUnderTest = [&memoryState]() {
+            auto& threadData = *memoryState->GetThreadData();
+            ASSERT_EQ(threadData.state(), ThreadState::kRunnable);
+            {
+                ThreadStateGuard outerGuard;
+                {
+                    ThreadStateGuard innerGuard(memoryState, ThreadState::kRunnable, /* reentrant = */ true);
+                    outerGuard = std::move(innerGuard);
+                }
+            }
+            EXPECT_EQ(threadData.state(), ThreadState::kRunnable);
+            exit(0);
+        };
+
+        EXPECT_EXIT({ blockUnderTest(); },
+                    testing::ExitedWithCode(0),
+                    testing::Not(testing::ContainsRegex("runtime assert: Illegal thread state switch.")));
+    });
 }

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

@@ -3,8 +3,8 @@
  * that can be found in the LICENSE file.
  */
 
-#ifndef RUNTIME_MM_THREAD_SUSPENSION_UTILS_H
-#define RUNTIME_MM_THREAD_SUSPENSION_UTILS_H
+#ifndef RUNTIME_MM_THREAD_SUSPENSION_H
+#define RUNTIME_MM_THREAD_SUSPENSION_H
 
 #include <atomic>
 
@@ -57,4 +57,4 @@ void ResumeThreads() noexcept;
 } // namespace mm
 } // namespace kotlin
 
-#endif // RUNTIME_MM_THREAD_SUSPENSION_UTILS_H
+#endif // RUNTIME_MM_THREAD_SUSPENSION_H