From e533a64f9b7245ba49596a889c4163b29f8d1eac Mon Sep 17 00:00:00 2001
From: fantasy-peak <1356346239@qq.com>
Date: Mon, 12 Aug 2024 11:24:18 +0800
Subject: [PATCH] Add coroutine spinlock

---
 lib/inc/drogon/utils/coroutine.h     | 66 ++++++++++++++++++++++++++++
 lib/tests/unittests/CoroutineTest.cc | 29 ++++++++++++
 2 files changed, 95 insertions(+)

diff --git a/lib/inc/drogon/utils/coroutine.h b/lib/inc/drogon/utils/coroutine.h
index 9c6678d2e4..9c1d363c76 100644
--- a/lib/inc/drogon/utils/coroutine.h
+++ b/lib/inc/drogon/utils/coroutine.h
@@ -19,6 +19,7 @@
 #include <algorithm>
 #include <atomic>
 #include <cassert>
+#include <chrono>
 #include <condition_variable>
 #include <coroutine>
 #include <exception>
@@ -987,4 +988,69 @@ class Mutex final
     CoroMutexAwaiter *waiters_;
 };
 
+class SpinLock final
+{
+  public:
+    explicit SpinLock(std::int32_t count = 1024) noexcept
+        : spinCount_(count), locked_(false)
+    {
+    }
+
+    bool try_lock() noexcept
+    {
+        return !locked_.exchange(true, std::memory_order_acquire);
+    }
+
+    Task<> coro_lock() noexcept
+    {
+        auto counter = spinCount_;
+        while (!try_lock())
+        {
+            while (locked_.load(std::memory_order_relaxed))
+            {
+                if (counter-- <= 0)
+                {
+                    trantor::EventLoop *loop =
+                        trantor::EventLoop::getEventLoopOfCurrentThread();
+                    assert(loop != nullptr);
+                    co_await sleepCoro(loop, std::chrono::milliseconds(1));
+                    counter = spinCount_;
+                }
+            }
+        }
+        co_return;
+    }
+
+    void lock() noexcept
+    {
+        auto counter = spinCount_;
+        while (!try_lock())
+        {
+            while (locked_.load(std::memory_order_relaxed))
+            {
+                if (counter-- <= 0)
+                {
+                    std::this_thread::yield();
+                    counter = spinCount_;
+                }
+            }
+        }
+    }
+
+    void unlock() noexcept
+    {
+        locked_.store(false, std::memory_order_release);
+    }
+
+    Task<std::unique_lock<SpinLock>> scoped_lock() noexcept
+    {
+        co_await coro_lock();
+        co_return std::unique_lock<SpinLock>{*this, std::adopt_lock};
+    }
+
+  private:
+    int32_t spinCount_;
+    std::atomic_bool locked_;
+};
+
 }  // namespace drogon
diff --git a/lib/tests/unittests/CoroutineTest.cc b/lib/tests/unittests/CoroutineTest.cc
index c2e0f9dcff..beea963e86 100644
--- a/lib/tests/unittests/CoroutineTest.cc
+++ b/lib/tests/unittests/CoroutineTest.cc
@@ -245,3 +245,32 @@ DROGON_TEST(Mutex)
         pool.getLoop(i)->quit();
     pool.wait();
 }
+
+DROGON_TEST(SpinLock)
+{
+    trantor::EventLoopThreadPool pool{3};
+    pool.start();
+    SpinLock spin_lock;
+    async_run([&]() -> Task<> {
+        co_await switchThreadCoro(pool.getLoop(0));
+        auto guard = co_await spin_lock.scoped_lock();
+        co_await sleepCoro(pool.getLoop(1), std::chrono::seconds(2));
+        co_return;
+    });
+    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    std::promise<void> done;
+    async_run([&]() -> Task<> {
+        co_await switchThreadCoro(pool.getLoop(2));
+        auto id = std::this_thread::get_id();
+        co_await spin_lock.coro_lock();
+        CHECK(id == std::this_thread::get_id());
+        spin_lock.unlock();
+        CHECK(id == std::this_thread::get_id());
+        done.set_value();
+        co_return;
+    });
+    done.get_future().wait();
+    for (int16_t i = 0; i < 3; i++)
+        pool.getLoop(i)->quit();
+    pool.wait();
+}