[K/N] Parallelize marking in CMS GC

This change parallelizes marking by making each non-native Kotlin thread
mark its own view of the heap. To make this safe, we make flipping the
marking bit atomic which ensures that threads do not try to trace the
same objects.

Co-authored-by: Johan Bay <jobay@google.com>

Merge-request: KOTLIN-MR-423
Merged-by: Alexander Shabalin <alexander.shabalin@jetbrains.com>

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/BinaryOptions.kt

@@ -31,6 +31,8 @@ object BinaryOptions : BinaryOptionRegistry() {
 
     val gcSchedulerType by option<GCSchedulerType>()
 
+    val gcMarkSingleThreaded by booleanOption()
+
     val linkRuntime by option<RuntimeLinkageStrategyBinaryOption>()
 
     val bundleId by stringOption()

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanConfig.kt

@@ -170,6 +170,9 @@ class KonanConfig(val project: Project, val configuration: CompilerConfiguration
         configuration.get(BinaryOptions.gcSchedulerType) ?: defaultGCSchedulerType
     }
 
+    val gcMarkSingleThreaded: Boolean
+        get() = configuration.get(BinaryOptions.gcMarkSingleThreaded) == true
+
     val needVerifyIr: Boolean
         get() = configuration.get(KonanConfigKeys.VERIFY_IR) == true
 

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IrToBitcode.kt

@@ -2773,6 +2773,7 @@ internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrE
             return
 
         overrideRuntimeGlobal("Kotlin_destroyRuntimeMode", Int32(context.config.destroyRuntimeMode.value))
+        overrideRuntimeGlobal("Kotlin_gcMarkSingleThreaded", Int32(if (context.config.gcMarkSingleThreaded) 1 else 0))
         overrideRuntimeGlobal("Kotlin_workerExceptionHandling", Int32(context.config.workerExceptionHandling.value))
         overrideRuntimeGlobal("Kotlin_suspendFunctionsFromAnyThreadFromObjC", Int32(if (context.config.suspendFunctionsFromAnyThreadFromObjC) 1 else 0))
         val getSourceInfoFunctionName = when (context.config.sourceInfoType) {

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

@@ -9,6 +9,7 @@
 
 #include "CompilerConstants.hpp"
 #include "GlobalData.hpp"
+#include "GCImpl.hpp"
 #include "Logging.hpp"
 #include "MarkAndSweepUtils.hpp"
 #include "Memory.h"
@@ -23,6 +24,9 @@
 using namespace kotlin;
 
 namespace {
+    [[clang::no_destroy]] std::mutex markingMutex;
+    [[clang::no_destroy]] std::condition_variable markingCondVar;
+    [[clang::no_destroy]] std::atomic<bool> markingRequested = false;
 
 struct MarkTraits {
     using MarkQueue = gc::ConcurrentMarkAndSweep::MarkQueue;
@@ -44,8 +48,7 @@ struct MarkTraits {
 
     static void enqueue(MarkQueue& queue, ObjHeader* object) noexcept {
         auto& objectData = mm::ObjectFactory<gc::ConcurrentMarkAndSweep>::NodeRef::From(object).ObjectData();
-        if (objectData.color() == gc::ConcurrentMarkAndSweep::ObjectData::Color::kBlack) return;
-        objectData.setColor(gc::ConcurrentMarkAndSweep::ObjectData::Color::kBlack);
+        if (!objectData.atomicSetToBlack()) return;
         queue.push_front(objectData);
     }
 };
@@ -92,6 +95,22 @@ void gc::ConcurrentMarkAndSweep::ThreadData::OnOOM(size_t size) noexcept {
     ScheduleAndWaitFullGC();
 }
 
+NO_EXTERNAL_CALLS_CHECK void gc::ConcurrentMarkAndSweep::ThreadData::PublishAndMark() noexcept {
+    std::unique_lock lock(markingMutex);
+    if (!markingRequested.load())
+        return;
+    AutoReset scopedAssignMarking(&marking_, true);
+    threadData_.Publish();
+    markingCondVar.wait(lock, []() { return !markingRequested.load(); });
+    // // Unlock while marking to allow mutliple threads to mark in parallel.
+    lock.unlock();
+    RuntimeLogDebug({kTagGC}, "Parallel marking in thread %d", konan::currentThreadId());
+    MarkQueue markQueue;
+    gc::collectRootSetForThread<MarkTraits>(markQueue, threadData_);
+    MarkStats stats = gc::Mark<MarkTraits>(markQueue);
+    gc_.MergeMarkStats(stats);
+}
+
 gc::ConcurrentMarkAndSweep::ConcurrentMarkAndSweep(
         mm::ObjectFactory<ConcurrentMarkAndSweep>& objectFactory, GCScheduler& gcScheduler) noexcept :
     objectFactory_(objectFactory),
@@ -113,6 +132,10 @@ gc::ConcurrentMarkAndSweep::ConcurrentMarkAndSweep(
             }
         }
     });
+    markingBehavior_ = kotlin::compiler::gcMarkSingleThreaded() ? MarkingBehavior::kDoNotMark : MarkingBehavior::kMarkOwnStack;
+    mm::SetOnSuspendCallback([](mm::ThreadData& thread) {
+        thread.gc().impl().gc().PublishAndMark();
+    });
     RuntimeLogDebug({kTagGC}, "Concurrent Mark & Sweep GC initialized");
 }
 
