Merge pull request #685 from pika-org/barrier-no-yield

Add timed busy-wait option to `pika::barrier` `wait` and `arrive_and_wait`
pika-org · May 12, 2023 · fdbb461 · fdbb461
2 parents f4a9cb1 + 83dc22e
commit fdbb461
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 4 deletions.
diff --git a/libs/pika/execution_base/include/pika/execution_base/this_thread.hpp b/libs/pika/execution_base/include/pika/execution_base/this_thread.hpp
@@ -199,5 +199,61 @@ namespace pika::util {
                 }
             }
         }
+
+        // yield_while_timeout is similar to yield_while, with the
+        // addition of a timeout parameter. If the timeout is exceeded, waiting
+        // is stopped and the function returns false. If the predicate is
+        // successfully waited for the function returns true.
+        template <typename Predicate>
+        [[nodiscard]] bool
+        yield_while_timeout(Predicate&& predicate, std::chrono::duration<double> timeout,
+            const char* thread_name = nullptr, bool allow_timed_suspension = true)
+        {
+            // Seconds represented using a double
+            using duration_type = std::chrono::duration<double>;
+
+            bool use_timeout = timeout >= duration_type(0.0);
+
+            pika::chrono::detail::high_resolution_timer t;
+
+            if (allow_timed_suspension)
+            {
+                for (std::size_t k = 0;; ++k)
+                {
+                    if (use_timeout && duration_type(t.elapsed()) > timeout)
+                    {
+                        return false;
+                    }
+
+                    if (!predicate())
+                    {
+                        return true;
+                    }
+                    else
+                    {
+                        pika::execution::this_thread::detail::yield_k(k, thread_name);
+                    }
+                }
+            }
+            else
+            {
+                for (std::size_t k = 0;; ++k)
+                {
+                    if (use_timeout && duration_type(t.elapsed()) > timeout)
+                    {
+                        return false;
+                    }
+
+                    if (!predicate())
+                    {
+                        return true;
+                    }
+                    else
+                    {
+                        pika::execution::this_thread::detail::yield_k(k % 16, thread_name);
+                    }
+                }
+            }
+        }
     }    // namespace detail
 }    // namespace pika::util
diff --git a/libs/pika/synchronization/include/pika/synchronization/barrier.hpp b/libs/pika/synchronization/include/pika/synchronization/barrier.hpp
@@ -15,6 +15,7 @@
 #include <pika/synchronization/detail/condition_variable.hpp>
 #include <pika/thread_support/assert_owns_lock.hpp>
 
+#include <chrono>
 #include <climits>
 #include <cstddef>
 #include <cstdint>
@@ -181,20 +182,35 @@ namespace pika {
         //                 ([thread.req.exception]).
         // Error conditions: Any of the error conditions allowed for mutex
         //                 types ([thread.mutex.requirements.mutex]).
-        void wait(arrival_token&& old_phase) const
+        void wait(arrival_token&& old_phase,
+            std::chrono::duration<double> busy_wait_timeout = std::chrono::duration<double>(
+                0.0)) const
         {
+            bool const do_busy_wait = busy_wait_timeout > std::chrono::duration<double>(0.0);
+            if (do_busy_wait &&
+                pika::util::detail::yield_while_timeout(
+                    [&]() {
+                        std::unique_lock<mutex_type> l(mtx_);
+                        return phase_ == old_phase;
+                    },
+                    busy_wait_timeout, "barrier::wait", false))
+            {
+                return;
+            }
+
             std::unique_lock<mutex_type> l(mtx_);
             if (phase_ == old_phase)
             {
                 cond_.wait(l, "barrier::wait");
-                PIKA_ASSERT(phase_ != old_phase);
             }
+            PIKA_ASSERT(phase_ != old_phase);
         }
 
         /// Effects:        Equivalent to: wait(arrive()).
-        void arrive_and_wait()
+        void arrive_and_wait(
+            std::chrono::duration<double> busy_wait_timeout = std::chrono::duration<double>(0.0))
         {
-            wait(arrive());
+            wait(arrive(), busy_wait_timeout);
         }
 
         // Preconditions:  The expected count for the current barrier phase is