Even better implementation

nv-morpheus · Jun 30, 2023 · cfad09b · cfad09b
1 parent d77458d
commit cfad09b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 17 deletions.
diff --git a/cpp/mrc/include/mrc/coroutines/ring_buffer.hpp b/cpp/mrc/include/mrc/coroutines/ring_buffer.hpp
@@ -158,25 +158,44 @@ class RingBuffer
             return (!m_stopped ? RingBufferOpStatus::Success : RingBufferOpStatus::Stopped);
         }
 
-        WriteOperation& use_scheduling_policy(SchedulePolicy policy)
+        WriteOperation& use_scheduling_policy(SchedulePolicy policy) &
         {
             m_policy = policy;
             return *this;
         }
 
-        WriteOperation& resume_immediately()
+        WriteOperation use_scheduling_policy(SchedulePolicy policy) &&
+        {
+            m_policy = policy;
+            return std::move(*this);
+        }
+
+        WriteOperation& resume_immediately() &
         {
             m_policy = SchedulePolicy::Immediate;
             return *this;
         }
 
-        WriteOperation& resume_on(ThreadPool* thread_pool)
+        WriteOperation resume_immediately() &&
+        {
+            m_policy = SchedulePolicy::Immediate;
+            return std::move(*this);
+        }
+
+        WriteOperation& resume_on(ThreadPool* thread_pool) &
         {
             m_policy = SchedulePolicy::Reschedule;
             set_resume_on_thread_pool(thread_pool);
             return *this;
         }
 
+        WriteOperation resume_on(ThreadPool* thread_pool) &&
+        {
+            m_policy = SchedulePolicy::Reschedule;
+            set_resume_on_thread_pool(thread_pool);
+            return std::move(*this);
+        }
+
       private:
         friend RingBuffer;
 

diff --git a/cpp/mrc/tests/coroutines/test_ring_buffer.cpp b/cpp/mrc/tests/coroutines/test_ring_buffer.cpp
@@ -216,25 +216,16 @@ TEST_F(TestCoroRingBuffer, MultiProducerMultiConsumer)
         {
             switch (i % 3)
             {
-            case 0: {
+            case 0:
                 co_await rb.write(i);
                 break;
-            }
-            case 1: {
-                // Must save this to a temporary to workaround compiler bug in GCC==11.3 (working in 11.2/11.4). See:
-                // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=103871#c10
-                auto& task = rb.write(i).resume_immediately();
-                co_await task;
+            case 1:
+                co_await rb.write(i).resume_immediately();
                 break;
-            }
-            case 2: {
-                // Must save this to a temporary to workaround compiler bug in GCC==11.3 (working in 11.2/11.4). See:
-                // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=103871#c10
-                auto& task = rb.write(i).resume_on(&tp);
-                co_await task;
+            case 2:
+                co_await rb.write(i).resume_on(&tp);
                 break;
             }
-            }
         }
 
         producers_latch.count_down();