@@ -135,18 +158,23 @@ bool gc::ConcurrentMarkAndSweep::FinalizersThreadIsRunning() noexcept {
     return finalizerProcessor_->IsRunning();
 }
 
+void gc::ConcurrentMarkAndSweep::SetMarkingBehaviorForTests(MarkingBehavior markingBehavior) noexcept {
+    markingBehavior_ = markingBehavior;
+}
+
 bool gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     RuntimeLogDebug({kTagGC}, "Attempt to suspend threads by thread %d", konan::currentThreadId());
+    SetMarkingRequested();
     auto timeStartUs = konan::getTimeMicros();
     bool didSuspend = mm::RequestThreadsSuspension();
     RuntimeAssert(didSuspend, "Only GC thread can request suspension");
     RuntimeLogDebug({kTagGC}, "Requested thread suspension by thread %d", konan::currentThreadId());
 
     RuntimeAssert(!kotlin::mm::IsCurrentThreadRegistered(), "Concurrent GC must run on unregistered thread");
-
-    mm::WaitForThreadsSuspension();
+    WaitForThreadsReadyToMark();
     auto timeSuspendUs = konan::getTimeMicros();
     RuntimeLogDebug({kTagGC}, "Suspended all threads in %" PRIu64 " microseconds", timeSuspendUs - timeStartUs);
+    lastGCMarkStats_ = MarkStats();
 
     auto& scheduler = gcScheduler_;
     scheduler.gcData().OnPerformFullGC();
@@ -154,28 +182,32 @@ bool gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     state_.start(epoch);
     RuntimeLogInfo(
             {kTagGC}, "Started GC epoch %" PRId64 ". Time since last GC %" PRIu64 " microseconds", epoch, timeStartUs - lastGCTimestampUs_);
-    gc::collectRootSet<MarkTraits>(markQueue_);
-    auto timeRootSetUs = konan::getTimeMicros();
-    // Can be unsafe, because we've stopped the world.
 
+    CollectRootSetAndStartMarking();
+
+    // Can be unsafe, because we've stopped the world.
     auto objectsCountBefore = objectFactory_.GetSizeUnsafe();
-    RuntimeLogInfo(
-            {kTagGC}, "Collected root set of size %zu in %" PRIu64 " microseconds", markQueue_.size(),
-            timeRootSetUs - timeSuspendUs);
+
     auto markStats = gc::Mark<MarkTraits>(markQueue_);
-    auto timeMarkUs = konan::getTimeMicros();
-    RuntimeLogDebug({kTagGC}, "Marked %zu objects in %" PRIu64 " microseconds", markStats.aliveHeapSet, timeMarkUs - timeRootSetUs);
-    scheduler.gcData().UpdateAliveSetBytes(markStats.aliveHeapSetBytes);
+    MergeMarkStats(markStats);
+
+    RuntimeLogDebug({kTagGC}, "Waiting for marking in threads");
+    mm::WaitForThreadsSuspension();
+    auto timeMarkingUs = konan::getTimeMicros();
+    RuntimeLogInfo({kTagGC}, "Collected root set of size %zu and marked %zu objects in all threads in %" PRIu64 " microseconds", lastGCMarkStats_.rootSetSize, lastGCMarkStats_.aliveHeapSet, timeMarkingUs - timeSuspendUs);
+
+    scheduler.gcData().UpdateAliveSetBytes(lastGCMarkStats_.aliveHeapSetBytes);
+
     gc::SweepExtraObjects<SweepTraits>(mm::GlobalData::Instance().extraObjectDataFactory());
     auto timeSweepExtraObjectsUs = konan::getTimeMicros();
-    RuntimeLogDebug({kTagGC}, "Sweeped extra objects in %" PRIu64 " microseconds", timeSweepExtraObjectsUs - timeMarkUs);
+    RuntimeLogDebug({kTagGC}, "Sweeped extra objects in %" PRIu64 " microseconds", timeSweepExtraObjectsUs - timeMarkingUs);
 
     auto objectFactoryIterable = objectFactory_.LockForIter();
 
     mm::ResumeThreads();
     auto timeResumeUs = konan::getTimeMicros();
 
-    RuntimeLogDebug({kTagGC},
+    RuntimeLogInfo({kTagGC},
                     "Resumed threads in %" PRIu64 " microseconds. Total pause for most threads is %"  PRIu64" microseconds",
                     timeResumeUs - timeSweepExtraObjectsUs, timeResumeUs - timeStartUs);
 
@@ -202,3 +234,51 @@ bool gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
     return true;
 }
 
+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() noexcept {
+    markingRequested = markingBehavior_ == MarkingBehavior::kMarkOwnStack;
+}
+
+void gc::ConcurrentMarkAndSweep::WaitForThreadsReadyToMark() noexcept {
+    while(!allThreads([](kotlin::mm::ThreadData& thread) { return isSuspendedOrNative(thread) || thread.gc().impl().gc().marking_.load(); })) {
+        yield();
+    }
+}
+
+NO_EXTERNAL_CALLS_CHECK void gc::ConcurrentMarkAndSweep::CollectRootSetAndStartMarking() noexcept {
+        std::unique_lock lock(markingMutex);
+        markingRequested = false;
+        gc::collectRootSet<MarkTraits>(markQueue_, [](mm::ThreadData& thread) { return !thread.gc().impl().gc().marking_.load(); });
+        RuntimeLogDebug({kTagGC}, "Requesting marking in threads");
+        markingCondVar.notify_all();
+}
+
+void gc::ConcurrentMarkAndSweep::MergeMarkStats(gc::MarkStats stats) noexcept {
+    std::unique_lock lock(markingMutex);
+    lastGCMarkStats_.merge(stats);
+}

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

