Improve implementation to be near to Dmitry's

Author: maikel

include/exec/timed_thread_scheduler.hpp

@@ -53,7 +53,7 @@ namespace exec {
         , set_value_{set_value} {
       }
 
-      STDEXEC::__std::atomic<void*> next_{nullptr};
+      STDEXEC::__std::atomic<timed_thread_operation_base*> next_{nullptr};
       command_type command_;
       void (*set_value_)(timed_thread_operation_base*) noexcept;
     };

include/stdexec/__detail/__intrusive_mpsc_queue.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) Dmitiy V'jukov
+ * Copyright (c) Dmitry V'jukov
  * Copyright (c) 2024 Maikel Nadolski
  * Copyright (c) 2024 NVIDIA Corporation
  *
@@ -24,6 +24,8 @@
 
 #include "__atomic.hpp"
 
+#include "stdexec/__detail/__config.hpp"
+
 #include "./__spin_loop_pause.hpp"
 
 namespace STDEXEC {
@@ -32,13 +34,13 @@ namespace STDEXEC {
 
   // _Node must be default_initializable only for the queue to construct an
   // internal "stub" node - only the _Next data element is accessed internally.
-  template <class _Node, __std::atomic<void*> _Node::* _Next>
+  template <class _Node, __std::atomic<_Node*> _Node::* _Next>
     requires __std::default_initializable<_Node>
   class __intrusive_mpsc_queue<_Next> {
 
-    __std::atomic<void*> __back_{&__stub_};
-    __std::atomic<void*> __head_{&__stub_};
-    _Node __stub_;
+    __std::atomic<_Node*> __head_{&__stub_};
+    _Node* __tail_{&__stub_};
+    _Node __stub_{};
 
    public:
 
@@ -47,38 +49,51 @@ namespace STDEXEC {
     }
 
     constexpr auto push_back(_Node* __new_node) noexcept -> bool {
-      (__new_node->*_Next).store(nullptr, __std::memory_order_release);
-      _Node* __prev = static_cast<_Node*>(
-          __head_.exchange(static_cast<void*>(__new_node), __std::memory_order_acq_rel)
-      );
-      bool was_stub = __prev == &__stub_;
-      (__prev->*_Next).store(static_cast<void*>(__new_node), __std::memory_order_release);
-      return was_stub;
+      (__new_node->*_Next).store(nullptr, __std::memory_order_relaxed);
+      _Node* __prev = __head_.exchange(__new_node, __std::memory_order_acq_rel);
+      (__prev->*_Next).store(__new_node, __std::memory_order_release);
+      return __prev == &__stub_;
     }
 
     constexpr auto pop_front() noexcept -> _Node* {
-      _Node* __back = static_cast<_Node*>(__back_.load(__std::memory_order_relaxed));
-      _Node* __next = static_cast<_Node*>((__back->*_Next).load(__std::memory_order_acquire));
-      if (__back == &__stub_) {
-        if (nullptr == __next)
-          return nullptr;
-        __back_.store(static_cast<void*>(__next), __std::memory_order_relaxed);
-        __back = __next;
-        __next = static_cast<_Node*>((__next->*_Next).load(__std::memory_order_acquire));
+      _Node* __tail = this->__tail_;
+      STDEXEC_ASSERT(__tail != nullptr);
+      _Node* __next = (__tail->*_Next).load(__std::memory_order_acquire);
+      // If tail is pointing to the stub node we need to advance it once more
+      if (&__stub_ == __tail) {
+        if (nullptr == __next) {
+            return nullptr;
+        }
+        this->__tail_ = __next;
+        __tail = __next;
+        __next = (__next->*_Next).load(__std::memory_order_acquire);
       }
-      if (__next) {
-        __back_.store(static_cast<void*>(__next), __std::memory_order_relaxed);
-        return __back;
+      // Normal case: there is a next node and we can just advance the tail
+      if (nullptr != __next) {
+        this->__tail_ = __next;
+        return __tail;
       }
-      _Node* __head = static_cast<_Node*>(__head_.load(__std::memory_order_relaxed));
-      if (__back != __head)
+      // Next is nullptr here means that either:
+      // 1) There are no more nodes in the queue
+      // 2) A producer is in the middle of adding a new node
+      const _Node* __head = this->__head_.load(__std::memory_order_acquire);
+      // A producer is in the middle of adding a new node
+      // we cannot return tail as we cannot link the next node yet
+      if (__tail != __head) {
         return nullptr;
+      }
+      // No more nodes in the queue - we need to insert a stub node
+      // to be able to link to an eventual empty state (or new nodes)
       push_back(&__stub_);
-      __next = static_cast<_Node*>((__back->*_Next).load(__std::memory_order_acquire));
-      if (__next) {
-        __back_.store(static_cast<void*>(__next), __std::memory_order_relaxed);
-        return __back;
+      // Now re-attempt to load next
+      __next = (__tail->*_Next).load(__std::memory_order_acquire);
+      if (nullptr != __next) {
+        // Successfully linked either a new node or the stub node
+        this->__tail_ = __next;
+        return __tail;
       }
+      // A producer is in the middle of adding a new node since next is still nullptr 
+      // and not our stub node, thus we cannot link the next node yet
       return nullptr;
     }
   };

test/rrd/intrusive_mpsc_queue.cpp

@@ -19,7 +19,7 @@
 #include <stdexec/__detail/__intrusive_mpsc_queue.hpp>
 
 struct test_node {
-  std::atomic<void*> next_{nullptr};
+  std::atomic<test_node*> next_{nullptr};
   int value_{0};
 
   test_node() = default;
@@ -271,11 +271,11 @@ struct mpsc_five_producers_ordered : rl::test_suite<mpsc_five_producers_ordered,
         queue.push_back(&nodes[start_idx + i]);
       }
     } else {
-      int count = 0;
+      // int count = 0;
       while (consumed_count < TOTAL_ITEMS) {
         test_node* node = queue.pop_front();
         if (node) {
-          consumed_values[count] = node->value_;
+          consumed_values[consumed_count] = node->value_;
           ++consumed_count;
         }
       }

test/stdexec/detail/test_intrusive_mpsc_queue.cpp

@@ -10,7 +10,7 @@
 namespace {
 
   struct test_node {
-    std::atomic<void*> next_{nullptr};
+    std::atomic<test_node*> next_{nullptr};
     int value_{0};
 
     test_node() = default;