[K/N] ExtraSafePointActionActivator to be used in concurrnt mark

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

@@ -9,6 +9,10 @@
 #include <utility>
 
 #include "Utils.hpp"
+#include "ThreadRegistry.hpp"
+#include "CallsChecker.hpp"
+
+#include <shared_mutex>
 
 namespace kotlin::mm {
 
@@ -36,6 +40,56 @@ private:
 void safePoint(std::memory_order fastPathOrder = std::memory_order_relaxed) noexcept;
 void safePoint(ThreadData& threadData, std::memory_order fastPathOrder = std::memory_order_relaxed) noexcept;
 
+/**
+ * A helper class template to implement custom safe point action activators.
+ * An implementation has to inherit from this template, providing itself as a CRTP template argument.
+ *
+ * It's guaranteed that the action executed through `doIfActive`
+ * will be fully completed before the activator destructor returns.
+ */
+template <typename Impl>
+class ExtraSafePointActionActivator : private MoveOnly {
+public:
+    template <typename Action>
+    static void doIfActive(Action&& action) {
+        CallsCheckerIgnoreGuard guard;
+
+        // Without this check and with many frequent-enough readers,
+        // a writer might never get a chance to obtain the lock.
+        if (!active_.load(std::memory_order_relaxed)) return;
+
+        std::shared_lock lock(mutex_);
+        if (active_.load(std::memory_order_relaxed)) {
+            action();
+        }
+    }
+
+    static bool isActive() noexcept {
+        return active_.load(std::memory_order_relaxed);
+    }
+
+    ExtraSafePointActionActivator() noexcept {
+        std::unique_lock lock(mutex_);
+        active_.store(true, std::memory_order_relaxed);
+    }
+
+    virtual ~ExtraSafePointActionActivator() noexcept = 0;
+
+private:
+    [[clang::no_destroy]] inline static std::shared_mutex mutex_{};
+    inline static std::atomic<bool> active_ = false;
+
+    SafePointActivator safePointActivator_{};
+};
+
+template <typename Impl>
+ExtraSafePointActionActivator<Impl>::~ExtraSafePointActionActivator() noexcept {
+    // First stop new incoming threads from acquiring the mutex.
+    active_.store(false, std::memory_order_relaxed);
+    // Then synchronize with those, already holding the mutex.
+    std::unique_lock lock(mutex_);
+}
+
 namespace test_support {
 
 bool safePointsAreActive() noexcept;

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

@@ -146,3 +146,50 @@ TEST_F(SafePointActionTest, SafePointWithExplicitThread) {
         mm::test_support::setSafePointAction(nullptr);
     });
 }
+
+class ExtraSafePointActionActivatorTest : public ::testing::Test {
+public:
+    class ActivatorImpl : public mm::ExtraSafePointActionActivator<ActivatorImpl> {};
+
+    auto& mockAction() noexcept { return mockAction_; }
+private:
+    testing::StrictMock<testing::MockFunction<void()>> mockAction_;
+};
+
+TEST_F(ExtraSafePointActionActivatorTest, ExtraActionActivator) {
+    ASSERT_FALSE(mm::test_support::safePointsAreActive());
+    ActivatorImpl::doIfActive([&] { mockAction().Call(); });
+    {
+        ActivatorImpl activator;
+        ASSERT_TRUE(mm::test_support::safePointsAreActive());
+
+        EXPECT_CALL(mockAction(), Call());
+        ActivatorImpl::doIfActive([&] { mockAction().Call(); });
+        testing::Mock::VerifyAndClearExpectations(&mockAction());
+    }
+    ActivatorImpl::doIfActive([&] { mockAction().Call(); });
+    ASSERT_FALSE(mm::test_support::safePointsAreActive());
+}
+
+TEST_F(ExtraSafePointActionActivatorTest, ExtraActionActivatorStress) {
+    std::atomic terminate = false;
+
+    std::vector<ScopedThread> threads;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        threads.emplace_back([&]() noexcept {
+            while (!terminate) {
+                ActivatorImpl::doIfActive([&] { mockAction().Call(); });
+                std::this_thread::yield();
+            }
+        });
+    }
+
+    EXPECT_CALL(mockAction(), Call()).Times(testing::AnyNumber());
+    {
+        ActivatorImpl activator;
+        std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    }
+    testing::Mock::VerifyAndClearExpectations(&mockAction());
+
+    terminate = true;
+}