@@ -5,11 +5,13 @@
 
 #pragma once
 
+#include <atomic>
 #include <cstddef>
 
 #include "Allocator.hpp"
 #include "GCScheduler.hpp"
 #include "IntrusiveList.hpp"
+#include "MarkAndSweepUtils.hpp"
 #include "ObjectFactory.hpp"
 #include "ScopedThread.hpp"
 #include "Types.h"
@@ -27,8 +29,8 @@ namespace gc {
 
 class FinalizerProcessor;
 
-// Stop-the-world mark + concurrent sweep. The GC runs in a separate thread, finalizers run in another thread of their own.
-// TODO: Also make mark concurrent.
+// Stop-the-world parallel mark + concurrent sweep. The GC runs in a separate thread, finalizers run in another thread of their own.
+// TODO: Also make marking run concurrently with Kotlin threads.
 class ConcurrentMarkAndSweep : private Pinned {
 public:
     class ObjectData {
@@ -41,19 +43,28 @@ public:
             kBlack, // Objects encountered during mark phase.
         };
 
-        Color color() const noexcept { return static_cast<Color>(getPointerBits(next_, colorMask)); }
-        void setColor(Color color) noexcept { next_ = setPointerBits(clearPointerBits(next_, colorMask), static_cast<unsigned>(color)); }
+        Color color() const noexcept { return static_cast<Color>(getPointerBits(next_.load(std::memory_order_relaxed), colorMask)); }
+        void setColor(Color color) noexcept { next_.store(setPointerBits(clearPointerBits(next_.load(std::memory_order_relaxed), colorMask), static_cast<unsigned>(color)), std::memory_order_relaxed); }
+        bool atomicSetToBlack() noexcept {
+            ObjectData* before = next_.load(std::memory_order_relaxed);
+            if (getPointerBits(before, colorMask) != static_cast<unsigned>(Color::kWhite))
+                return false;
+            ObjectData* black = setPointerBits(before, static_cast<unsigned>(Color::kBlack));
+            bool success = next_.compare_exchange_strong(before, black, std::memory_order_relaxed);
+            RuntimeAssert(success || hasPointerBits(before, colorMask), "next_ must have been marked black");
+            return success;
+        }
 
-        ObjectData* next() const noexcept { return clearPointerBits(next_, colorMask); }
+        ObjectData* next() const noexcept { return clearPointerBits(next_.load(std::memory_order_relaxed), colorMask); }
         void setNext(ObjectData* next) noexcept {
             RuntimeAssert(!hasPointerBits(next, colorMask), "next must be untagged: %p", next);
-            auto bits = getPointerBits(next_, colorMask);
-            next_ = setPointerBits(next, bits);
+            auto bits = getPointerBits(next_.load(std::memory_order_relaxed), colorMask);
+            next_.store(setPointerBits(next, bits), std::memory_order_relaxed);
         }
 
     private:
         // Color is encoded in low bits.
-        ObjectData* next_ = nullptr;
+        std::atomic<ObjectData*> next_ = nullptr;
     };
 
     struct MarkQueueTraits {
@@ -62,6 +73,8 @@ public:
         static void setNext(ObjectData& value, ObjectData* next) noexcept { value.setNext(next); }
     };
 
+    enum MarkingBehavior { kMarkOwnStack, kDoNotMark };
+
     using MarkQueue = intrusive_forward_list<ObjectData, MarkQueueTraits>;
 
     class ThreadData : private Pinned {
@@ -70,7 +83,7 @@ public:
         using Allocator = AllocatorWithGC<Allocator, ThreadData>;
 
         explicit ThreadData(ConcurrentMarkAndSweep& gc, mm::ThreadData& threadData, GCSchedulerThreadData& gcScheduler) noexcept :
-            gc_(gc), gcScheduler_(gcScheduler) {}
+            gc_(gc), threadData_(threadData), gcScheduler_(gcScheduler) {}
         ~ThreadData() = default;
 
         void SafePointAllocation(size_t size) noexcept;
@@ -80,11 +93,16 @@ public:
 
         void OnOOM(size_t size) noexcept;
 
+        void PublishAndMark() noexcept;
+
         Allocator CreateAllocator() noexcept { return Allocator(gc::Allocator(), *this); }
 
     private:
+        friend ConcurrentMarkAndSweep;
         ConcurrentMarkAndSweep& gc_;
+        mm::ThreadData& threadData_;
         GCSchedulerThreadData& gcScheduler_;
+        std::atomic<bool> marking_;
     };
 
     using Allocator = ThreadData::Allocator;
@@ -95,10 +113,15 @@ public:
     void StartFinalizerThreadIfNeeded() noexcept;
     void StopFinalizerThreadIfRunning() noexcept;
     bool FinalizersThreadIsRunning() noexcept;
+    void SetMarkingBehaviorForTests(MarkingBehavior markingBehavior) noexcept;
+    void SetMarkingRequested() noexcept;
+    void WaitForThreadsReadyToMark() noexcept;
+    void CollectRootSetAndStartMarking() noexcept;
 
 private:
     // Returns `true` if GC has happened, and `false` if not (because someone else has suspended the threads).
     bool PerformFullGC(int64_t epoch) noexcept;
+    void MergeMarkStats(MarkStats stats) noexcept;
 
     mm::ObjectFactory<ConcurrentMarkAndSweep>& objectFactory_;
     GCScheduler& gcScheduler_;
@@ -109,6 +132,8 @@ private:
     std_support::unique_ptr<FinalizerProcessor> finalizerProcessor_;
 
     MarkQueue markQueue_;
+    MarkStats lastGCMarkStats_;
+    MarkingBehavior markingBehavior_;
 };
 
 } // namespace gc

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

