[K/N][New MM] Disallow getting thread data for detached threads

This patch fixes a performance degradation introduced by
99bd26c2ef (Switch thread states
in termination handlers).

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

@@ -3739,9 +3739,14 @@ ALWAYS_INLINE void kotlin::AssertThreadState(MemoryState* thread, std::initializ
 }
 
 MemoryState* kotlin::mm::GetMemoryState() noexcept {
+    RuntimeAssert(memoryState != nullptr, "Thread is not attached to the runtime");
     return ::memoryState;
 }
 
+bool kotlin::mm::IsCurrentThreadRegistered() noexcept {
+    return ::memoryState != nullptr;
+}
+
 kotlin::ThreadState kotlin::GetThreadState(MemoryState* thread) noexcept {
     // Assume that we are always in the Runnable thread state.
     return ThreadState::kRunnable;

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

@@ -47,10 +47,10 @@ 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::GetMemoryState).
+// 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) {
-    auto* memoryState = kotlin::mm::GetMemoryState();
+    if (kotlin::mm::IsCurrentThreadRegistered()) {
-    if (memoryState) {
         return kotlin::ThreadStateGuard(kotlin::ThreadState::kNative, reentrant);
     } else {
         // The current thread is not registered in the Kotlin runtime,

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

@@ -337,10 +337,10 @@ void setupMocks(bool expectRegisteredThread = true) noexcept {
     ON_CALL(nativeHandlerMock, Call)
             .WillByDefault([expectRegisteredThread]() {
                 if (expectRegisteredThread) {
-                    loggingAssert(mm::GetMemoryState() != nullptr, "Expected registered thread in the native handler");
+                    loggingAssert(mm::IsCurrentThreadRegistered(), "Expected registered thread in the native handler");
                     loggingAssert(GetThreadState() == ThreadState::kNative, "Expected kNative thread state in the native handler");
                 } else {
-                    loggingAssert(mm::GetMemoryState() == nullptr, "Expected unregistered thread in the native handler");
+                    loggingAssert(!mm::IsCurrentThreadRegistered(), "Expected unregistered thread in the native handler");
                 }
                 log("Native handler");
             });
@@ -399,7 +399,7 @@ TEST(TerminationThreadStateDeathTest, TerminationInForeignThread) {
     auto testBlock = []() {
         setupMocks(/* expectRegisteredThread = */ false);
 
-        loggingAssert(mm::GetMemoryState() == nullptr, "Expected unregistered thread before std::terminate");
+        loggingAssert(!mm::IsCurrentThreadRegistered(), "Expected unregistered thread before std::terminate");
         std::terminate();
     };
 
@@ -458,7 +458,7 @@ TEST(TerminationThreadStateDeathTest, UnhandledKotlinExceptionInForeignThread) {
         // It is possible if a Kotlin exception thrown by a Kotlin callback is re-thrown in
         // another thread which is not attached to the Kotlin runtime at all.
         std::thread foreignThread([]() {
-            loggingAssert(mm::GetMemoryState() == nullptr, "Expected unregistered thread before throwing");
+            loggingAssert(!mm::IsCurrentThreadRegistered(), "Expected unregistered thread before throwing");
 
             auto future = std::async(std::launch::async, []() {
                 // Initial Kotlin exception throwing requires the runtime to be initialized.
@@ -500,7 +500,7 @@ TEST(TerminationThreadStateDeathTest, UnhandledForeignExceptionInForeignThread)
         setupMocks(/* expectRegisteredThread = */ false);
 
         std::thread foreignThread([]() {
-            loggingAssert(mm::GetMemoryState() == nullptr, "Expected unregistered thread before throwing");
+            loggingAssert(!mm::IsCurrentThreadRegistered(), "Expected unregistered thread before throwing");
             throw std::runtime_error("Foreign exception");
         });
         foreignThread.join();

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

@@ -389,11 +389,17 @@ namespace kotlin {
 namespace mm {
 
 // Returns the MemoryState for the current thread.
-// If the memory subsystem isn't initialized for the current thread, returns nullptr.
+// The current thread must be attached to the runtime.
 // Try not to use it very often, as (1) thread local access can be slow on some platforms,
 // (2) TLS gets deallocated before our thread destruction hooks run.
 MemoryState* GetMemoryState() noexcept;
 
+
+// TODO: Replace with direct access to ThreadRegistry when the legacy MM is gone.
+// Checks if the current thread is attached to the runtime.
+// This function accesses a TLS variable, so it must not be called from a thread destructor.
+bool IsCurrentThreadRegistered() noexcept;
+
 } // namespace mm
 
 enum class ThreadState {

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

@@ -11,8 +11,8 @@
 
 using namespace kotlin;
 
-TEST(MemoryStateTest, GetMemoryStateForUnregisteredThread) {
+TEST(MemoryStateTestDeathTest, GetMemoryStateForUnregisteredThread) {
-    EXPECT_EQ(mm::GetMemoryState(), nullptr);
+    EXPECT_DEATH(mm::GetMemoryState(), "Thread is not attached to the runtime");
 }
 
 TEST(MemoryStateTest, GetMemoryStateForRegisteredThread) {
@@ -22,3 +22,10 @@ TEST(MemoryStateTest, GetMemoryStateForRegisteredThread) {
         EXPECT_EQ(actualState, expectedState);
     });
 }
+
+TEST(MemoryStateTest, IsCurrentThreadRegistered) {
+    EXPECT_FALSE(mm::IsCurrentThreadRegistered());
+    RunInNewThread([]() {
+        EXPECT_TRUE(mm::IsCurrentThreadRegistered());
+    });
+}

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

@@ -541,6 +541,10 @@ MemoryState* kotlin::mm::GetMemoryState() noexcept {
     return ToMemoryState(ThreadRegistry::Instance().CurrentThreadDataNode());
 }
 
+bool kotlin::mm::IsCurrentThreadRegistered() noexcept {
+    return ThreadRegistry::Instance().IsCurrentThreadRegistered();
+}
+
 ALWAYS_INLINE kotlin::CalledFromNativeGuard::CalledFromNativeGuard(bool reentrant) noexcept : reentrant_(reentrant) {
     Kotlin_initRuntimeIfNeeded();
     thread_ = mm::GetMemoryState();

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

@@ -41,8 +41,7 @@ std::unique_lock<mm::ThreadRegistry::Mutex> mm::ThreadRegistry::Lock() noexcept
 }
 
 ALWAYS_INLINE mm::ThreadData* mm::ThreadRegistry::CurrentThreadData() const noexcept {
-    auto* threadDataNode = CurrentThreadDataNode();
+    return CurrentThreadDataNode()->Get();
-    return threadDataNode ? threadDataNode->Get() : nullptr;
 }
 
 mm::ThreadRegistry::ThreadRegistry() = default;

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

@@ -38,8 +38,12 @@ public:
     // Try not to use these methods very often, as (1) thread local access can be slow on some platforms,
     // (2) TLS gets deallocated before our thread destruction hooks run.
     // Using this after `Unregister` for the thread has been called is undefined behaviour.
+    // Using this by a thread which is not attached to the Kotlin runtime is undefined behaviour.
     ALWAYS_INLINE ThreadData* CurrentThreadData() const noexcept;
-    Node* CurrentThreadDataNode() const noexcept { return currentThreadDataNode_; }
+    Node* CurrentThreadDataNode() const noexcept {
+        RuntimeAssert(currentThreadDataNode_ != nullptr, "Thread is not attached to the runtime");
+        return currentThreadDataNode_;
+    }
 
     bool IsCurrentThreadRegistered() const noexcept { return currentThreadDataNode_ != nullptr; }
 

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

@@ -175,14 +175,14 @@ TEST(ThreadStateDeathTest, StateAssertsForDetachedThread) {
     EXPECT_DEATH(AssertThreadState(static_cast<mm::ThreadData*>(nullptr), ThreadState::kNative),
                  "runtime assert: threadData must not be nullptr");
     EXPECT_DEATH(AssertThreadState(ThreadState::kNative),
-                 "runtime assert: thread must not be nullptr");
+                 "runtime assert: Thread is not attached to the runtime");
 
     EXPECT_DEATH(AssertThreadState(static_cast<MemoryState*>(nullptr), {ThreadState::kNative}),
                  "runtime assert: thread must not be nullptr");
     EXPECT_DEATH(AssertThreadState(static_cast<mm::ThreadData*>(nullptr), {ThreadState::kNative}),
                  "runtime assert: threadData must not be nullptr");
     EXPECT_DEATH(AssertThreadState({ThreadState::kNative}),
-                 "runtime assert: thread must not be nullptr");
+                 "runtime assert: Thread is not attached to the runtime");
 
 }
 
@@ -223,8 +223,8 @@ TEST(ThreadStateDeathTest, StateSwitchForDetachedThread) {
     EXPECT_DEATH(SwitchThreadState(static_cast<MemoryState*>(nullptr), ThreadState::kNative), "thread must not be nullptr");
     EXPECT_DEATH(SwitchThreadState(static_cast<mm::ThreadData*>(nullptr), ThreadState::kNative), "threadData must not be nullptr");
 
-    EXPECT_DEATH(Kotlin_mm_switchThreadStateNative(), "threadData must not be nullptr");
+    EXPECT_DEATH(Kotlin_mm_switchThreadStateNative(), "Thread is not attached to the runtime");
-    EXPECT_DEATH(Kotlin_mm_switchThreadStateRunnable(), "threadData must not be nullptr" );
+    EXPECT_DEATH(Kotlin_mm_switchThreadStateRunnable(), "Thread is not attached to the runtime" );
 }
 
 TEST(ThreadStateDeathTest, ReentrantStateSwitch) {
@@ -267,6 +267,7 @@ TEST(ThreadStateDeathTest, GuardForDetachedThread) {
     EXPECT_DEATH({ ThreadStateGuard guard(nullptr, ThreadState::kRunnable, true); }, expectedError);
     EXPECT_DEATH({ ThreadStateGuard guard(nullptr, ThreadState::kNative, true); }, expectedError);
 
+    expectedError = "Thread is not attached to the runtime";
     EXPECT_DEATH({ ThreadStateGuard guard(ThreadState::kRunnable, false); }, expectedError);
     EXPECT_DEATH({ ThreadStateGuard guard(ThreadState::kNative, false); }, expectedError);
     EXPECT_DEATH({ ThreadStateGuard guard(ThreadState::kRunnable, true); }, expectedError);