[K/N] Extract ThreadRegistry::waitAllThreads

kotlin-native/runtime/src/gc/cms/cpp/Barriers.cpp

@@ -44,15 +44,11 @@ void waitForThreadsToReachCheckpoint() {
 
     mm::SafePointActivator safePointActivator;
 
-    // Disable new threads coming and going.
+    // Wait for all threads to either have passed safepoint or to be in the native state.
-    auto threads = mm::ThreadRegistry::Instance().LockForIter();
-    // And wait for all threads to either have passed safepoint or to be in the native state.
     // Either of these mean that none of them are inside a weak reference accessing code.
-    while (!std::all_of(threads.begin(), threads.end(), [](mm::ThreadData& thread) noexcept {
+    mm::ThreadRegistry::Instance().waitAllThreads([](mm::ThreadData& thread) noexcept {
         return thread.gc().impl().gc().barriers().visitedCheckpoint() || thread.suspensionData().suspendedOrNative();
-    })) {
+    });
-        std::this_thread::yield();
-    }
 }
 
 } // namespace

kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMarkAndSweep.cpp

@@ -162,7 +162,6 @@ void gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     RuntimeAssert(didSuspend, "Only GC thread can request suspension");
     gcHandle.suspensionRequested();
 
-    RuntimeAssert(!kotlin::mm::IsCurrentThreadRegistered(), "GC must run on unregistered thread");
     WaitForThreadsReadyToMark();
     gcHandle.threadsAreSuspended();
 
@@ -236,41 +235,16 @@ void gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     finalizerProcessor_.ScheduleTasks(std::move(finalizerQueue), epoch);
 }
 
-namespace {
-    bool isSuspendedOrNative(kotlin::mm::ThreadData& thread) noexcept {
-        auto& suspensionData = thread.suspensionData();
-        return suspensionData.suspended() || suspensionData.state() == kotlin::ThreadState::kNative;
-    }
-
-    template <typename F>
-    bool allThreads(F predicate) noexcept {
-        auto& threadRegistry = kotlin::mm::ThreadRegistry::Instance();
-        auto* currentThread = (threadRegistry.IsCurrentThreadRegistered()) ? threadRegistry.CurrentThreadData() : nullptr;
-        kotlin::mm::ThreadRegistry::Iterable threads = kotlin::mm::ThreadRegistry::Instance().LockForIter();
-        for (auto& thread : threads) {
-            // Handle if suspension was initiated by the mutator thread.
-            if (&thread == currentThread) continue;
-            if (!predicate(thread)) {
-                return false;
-            }
-        }
-        return true;
-    }
-
-    void yield() noexcept {
-        std::this_thread::yield();
-    }
-} // namespace
-
 void gc::ConcurrentMarkAndSweep::SetMarkingRequested(uint64_t epoch) noexcept {
     markingRequested = markingBehavior_ == MarkingBehavior::kMarkOwnStack;
     markingEpoch = epoch;
 }
 
 void gc::ConcurrentMarkAndSweep::WaitForThreadsReadyToMark() noexcept {
-    while(!allThreads([](kotlin::mm::ThreadData& thread) { return isSuspendedOrNative(thread) || thread.gc().impl().gc().marking_.load(); })) {
+    RuntimeAssert(!kotlin::mm::IsCurrentThreadRegistered(), "GC must run on unregistered thread");
-        yield();
+    mm::ThreadRegistry::Instance().waitAllThreads([](mm::ThreadData& thread) noexcept {
-    }
+        return thread.suspensionData().suspendedOrNative() || thread.gc().impl().gc().marking_.load();
+    });
 }
 
 void gc::ConcurrentMarkAndSweep::CollectRootSetAndStartMarking(GCHandle gcHandle) noexcept {

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

@@ -50,6 +50,15 @@ public:
 
     static void ClearCurrentThreadData() { currentThreadDataNode_ = nullptr; }
 
+    template <typename F>
+    void waitAllThreads(F&& f) noexcept {
+        // Disable new threads coming and going.
+        auto iter = LockForIter();
+        while (!std::all_of(iter.begin(), iter.end(), std::forward<F>(f))) {
+            std::this_thread::yield();
+        }
+    }
+
 private:
     friend class GlobalData;
 

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

@@ -19,26 +19,6 @@ using namespace kotlin;
 
 namespace {
 
-template<typename F>
-bool allThreads(F predicate) noexcept {
-    auto& threadRegistry = kotlin::mm::ThreadRegistry::Instance();
-    auto* currentThread = (threadRegistry.IsCurrentThreadRegistered()) ? threadRegistry.CurrentThreadData() : nullptr;
-    kotlin::mm::ThreadRegistry::Iterable threads = kotlin::mm::ThreadRegistry::Instance().LockForIter();
-    for (auto& thread : threads) {
-        // Handle if suspension was initiated by the mutator thread.
-        if (&thread == currentThread)
-            continue;
-        if (!predicate(thread)) {
-            return false;
-        }
-    }
-    return true;
-}
-
-void yield() noexcept {
-    std::this_thread::yield();
-}
-
 [[clang::no_destroy]] thread_local std::optional<mm::SafePointActivator> gSafePointActivator = std::nullopt;
 [[clang::no_destroy]] std::mutex gSuspensionMutex;
 [[clang::no_destroy]] std::condition_variable gSuspensionCondVar;
@@ -89,10 +69,12 @@ bool kotlin::mm::RequestThreadsSuspension() noexcept {
 }
 
 void kotlin::mm::WaitForThreadsSuspension() noexcept {
+    auto& threadRegistry = kotlin::mm::ThreadRegistry::Instance();
+    auto* currentThread = (threadRegistry.IsCurrentThreadRegistered()) ? threadRegistry.CurrentThreadData() : nullptr;
     // Spin waiting for threads to suspend. Ignore Native threads.
-    while (!allThreads([](mm::ThreadData& thread) { return thread.suspensionData().suspendedOrNative(); })) {
+    threadRegistry.waitAllThreads([currentThread](mm::ThreadData& thread) noexcept {
-        yield();
+        return &thread == currentThread || thread.suspensionData().suspendedOrNative();
-    }
+    });
 }
 
 void kotlin::mm::ResumeThreads() noexcept {