mutex.hpp: Change the mutex implementation

rt/include/typon/mutex.hpp

@@ -3,6 +3,7 @@
 
 #include <atomic>
 #include <coroutine>
+#include <cstdint>
 
 #include <typon/scheduler.hpp>
 #include <typon/stack.hpp>
@@ -13,42 +14,45 @@ namespace typon
 
   /* An asynchronous mutex.
 
-     Based on the MCS lock, but without the spinning. This means it should be
-     easy to implement spinning before suspension.
-
-     Disadvantage: each lock acquisition incurs an allocation. The alternative
-     is to potentially spin when unlocking.
+     Inspired by the implementation in CppCoro by Lewis Baker.
+     https://github.com/lewissbaker/cppcoro
   */
   struct Mutex
   {
     struct Node
     {
-      std::atomic<Node *> _next { nullptr };
-      Stack * _stack;
+      Node * _next;
+      std::atomic<std::uintptr_t> _stack { 0 };
     };
 
     std::atomic<Node *> _state { nullptr };
+    Node * _waiters { nullptr };
 
     [[nodiscard]] auto lock() noexcept
     {
-      struct awaitable
+      struct awaitable : Node
       {
-        std::atomic<Node *> & _state;
-        Node * _prev;
-        Node * _node;
+        Mutex * _mutex;
 
         bool await_ready() noexcept
         {
-          return !_prev;
+          Node * next = nullptr;
+          for(;;)
+          {
+            // _next must be properly set before updating _state
+            _next = next;
+            if(_mutex->_state.compare_exchange_weak(next, this))
+            {
+              return !next;
+            }
+          }
         }
 
         void await_suspend(std::coroutine_handle<> coroutine) noexcept
         {
-          auto stack = _node->_stack = Scheduler::suspend(coroutine);
-          auto ready = _prev->_next.exchange(_node);
-          if (ready)
+          auto stack = Scheduler::suspend(coroutine);
+          if (_stack.exchange(reinterpret_cast<std::uintptr_t>(stack)))
           {
-            delete _prev;
             Scheduler::enable(stack);
           }
         }
@@ -57,28 +61,33 @@ namespace typon
 
         ~awaitable()
         {
-          Node * node = _node;
-          bool waiters = !_state.compare_exchange_strong(node, nullptr);
-          if (waiters)
+          Node * waiter = _mutex->_waiters;
+          if (!waiter)
+          {
+            Node * state = this;
+            if (_mutex->_state.compare_exchange_strong(state, nullptr))
             {
-            node = _node;
-            auto next = node->_next.exchange(_node);
-            if (next)
+              return;
+            }
+            auto next = state;
+            while (next != this)
             {
-              delete node;
-              Scheduler::enable(next->_stack);
+              auto tmp = next->_next;
+              next->_next = waiter;
+              waiter = next;
+              next = tmp;
             }
           }
-          else
+          _mutex->_waiters = waiter->_next;
+          auto stack = waiter->_stack.exchange(1);
+          if (stack)
           {
-            delete node;
+            Scheduler::enable(reinterpret_cast<Stack *>(stack));
           }
         }
       };
 
-      auto node = new Node();
-      auto prev = _state.exchange(node);
-      return awaitable { this->_state, prev, node };
+      return awaitable { {}, this };
     }
   };