[K/N] Generalize GC-assistance pause logic

Merge-request: KT-MR-11790
Merged-by: Alexey Glushko <aleksei.glushko@jetbrains.com>

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

@@ -71,14 +71,6 @@ bool gc::ConcurrentMarkAndSweep::ThreadData::tryLockRootSet() {
     return locked;
 }
 
-void gc::ConcurrentMarkAndSweep::ThreadData::beginCooperation() {
-    cooperative_.store(true, std::memory_order_release);
-}
-
-bool gc::ConcurrentMarkAndSweep::ThreadData::cooperative() const {
-    return cooperative_.load(std::memory_order_relaxed);
-}
-
 void gc::ConcurrentMarkAndSweep::ThreadData::publish() {
     threadData_.Publish();
     published_.store(true, std::memory_order_release);
@@ -90,7 +82,6 @@ bool gc::ConcurrentMarkAndSweep::ThreadData::published() const {
 
 void gc::ConcurrentMarkAndSweep::ThreadData::clearMarkFlags() {
     published_.store(false, std::memory_order_relaxed);
-    cooperative_.store(false, std::memory_order_relaxed);
     rootSetLocked_.store(false, std::memory_order_release);
 }
 
@@ -170,14 +161,7 @@ void gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     markDispatcher_.beginMarkingEpoch(gcHandle);
     GCLogDebug(epoch, "Main GC requested marking in mutators");
 
-    // Request STW
+    stopTheWorld(gcHandle);
-    bool didSuspend = mm::RequestThreadsSuspension();
-    RuntimeAssert(didSuspend, "Only GC thread can request suspension");
-    gcHandle.suspensionRequested();
-
-    markDispatcher_.waitForThreadsPauseMutation();
-    GCLogDebug(epoch, "All threads have paused mutation");
-    gcHandle.threadsAreSuspended();
 
     auto& scheduler = gcScheduler_;
     scheduler.onGCStart();
@@ -188,26 +172,17 @@ void gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
 
     markDispatcher_.endMarkingEpoch();
 
-    mm::WaitForThreadsSuspension();
-
     if (compiler::concurrentWeakSweep()) {
-        EnableWeakRefBarriers(epoch);
+        // Expected to happen inside STW.
-
+        gc::EnableWeakRefBarriers(epoch);
-        mm::ResumeThreads();
+        resumeTheWorld(gcHandle);
-        gcHandle.threadsAreResumed();
     }
 
     gc::processWeaks<DefaultProcessWeaksTraits>(gcHandle, mm::SpecialRefRegistry::instance());
 
     if (compiler::concurrentWeakSweep()) {
-        bool didSuspend = mm::RequestThreadsSuspension();
+        stopTheWorld(gcHandle);
-        RuntimeAssert(didSuspend, "Only GC thread can request suspension");
+        gc::DisableWeakRefBarriers();
-        gcHandle.suspensionRequested();
-
-        mm::WaitForThreadsSuspension();
-        GCLogDebug(gcHandle.getEpoch(), "All threads have paused mutation");
-        gcHandle.threadsAreSuspended();
-        DisableWeakRefBarriers();
     }
 
     // TODO outline as mark_.isolateMarkedHeapAndFinishMark()
@@ -233,8 +208,7 @@ void gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     checkMarkCorrectness(*objectFactoryIterable);
 #endif
 
-    mm::ResumeThreads();
+    resumeTheWorld(gcHandle);
-    gcHandle.threadsAreResumed();
 
 #ifndef CUSTOM_ALLOCATOR
     gc::SweepExtraObjects<DefaultSweepTraits<ObjectFactory>>(gcHandle, *extraObjectFactoryIterable);

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

@@ -45,8 +45,6 @@ public:
         BarriersThreadData& barriers() noexcept { return barriers_; }
 
         bool tryLockRootSet();
-        void beginCooperation();
-        bool cooperative() const;
         void publish();
         bool published() const;
         void clearMarkFlags();
@@ -61,7 +59,6 @@ public:
 
         std::atomic<bool> rootSetLocked_ = false;
         std::atomic<bool> published_ = false;
-        std::atomic<bool> cooperative_ = false;
     };
 
 #ifdef CUSTOM_ALLOCATOR

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

@@ -85,15 +85,6 @@ void gc::mark::ParallelMark::beginMarkingEpoch(gc::GCHandle gcHandle) {
     }
 }
 