@@ -216,9 +216,13 @@ WeakCounter& InstallWeakCounter(mm::ThreadData& threadData, ObjHeader* objHeader
     return weakCounter;
 }
 
-class ConcurrentMarkAndSweepTest : public testing::Test {
+class ConcurrentMarkAndSweepTest : public testing::TestWithParam<gc::ConcurrentMarkAndSweep::MarkingBehavior> {
 public:
 
+    ConcurrentMarkAndSweepTest() {
+        mm::GlobalData::Instance().gc().impl().gc().SetMarkingBehaviorForTests(GetParam());
+    }
+
     ~ConcurrentMarkAndSweepTest() {
         mm::GlobalsRegistry::Instance().ClearForTests();
         mm::GlobalData::Instance().extraObjectDataFactory().ClearForTests();
@@ -233,7 +237,7 @@ private:
 
 } // namespace
 
-TEST_F(ConcurrentMarkAndSweepTest, RootSet) {
+TEST_P(ConcurrentMarkAndSweepTest, RootSet) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalObjectHolder global1{threadData};
         GlobalObjectArrayHolder global2{threadData};
@@ -268,7 +272,7 @@ TEST_F(ConcurrentMarkAndSweepTest, RootSet) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, InterconnectedRootSet) {
+TEST_P(ConcurrentMarkAndSweepTest, InterconnectedRootSet) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalObjectHolder global1{threadData};
         GlobalObjectArrayHolder global2{threadData};
@@ -314,7 +318,7 @@ TEST_F(ConcurrentMarkAndSweepTest, InterconnectedRootSet) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, FreeObjects) {
+TEST_P(ConcurrentMarkAndSweepTest, FreeObjects) {
     RunInNewThread([](mm::ThreadData& threadData) {
         auto& object1 = AllocateObject(threadData);
         auto& object2 = AllocateObject(threadData);
@@ -329,7 +333,7 @@ TEST_F(ConcurrentMarkAndSweepTest, FreeObjects) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, FreeObjectsWithFinalizers) {
+TEST_P(ConcurrentMarkAndSweepTest, FreeObjectsWithFinalizers) {
     RunInNewThread([this](mm::ThreadData& threadData) {
         auto& object1 = AllocateObjectWithFinalizer(threadData);
         auto& object2 = AllocateObjectWithFinalizer(threadData);
@@ -346,7 +350,7 @@ TEST_F(ConcurrentMarkAndSweepTest, FreeObjectsWithFinalizers) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeak) {
+TEST_P(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeak) {
     RunInNewThread([](mm::ThreadData& threadData) {
         auto& object1 = AllocateObject(threadData);
         auto& weak1 = ([&threadData, &object1]() -> WeakCounter& {
@@ -365,7 +369,7 @@ TEST_F(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeak) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, FreeObjectWithHoldedWeak) {
+TEST_P(ConcurrentMarkAndSweepTest, FreeObjectWithHoldedWeak) {
     RunInNewThread([](mm::ThreadData& threadData) {
         auto& object1 = AllocateObject(threadData);
         StackObjectHolder stack{threadData};
@@ -384,7 +388,7 @@ TEST_F(ConcurrentMarkAndSweepTest, FreeObjectWithHoldedWeak) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, ObjectReferencedFromRootSet) {
+TEST_P(ConcurrentMarkAndSweepTest, ObjectReferencedFromRootSet) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalObjectHolder global{threadData};
         StackObjectHolder stack{threadData};
@@ -424,7 +428,7 @@ TEST_F(ConcurrentMarkAndSweepTest, ObjectReferencedFromRootSet) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, ObjectsWithCycles) {
+TEST_P(ConcurrentMarkAndSweepTest, ObjectsWithCycles) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalObjectHolder global{threadData};
         StackObjectHolder stack{threadData};
@@ -473,7 +477,7 @@ TEST_F(ConcurrentMarkAndSweepTest, ObjectsWithCycles) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, ObjectsWithCyclesAndFinalizers) {
+TEST_P(ConcurrentMarkAndSweepTest, ObjectsWithCyclesAndFinalizers) {
     RunInNewThread([this](mm::ThreadData& threadData) {
         GlobalObjectHolder global{threadData};
         StackObjectHolder stack{threadData};
@@ -524,7 +528,7 @@ TEST_F(ConcurrentMarkAndSweepTest, ObjectsWithCyclesAndFinalizers) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, ObjectsWithCyclesIntoRootSet) {
+TEST_P(ConcurrentMarkAndSweepTest, ObjectsWithCyclesIntoRootSet) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalObjectHolder global{threadData};
         StackObjectHolder stack{threadData};
@@ -552,7 +556,7 @@ TEST_F(ConcurrentMarkAndSweepTest, ObjectsWithCyclesIntoRootSet) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, RunGCTwice) {
+TEST_P(ConcurrentMarkAndSweepTest, RunGCTwice) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalObjectHolder global{threadData};
         StackObjectHolder stack{threadData};
@@ -602,7 +606,7 @@ TEST_F(ConcurrentMarkAndSweepTest, RunGCTwice) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, PermanentObjects) {
+TEST_P(ConcurrentMarkAndSweepTest, PermanentObjects) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalPermanentObjectHolder global1{threadData};
         GlobalObjectHolder global2{threadData};
@@ -624,7 +628,7 @@ TEST_F(ConcurrentMarkAndSweepTest, PermanentObjects) {
     });
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, SameObjectInRootSet) {
+TEST_P(ConcurrentMarkAndSweepTest, SameObjectInRootSet) {
     RunInNewThread([](mm::ThreadData& threadData) {
         GlobalObjectHolder global{threadData};
         StackObjectHolder stack(*global);
@@ -710,7 +714,7 @@ private:
 
 } // namespace
 
-TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsCollect) {
+TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsCollect) {
     std_support::vector<Mutator> mutators(kDefaultThreadCount);
     std_support::vector<ObjHeader*> globals(kDefaultThreadCount);
     std_support::vector<ObjHeader*> locals(kDefaultThreadCount);
@@ -766,7 +770,7 @@ TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsCollect) {
     }
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsAllCollect) {
+TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsAllCollect) {
     std_support::vector<Mutator> mutators(kDefaultThreadCount);
     std_support::vector<ObjHeader*> globals(kDefaultThreadCount);
     std_support::vector<ObjHeader*> locals(kDefaultThreadCount);
@@ -828,7 +832,7 @@ TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsAllCollect) {
     }
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsAddToRootSetAfterCollectionRequested) {
+TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsAddToRootSetAfterCollectionRequested) {
     std_support::vector<Mutator> mutators(kDefaultThreadCount);
     std_support::vector<ObjHeader*> globals(kDefaultThreadCount);
     std_support::vector<ObjHeader*> locals(kDefaultThreadCount);
@@ -900,7 +904,7 @@ TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsAddToRootSetAfterCollectionRe
     }
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, CrossThreadReference) {
+TEST_P(ConcurrentMarkAndSweepTest, CrossThreadReference) {
     std_support::vector<Mutator> mutators(kDefaultThreadCount);
     std_support::vector<ObjHeader*> globals(kDefaultThreadCount);
     std_support::vector<ObjHeader*> locals(kDefaultThreadCount);
@@ -966,7 +970,7 @@ TEST_F(ConcurrentMarkAndSweepTest, CrossThreadReference) {
     }
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsWeaks) {
+TEST_P(ConcurrentMarkAndSweepTest, MultipleMutatorsWeaks) {
     std_support::vector<Mutator> mutators(kDefaultThreadCount);
     ObjHeader* globalRoot = nullptr;
     WeakCounter* weak = nullptr;
@@ -1018,7 +1022,7 @@ TEST_F(ConcurrentMarkAndSweepTest, MultipleMutatorsWeaks) {
     }
 }
 
-TEST_F(ConcurrentMarkAndSweepTest, NewThreadsWhileRequestingCollection) {
+TEST_P(ConcurrentMarkAndSweepTest, NewThreadsWhileRequestingCollection) {
     std_support::vector<Mutator> mutators(kDefaultThreadCount);
     std_support::vector<ObjHeader*> globals(2 * kDefaultThreadCount);
     std_support::vector<ObjHeader*> locals(2 * kDefaultThreadCount);
@@ -1100,7 +1104,7 @@ TEST_F(ConcurrentMarkAndSweepTest, NewThreadsWhileRequestingCollection) {
 }
 
 
-TEST_F(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeakReversedOrder) {
+TEST_P(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeakReversedOrder) {
     std_support::vector<Mutator> mutators(2);
     std::atomic<test_support::Object<Payload>*> object1 = nullptr;
     std::atomic<WeakCounter*> weak = nullptr;
@@ -1139,3 +1143,61 @@ TEST_F(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeakReversedOrder) {
     f0.wait();
     f1.wait();
 }
+
+TEST_P(ConcurrentMarkAndSweepTest, MutatorsCanMarkOwnLocals) {
+    std_support::vector<Mutator> mutators(kDefaultThreadCount);
+    std_support::vector<ObjHeader*> globals(kDefaultThreadCount);
+    std_support::vector<ObjHeader*> locals(kDefaultThreadCount);
+    std_support::vector<ObjHeader*> reachablesLocals(kDefaultThreadCount);
+    std_support::vector<ObjHeader*> reachablesGlobals(kDefaultThreadCount);
+
+    auto expandRootSet = [&globals, &locals, &reachablesLocals, &reachablesGlobals](mm::ThreadData& threadData, Mutator& mutator, int i) {
+        auto& global = mutator.AddGlobalRoot();
+        auto& local = mutator.AddStackRoot();
+        auto& reachableLocal = AllocateObject(threadData);
+        auto& reachableGlobal = AllocateObject(threadData);
+        local->field1 = reachableLocal.header();
+        global->field1 = reachableGlobal.header();
+        globals[i] = global.header();
+        locals[i] = local.header();
+        reachablesLocals[i] = reachableLocal.header();
+        reachablesGlobals[i] = reachableGlobal.header();
+    };
+
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        mutators[i]
+                .Execute([i, expandRootSet](mm::ThreadData& threadData, Mutator& mutator) { expandRootSet(threadData, mutator, i); })
+                .wait();
+    }
+
+    std_support::vector<std::future<void>> gcFutures(kDefaultThreadCount);
+
+    mm::GlobalData::Instance().gc().impl().gc().SetMarkingRequested();
+
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        gcFutures[i] = mutators[i]
+            .Execute([](mm::ThreadData& threadData, Mutator& mutator) { threadData.gc().impl().gc().PublishAndMark(); });
+    }
+
+    if (GetParam() == gc::ConcurrentMarkAndSweep::kMarkOwnStack) {
+        mm::GlobalData::Instance().gc().impl().gc().WaitForThreadsReadyToMark();
+        mm::GlobalData::Instance().gc().impl().gc().CollectRootSetAndStartMarking();
+    }
+
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        gcFutures[i].wait();
+        // Verify that threads marked their own locals (but not their globals) according to the configured marking behavior:
+        ASSERT_THAT(GetColor(reachablesLocals[i]), GetParam() == kotlin::gc::ConcurrentMarkAndSweep::kMarkOwnStack ? Color::kBlack : Color::kWhite);
+        ASSERT_THAT(GetColor(reachablesGlobals[i]), Color::kWhite);
+    }
+
+    for (auto& future : gcFutures) {
+        future.wait();
+    }
+}
+
+INSTANTIATE_TEST_SUITE_P(,
+                         ConcurrentMarkAndSweepTest,
+                         testing::Values(gc::ConcurrentMarkAndSweep::MarkingBehavior::kDoNotMark, gc::ConcurrentMarkAndSweep::MarkingBehavior::kMarkOwnStack),
+                         [] (const testing::TestParamInfo<gc::ConcurrentMarkAndSweep::MarkingBehavior>& behavior) { return (behavior.param == gc::ConcurrentMarkAndSweep::MarkingBehavior::kDoNotMark) ? "SingleThreadedMarking" : "MutatorsMarkOwnStack"; });
+

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

@@ -27,11 +27,21 @@ struct MarkStats {
     size_t aliveHeapSet = 0;
     // How many objects are alive in bytes. Note: this does not include overhead of malloc/mimalloc itself.
     size_t aliveHeapSetBytes = 0;
+    // How many roots are were marked.
+    size_t rootSetSize = 0;
+
+    void merge(MarkStats other) {
+        aliveHeapSet += other.aliveHeapSet;
+        aliveHeapSetBytes += other.aliveHeapSetBytes;
+        rootSetSize += other.rootSetSize;
+    }
 };
 
 template <typename Traits>
 MarkStats Mark(typename Traits::MarkQueue& markQueue) noexcept {
     MarkStats stats;
+    stats.rootSetSize = markQueue.size();
+    auto timeStart = konan::getTimeMicros();
     while (!Traits::isEmpty(markQueue)) {
         ObjHeader* top = Traits::dequeue(markQueue);
 
@@ -57,6 +67,8 @@ MarkStats Mark(typename Traits::MarkQueue& markQueue) noexcept {
             }
         }
     }
+    auto timeEnd = konan::getTimeMicros();
+    RuntimeLogDebug({kTagGC}, "Marked %zu objects in %" PRIu64 " microseconds in thread %d", stats.aliveHeapSet, timeEnd - timeStart, konan::currentThreadId());
     return stats;
 }
 
@@ -108,42 +120,40 @@ typename Traits::ObjectFactory::FinalizerQueue Sweep(typename Traits::ObjectFact
     return Sweep<Traits>(iter);
 }
 
-// TODO: This needs some tests now.
 template <typename Traits>
-void collectRootSet(typename Traits::MarkQueue& markQueue) noexcept {
-    Traits::clear(markQueue);
-    for (auto& thread : mm::GlobalData::Instance().threadRegistry().LockForIter()) {
-        // TODO: Maybe it's more efficient to do by the suspending thread?
-        thread.Publish();
-        thread.gc().OnStoppedForGC();
-        size_t stack = 0;
-        size_t tls = 0;
-        for (auto value : mm::ThreadRootSet(thread)) {
-            auto* object = value.object;
-            if (!isNullOrMarker(object)) {
-                if (object->heap()) {
-                    Traits::enqueue(markQueue, object);
-                } else {
-                    traverseReferredObjects(object, [&](ObjHeader* field) noexcept {
-                        // Each permanent and stack object has own entry in the root set.
-                        if (field->heap() && !isNullOrMarker(field)) {
-                            Traits::enqueue(markQueue, field);
-                        }
-                    });
-                    RuntimeAssert(!object->has_meta_object(), "Non-heap object %p may not have an extra object data", object);
-                }
-                switch (value.source) {
-                    case mm::ThreadRootSet::Source::kStack:
-                        ++stack;
-                        break;
-                    case mm::ThreadRootSet::Source::kTLS:
-                        ++tls;
-                        break;
-                }
+void collectRootSetForThread(typename Traits::MarkQueue& markQueue, mm::ThreadData& thread) {
+    thread.gc().OnStoppedForGC();
+    size_t stack = 0;
+    size_t tls = 0;
+    for (auto value : mm::ThreadRootSet(thread)) {
+        auto* object = value.object;
+        if (!isNullOrMarker(object)) {
+            if (object->heap()) {
+                Traits::enqueue(markQueue, object);
+            } else {
+                traverseReferredObjects(object, [&](ObjHeader* field) noexcept {
+                    // Each permanent and stack object has own entry in the root set.
+                    if (field->heap() && !isNullOrMarker(field)) {
+                        Traits::enqueue(markQueue, field);
+                    }
+                });
+                RuntimeAssert(!object->has_meta_object(), "Non-heap object %p may not have an extra object data", object);
+            }
+            switch (value.source) {
+                case mm::ThreadRootSet::Source::kStack:
+                    ++stack;
+                    break;
+                case mm::ThreadRootSet::Source::kTLS:
+                    ++tls;
+                    break;
             }
         }
-        RuntimeLogDebug({kTagGC}, "Collected root set for thread stack=%zu tls=%zu", stack, tls);
     }
+    RuntimeLogDebug({kTagGC}, "Collected root set for thread stack=%zu tls=%zu", stack, tls);
+}
+
+template <typename Traits>
+void collectRootSetGlobals(typename Traits::MarkQueue& markQueue) noexcept {
     mm::StableRefRegistry::Instance().ProcessDeletions();
     size_t global = 0;
     size_t stableRef = 0;
@@ -174,6 +184,19 @@ void collectRootSet(typename Traits::MarkQueue& markQueue) noexcept {
     RuntimeLogDebug({kTagGC}, "Collected global root set global=%zu stableRef=%zu", global, stableRef);
 }
 
+// TODO: This needs some tests now.
+template <typename Traits, typename F>
+void collectRootSet(typename Traits::MarkQueue& markQueue, F&& filter) noexcept {
+    Traits::clear(markQueue);
+    for (auto& thread : mm::GlobalData::Instance().threadRegistry().LockForIter()) {
+        if (!filter(thread))
+            continue;
+        thread.Publish();
+        collectRootSetForThread<Traits>(markQueue, thread);
+    }
+    collectRootSetGlobals<Traits>(markQueue);
+}
+
 } // namespace gc
 } // namespace kotlin
 

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

@@ -154,7 +154,7 @@ bool gc::SameThreadMarkAndSweep::PerformFullGC() noexcept {
 
         RuntimeLogInfo(
                 {kTagGC}, "Started GC epoch %zu. Time since last GC %" PRIu64 " microseconds", epoch_, timeStartUs - lastGCTimestampUs_);
-        gc::collectRootSet<MarkTraits>(markQueue_);
+        gc::collectRootSet<MarkTraits>(markQueue_, [] (mm::ThreadData&) { return true; });
         auto timeRootSetUs = konan::getTimeMicros();
         // Can be unsafe, because we've stopped the world.
         auto objectsCountBefore = objectFactory_.GetSizeUnsafe();

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

@@ -19,6 +19,7 @@ using Kotlin_getSourceInfo_FunctionType = int(*)(void * /*addr*/, SourceInfo* /*
  * but can be changed after compiling caches. So use this way for variables, which will be rarely accessed.
  */
 RUNTIME_WEAK int32_t Kotlin_destroyRuntimeMode = 1;
+RUNTIME_WEAK int32_t Kotlin_gcMarkSingleThreaded = 1;
 RUNTIME_WEAK int32_t Kotlin_workerExceptionHandling = 0;
 RUNTIME_WEAK int32_t Kotlin_suspendFunctionsFromAnyThreadFromObjC = 0;
 RUNTIME_WEAK Kotlin_getSourceInfo_FunctionType Kotlin_getSourceInfo_Function = nullptr;
@@ -33,6 +34,10 @@ ALWAYS_INLINE compiler::DestroyRuntimeMode compiler::destroyRuntimeMode() noexce
     return static_cast<compiler::DestroyRuntimeMode>(Kotlin_destroyRuntimeMode);
 }
 
+ALWAYS_INLINE bool compiler::gcMarkSingleThreaded() noexcept {
+    return Kotlin_gcMarkSingleThreaded != 0;
+}
+
 ALWAYS_INLINE compiler::WorkerExceptionHandling compiler::workerExceptionHandling() noexcept {
     return static_cast<compiler::WorkerExceptionHandling>(Kotlin_workerExceptionHandling);
 }

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

@@ -105,6 +105,7 @@ ALWAYS_INLINE inline GCSchedulerType getGCSchedulerType() noexcept {
 
 WorkerExceptionHandling workerExceptionHandling() noexcept;
 DestroyRuntimeMode destroyRuntimeMode() noexcept;
+bool gcMarkSingleThreaded() noexcept;
 bool suspendFunctionsFromAnyThreadFromObjCEnabled() noexcept;
 AppStateTracking appStateTracking() noexcept;
 int getSourceInfo(void* addr, SourceInfo *result, int result_size) noexcept;

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

@@ -42,6 +42,7 @@ extern "C" const char* Kotlin_callsCheckerGoodFunctionNames[] = {
         "_ZNSt6chrono3_V212steady_clock3nowEv", // std::chrono::_V2::steady_clock::now()
         "_ZNSt8__detail15_List_node_base7_M_hookEPS0_", // std::__detail::_List_node_base::_M_hook(std::__detail::_List_node_base*)
         "_ZNSt8__detail15_List_node_base9_M_unhookEv", // std::__detail::_List_node_base::_M_unhook()
+        "_ZNSt8__detail15_List_node_base11_M_transferEPS0_S1_", // std::__detail::_List_node_base::_M_transfer(std::__detail::_List_node_base*, std::__detail::_List_node_base*)
         "_ZNSt9exceptionD2Ev", // std::exception::~exception()
         "_ZSt17current_exceptionv", // std::current_exception()
         "_ZSt17rethrow_exceptionSt13exception_ptr", // std::rethrow_exception(std::exception_ptr)

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

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

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

@@ -44,20 +44,23 @@ void yield() noexcept {
 
 THREAD_LOCAL_VARIABLE bool gSuspensionRequestedByCurrentThread = false;
 [[clang::no_destroy]] std::mutex gSuspensionMutex;
-[[clang::no_destroy]] std::condition_variable gSuspendsionCondVar;
+[[clang::no_destroy]] std::condition_variable gSuspensionCondVar;
+[[clang::no_destroy]] std::function<void(kotlin::mm::ThreadData&)> gPreSuspend = nullptr;
 
 } // namespace
 
 std::atomic<bool> kotlin::mm::internal::gSuspensionRequested = false;
 
 NO_EXTERNAL_CALLS_CHECK void kotlin::mm::ThreadSuspensionData::suspendIfRequestedSlowPath() noexcept {
-    std::unique_lock lock(gSuspensionMutex);
     if (IsThreadSuspensionRequested()) {
+        if (gPreSuspend)
+            gPreSuspend(threadData_);
+        std::unique_lock lock(gSuspensionMutex);
         auto threadId = konan::currentThreadId();
         auto suspendStartMs = konan::getTimeMicros();
         RuntimeLogDebug({kTagGC, kTagMM}, "Suspending thread %d", threadId);
-        AutoReset scopedAssign(&suspended_, true);
-        gSuspendsionCondVar.wait(lock, []() { return !IsThreadSuspensionRequested(); });
+        AutoReset scopedAssignSuspended(&suspended_, true);
+        gSuspensionCondVar.wait(lock, []() { return !IsThreadSuspensionRequested(); });
         auto suspendEndMs = konan::getTimeMicros();
         RuntimeLogDebug({kTagGC, kTagMM}, "Resuming thread %d after %" PRIu64 " microseconds of suspension",
                         threadId, suspendEndMs - suspendStartMs);
@@ -96,6 +99,10 @@ ALWAYS_INLINE void kotlin::mm::SuspendIfRequested() noexcept {
     }
 }
 
+void kotlin::mm::SetOnSuspendCallback(std::function<void(mm::ThreadData&)> preSuspend) noexcept {
+    gPreSuspend = preSuspend;
+}
+
 void kotlin::mm::ResumeThreads() noexcept {
     // From the std::condition_variable docs:
     // Even if the shared variable is atomic, it must be modified under
@@ -106,5 +113,5 @@ void kotlin::mm::ResumeThreads() noexcept {
         internal::gSuspensionRequested = false;
     }
     gSuspensionRequestedByCurrentThread = false;
-    gSuspendsionCondVar.notify_all();
+    gSuspensionCondVar.notify_all();
 }

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

@@ -13,6 +13,8 @@
 namespace kotlin {
 namespace mm {
 
+class ThreadData;
+
 namespace internal {
 
 extern std::atomic<bool> gSuspensionRequested;
@@ -26,7 +28,7 @@ inline bool IsThreadSuspensionRequested() noexcept {
 
 class ThreadSuspensionData : private Pinned {
 public:
-    explicit ThreadSuspensionData(ThreadState initialState) noexcept : state_(initialState), suspended_(false) {}
+    explicit ThreadSuspensionData(ThreadState initialState, mm::ThreadData& threadData) noexcept : state_(initialState), threadData_(threadData), suspended_(false) {}
 
     ~ThreadSuspensionData() = default;
 
@@ -52,6 +54,7 @@ private:
     friend void SuspendIfRequestedSlowPath() noexcept;
 
     std::atomic<ThreadState> state_;
+    mm::ThreadData& threadData_;
     std::atomic<bool> suspended_;
     void suspendIfRequestedSlowPath() noexcept;
 };
@@ -75,6 +78,11 @@ inline bool SuspendThreads() noexcept {
     return true;
 }
 
+/**
+ * Set function `preSuspend` which will be called in threads before suspending.
+ */
+void SetOnSuspendCallback(std::function<void(mm::ThreadData&)>) noexcept;
+
 /**
  * Resumes all threads registered in ThreadRegistry that were suspended by the SuspendThreads call.
  * Does not wait until all such threads are actually resumed.