From 38e7a42a965384c3f6f19d751e70c9e921f03ad8 Mon Sep 17 00:00:00 2001
From: Thomas Heller <thom.heller@gmail.com>
Date: Fri, 1 Dec 2023 08:10:55 +0100
Subject: [PATCH] Fixing exec::__atomic_intrusive_queue (#1160)

* Fixing exec::__atomic_intrusive_queue

exec::__atomic_intrusive_queue::pop_all and
exec::__atomic_intrusive_queue::pop_all_reversed used a simple atomic
exchange which might lead to a race with a parallel insertion.

This patch fixes that by introducing a CAS loop to account for potential
parallel insertions.
---
 .../__detail/__atomic_intrusive_queue.hpp     | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/include/exec/__detail/__atomic_intrusive_queue.hpp b/include/exec/__detail/__atomic_intrusive_queue.hpp
index f5cdf50f9..8b8beee19 100644
--- a/include/exec/__detail/__atomic_intrusive_queue.hpp
+++ b/include/exec/__detail/__atomic_intrusive_queue.hpp
@@ -22,11 +22,11 @@ namespace exec {
   template <auto _NextPtr>
   class __atomic_intrusive_queue;
 
-  template <class _Tp, _Tp* _Tp::*_NextPtr>
+  template <class _Tp, _Tp *_Tp::*_NextPtr>
   class alignas(64) __atomic_intrusive_queue<_NextPtr> {
    public:
-    using __node_pointer = _Tp*;
-    using __atomic_node_pointer = std::atomic<_Tp*>;
+    using __node_pointer = _Tp *;
+    using __atomic_node_pointer = std::atomic<_Tp *>;
 
     [[nodiscard]] bool empty() const noexcept {
       return __head_.load(std::memory_order_relaxed) == nullptr;
@@ -65,16 +65,22 @@ namespace exec {
     }
 
     stdexec::__intrusive_queue<_NextPtr> pop_all() noexcept {
-      return stdexec::__intrusive_queue<_NextPtr>::make(
-        __head_.exchange(nullptr, std::memory_order_acq_rel));
+      return stdexec::__intrusive_queue<_NextPtr>::make(reset_head());
     }
 
     stdexec::__intrusive_queue<_NextPtr> pop_all_reversed() noexcept {
-      return stdexec::__intrusive_queue<_NextPtr>::make_reversed(
-        __head_.exchange(nullptr, std::memory_order_acq_rel));
+      return stdexec::__intrusive_queue<_NextPtr>::make_reversed(reset_head());
     }
 
    private:
+    __node_pointer reset_head() noexcept {
+      __node_pointer __old_head = __head_.load(std::memory_order_relaxed);
+      while (!__head_.compare_exchange_weak(__old_head, nullptr, std::memory_order_acq_rel)) {
+        ;
+      }
+      return __old_head;
+    }
+
     __atomic_node_pointer __head_{nullptr};
   };
 }
\ No newline at end of file