-void gc::mark::ParallelMark::waitForThreadsPauseMutation() noexcept {
-    RuntimeAssert(!kotlin::mm::IsCurrentThreadRegistered(), "Dispatcher thread must not be registered");
-    spinWait([this] {
-        return allMutators([](mm::ThreadData& mut) {
-            return mm::isSuspendedOrNative(mut) || mut.gc().impl().gc().cooperative();
-        });
-    });
-}
-
 void gc::mark::ParallelMark::endMarkingEpoch() {
     if (!compiler::gcMarkSingleThreaded()) {
         // We must now wait for every worker to finish the Mark procedure:
@@ -139,8 +130,6 @@ void gc::mark::ParallelMark::runOnMutator(mm::ThreadData& mutatorThread) {
     auto epoch = gcHandle().getEpoch();
     auto parallelWorker = createWorker();
     if (parallelWorker) {
-        auto& gcData = mutatorThread.gc().impl().gc();
-        gcData.beginCooperation();
         GCLogDebug(epoch, "Mutator thread cooperates in marking");
 
         tryCollectRootSet(mutatorThread, *parallelWorker);

kotlin-native/runtime/src/gc/cms/cpp/ParallelMark.hpp

@@ -110,7 +110,6 @@ public:
     ParallelMark(bool mutatorsCooperate);
 
     void beginMarkingEpoch(gc::GCHandle gcHandle);
-    void waitForThreadsPauseMutation() noexcept;
     void endMarkingEpoch();
 
     /** To be run by a single "main" GC thread during STW. */

kotlin-native/runtime/src/gc/common/cpp/MarkAndSweepUtils.cpp

@@ -4,3 +4,17 @@
  */
 
 #include "MarkAndSweepUtils.hpp"
+
+void kotlin::gc::stopTheWorld(kotlin::gc::GCHandle gcHandle) noexcept {
+    bool didSuspend = mm::RequestThreadsSuspension();
+    RuntimeAssert(didSuspend, "Only GC thread can request suspension");
+    gcHandle.suspensionRequested();
+
+    mm::WaitForThreadsSuspension();
+    gcHandle.threadsAreSuspended();
+}
+
+void kotlin::gc::resumeTheWorld(kotlin::gc::GCHandle gcHandle) noexcept {
+    mm::ResumeThreads();
+    gcHandle.threadsAreResumed();
+}

kotlin-native/runtime/src/gc/common/cpp/MarkAndSweepUtils.hpp

@@ -254,6 +254,13 @@ struct DefaultProcessWeaksTraits {
     static bool IsMarked(ObjHeader* obj) noexcept { return gc::isMarked(obj); }
 };
 
+void stopTheWorld(GCHandle gcHandle) noexcept;
+void resumeTheWorld(GCHandle gcHandle) noexcept;
+
+[[nodiscard]] inline auto stopTheWorldInScope(GCHandle gcHandle) noexcept {
+    return ScopeGuard([=]() { stopTheWorld(gcHandle); }, [=]() { resumeTheWorld(gcHandle); });
+}
+
 } // namespace gc
 } // namespace kotlin
 

kotlin-native/runtime/src/gc/stms/cpp/SameThreadMarkAndSweep.cpp

@@ -71,13 +71,8 @@ bool gc::SameThreadMarkAndSweep::FinalizersThreadIsRunning() noexcept {
 
 void gc::SameThreadMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     auto gcHandle = GCHandle::create(epoch);
-    bool didSuspend = mm::RequestThreadsSuspension();
-    RuntimeAssert(didSuspend, "Only GC thread can request suspension");
-    gcHandle.suspensionRequested();
 
-    RuntimeAssert(!kotlin::mm::IsCurrentThreadRegistered(), "GC must run on unregistered thread");
+    stopTheWorld(gcHandle);
-    mm::WaitForThreadsSuspension();
-    gcHandle.threadsAreSuspended();
 
     auto& scheduler = gcScheduler_;
     scheduler.onGCStart();
@@ -124,8 +119,8 @@ void gc::SameThreadMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
 
     scheduler.onGCFinish(epoch, allocatedBytes());
 
-    mm::ResumeThreads();
+    resumeTheWorld(gcHandle);
-    gcHandle.threadsAreResumed();
+
     state_.finish(epoch);
     gcHandle.finalizersScheduled(finalizerQueue.size());
     gcHandle.finished();

kotlin-native/runtime/src/gcScheduler/common/cpp/MutatorAssists.cpp

@@ -16,8 +16,7 @@ void gcScheduler::internal::MutatorAssists::ThreadData::safePoint() noexcept {
     Epoch epoch = owner_.assistsEpoch_.load(std::memory_order_acquire);
     auto noNeedToWait = [this, epoch] { return owner_.completedEpoch_.load(std::memory_order_acquire) >= epoch; };
     if (noNeedToWait()) return;
-    auto prevState = thread_.suspensionData().setStateNoSafePoint(ThreadState::kNative);
+    auto pauseHandle = thread_.suspensionData().pauseMutationInScope("assisting GC");
-    RuntimeAssert(prevState == ThreadState::kRunnable, "Expected runnable state");
     startedWaiting_.store(epoch * 2, std::memory_order_release);
     {
         std::unique_lock guard(owner_.m_);
@@ -27,8 +26,6 @@ void gcScheduler::internal::MutatorAssists::ThreadData::safePoint() noexcept {
     }
     startedWaiting_.store(epoch * 2 + 1, std::memory_order_release);
     // Not doing a safe point. We're a safe point.
-    prevState = thread_.suspensionData().setStateNoSafePoint(ThreadState::kRunnable);
-    RuntimeAssert(prevState == ThreadState::kNative, "Expected native state");
 }
 
 bool gcScheduler::internal::MutatorAssists::ThreadData::completedEpoch(Epoch epoch) const noexcept {

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

@@ -89,6 +89,7 @@ void VLog(Level level, std::initializer_list<const char*> tags, const char* form
 inline constexpr const char* kTagGC = "gc";
 inline constexpr const char* kTagMM = "mm";
 inline constexpr const char* kTagTLS = "tls";
+inline constexpr const char* kTagPause = "pause";
 
 } // namespace kotlin
 

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

@@ -7,14 +7,11 @@
 #include "ThreadSuspension.hpp"
 
 #include <condition_variable>
-#include <thread>
-#include <mutex>
 
 #include "CallsChecker.hpp"
 #include "Logging.hpp"
 #include "Porting.h"
 #include "SafePoint.hpp"
-#include "StackTrace.hpp"
 
 using namespace kotlin;
 
@@ -26,12 +23,29 @@ namespace {
 
 } // namespace
 
-bool kotlin::mm::isSuspendedOrNative(kotlin::mm::ThreadData& thread) noexcept {
+std::atomic<bool> kotlin::mm::internal::gSuspensionRequested = false;
-    auto& suspensionData = thread.suspensionData();
+
-    return suspensionData.suspended() || suspensionData.state() == kotlin::ThreadState::kNative;
+ALWAYS_INLINE mm::ThreadSuspensionData::MutatorPauseHandle::MutatorPauseHandle(const char* reason, mm::ThreadData& threadData) noexcept
+    : reason_(reason), threadData_(threadData), pauseStartTimeMicros_(konan::getTimeMicros())
+{
+    auto prevState = threadData_.suspensionData().setStateNoSafePoint(ThreadState::kNative);
+    // no special reason, fill free to implement pause from native if needed
+    RuntimeAssert(prevState == ThreadState::kRunnable, "Expected runnable state");
+    RuntimeLogDebug({kTagPause}, "Suspending mutation (%s)", reason_);
 }
 
-std::atomic<bool> kotlin::mm::internal::gSuspensionRequested = false;
+ALWAYS_INLINE mm::ThreadSuspensionData::MutatorPauseHandle::~MutatorPauseHandle() noexcept {
+    if (!resumed) resume();
+}
+
+ALWAYS_INLINE void mm::ThreadSuspensionData::MutatorPauseHandle::resume() noexcept {
+    RuntimeAssert(!resumed, "Must not be resumed yet");
+    auto prevState = threadData_.suspensionData().setStateNoSafePoint(ThreadState::kRunnable);
+    RuntimeAssert(prevState == ThreadState::kNative, "Expected native state");
+    auto pauseTimeMicros = konan::getTimeMicros() - pauseStartTimeMicros_;
+    RuntimeLogInfo({kTagPause}, "Resuming mutation after %" PRIu64 " microseconds of suspension (%s)", pauseTimeMicros, reason_);
+    resumed = true;
+}
 
 kotlin::ThreadState kotlin::mm::ThreadSuspensionData::setState(kotlin::ThreadState newState) noexcept {
     ThreadState oldState = state_.exchange(newState);
@@ -51,19 +65,21 @@ kotlin::ThreadState kotlin::mm::ThreadSuspensionData::setState(kotlin::ThreadSta
 
 NO_EXTERNAL_CALLS_CHECK void kotlin::mm::ThreadSuspensionData::suspendIfRequested() noexcept {
     if (IsThreadSuspensionRequested()) {
-        auto suspendStartMs = konan::getTimeMicros();
+        auto pauseHandle = pauseMutationInScope("stop the world");
+
         threadData_.gc().OnSuspendForGC();
         std::unique_lock lock(gSuspensionMutex);
-        auto threadId = konan::currentThreadId();
-        RuntimeLogDebug({kTagGC, kTagMM}, "Suspending thread %d", threadId);
-        AutoReset scopedAssignSuspended(&suspended_, true);
         gSuspensionCondVar.wait(lock, []() { return !IsThreadSuspensionRequested(); });
-        auto suspendEndMs = konan::getTimeMicros();
+
-        RuntimeLogDebug({kTagGC, kTagMM}, "Resuming thread %d after %" PRIu64 " microseconds of suspension",
+        // Must return to running state under the lock.
-                        threadId, suspendEndMs - suspendStartMs);
+        pauseHandle.resume();
     }
 }
 
+ALWAYS_INLINE mm::ThreadSuspensionData::MutatorPauseHandle mm::ThreadSuspensionData::pauseMutationInScope(const char* reason) noexcept {
+    return MutatorPauseHandle(reason, threadData_);
+}
+
 bool kotlin::mm::RequestThreadsSuspension() noexcept {
     CallsCheckerIgnoreGuard guard;
 

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

@@ -28,8 +28,21 @@ inline bool IsThreadSuspensionRequested() noexcept {
 }
 
 class ThreadSuspensionData : private Pinned {
+private:
+    class MutatorPauseHandle : private Pinned {
+    public:
+        explicit MutatorPauseHandle(const char* reason, ThreadData& threadData) noexcept;
+        ~MutatorPauseHandle() noexcept;
+        void resume() noexcept;
+    private:
+        const char* reason_;
+        ThreadData& threadData_;
+        uint64_t pauseStartTimeMicros_;
+        bool resumed = false;
+    };
+
 public:
-    explicit ThreadSuspensionData(ThreadState initialState, mm::ThreadData& threadData) noexcept : state_(initialState), threadData_(threadData), suspended_(false) {}
+    explicit ThreadSuspensionData(ThreadState initialState, mm::ThreadData& threadData) noexcept : state_(initialState), threadData_(threadData) {}
 
     ~ThreadSuspensionData() = default;
 
@@ -38,22 +51,29 @@ public:
     ThreadState setState(ThreadState newState) noexcept;
     ThreadState setStateNoSafePoint(ThreadState newState) noexcept { return state_.exchange(newState, std::memory_order_acq_rel); }
 
-    bool suspended() noexcept { return suspended_; }
+    bool suspendedOrNative() noexcept { return state() == kotlin::ThreadState::kNative; }
-    bool suspendedOrNative() noexcept { return suspended() || state() == kotlin::ThreadState::kNative; }
 
     void suspendIfRequested() noexcept;
 
+    /**
+     * Signals that the thread would not mutate a heap during a relatively long time.
+     * For example while waiting for or participating in the GC.
+     * Effectively sets the thread's state to `kNative`.
+     *
+     * The pause is considered completed upon destruction of a returned pause-handle object.
+     *
+     * NOTE: the safe point actions will NOT be automatically executed after the pause.
+     */
+    [[nodiscard]] MutatorPauseHandle pauseMutationInScope(const char* reason) noexcept;
+
 private:
     std::atomic<ThreadState> state_;
     mm::ThreadData& threadData_;
-    std::atomic<bool> suspended_;
 };
 
 bool RequestThreadsSuspension() noexcept;
 void WaitForThreadsSuspension() noexcept;
 
-bool isSuspendedOrNative(kotlin::mm::ThreadData& thread) noexcept;
-
 /**
  * Suspends all threads registered in ThreadRegistry except threads that are in the Native state.
  * Blocks until all such threads are suspended. Threads that are in the Native state on the moment

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

@@ -52,7 +52,7 @@ std_support::vector<T> collectFromThreadData(F extractFunction) {
 
 std_support::vector<bool> collectSuspended() {
     return collectFromThreadData<bool>(
-            [](mm::ThreadData* threadData) { return threadData->suspensionData().suspended(); });
+            [](mm::ThreadData* threadData) { return threadData->suspensionData().suspendedOrNative(); });
 }
 
 void reportProgress(size_t currentIteration, size_t totalIterations) {
@@ -116,9 +116,9 @@ TEST_F(ThreadSuspensionTest, SimpleStartStop) {
             while(!shouldStop) {
                 waitUntilCanStart(i);
 
-                EXPECT_FALSE(suspensionData.suspended());
+                EXPECT_FALSE(suspensionData.suspendedOrNative());
                 suspensionData.suspendIfRequested();
-                EXPECT_FALSE(suspensionData.suspended());
+                EXPECT_FALSE(suspensionData.suspendedOrNative());
            }
         });
     }
@@ -201,10 +201,10 @@ TEST_F(ThreadSuspensionTest, ConcurrentSuspend) {
                 successCount++;
                 auto allThreadData = collectThreadData();
                 auto isCurrentOrSuspended = [currentThreadData](mm::ThreadData* data) {
-                    return data == currentThreadData || data->suspensionData().suspended();
+                    return data == currentThreadData || data->suspensionData().suspendedOrNative();
                 };
                 EXPECT_THAT(allThreadData, testing::Each(testing::Truly(isCurrentOrSuspended)));
-                EXPECT_FALSE(currentThreadData->suspensionData().suspended());
+                EXPECT_FALSE(currentThreadData->suspensionData().suspendedOrNative());
                 mm::ResumeThreads();
             } else {
                 EXPECT_TRUE(mm::IsThreadSuspensionRequested());