1 files changed, 84 insertions, 39 deletions

diff --git a/src/common/polyfill_thread.h b/src/common/polyfill_thread.h
index 5a8d1ce08..b5ef055db 100644
--- a/src/common/polyfill_thread.h
+++ b/src/common/polyfill_thread.h
@@ -11,6 +11,8 @@
 
 #ifdef __cpp_lib_jthread
 
+#include <chrono>
+#include <condition_variable>
 #include <stop_token>
 #include <thread>
 
@@ -21,23 +23,36 @@ void CondvarWait(Condvar& cv, Lock& lock, std::stop_token token, Pred&& pred) {
     cv.wait(lock, token, std::move(pred));
 }
 
+template <typename Rep, typename Period>
+bool StoppableTimedWait(std::stop_token token, const std::chrono::duration<Rep, Period>& rel_time) {
+    std::condition_variable_any cv;
+    std::mutex m;
+
+    // Perform the timed wait.
+    std::unique_lock lk{m};
+    return !cv.wait_for(lk, token, rel_time, [&] { return token.stop_requested(); });
+}
+
 } // namespace Common
 
 #else
 
 #include <atomic>
+#include <chrono>
+#include <condition_variable>
 #include <functional>
-#include <list>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <optional>
 #include <thread>
 #include <type_traits>
+#include <utility>
 
 namespace std {
 namespace polyfill {
 
-using stop_state_callbacks = list<function<void()>>;
+using stop_state_callback = size_t;
 
 class stop_state {
 public:
@@ -45,61 +60,69 @@ public:
     ~stop_state() = default;
 
     bool request_stop() {
-        stop_state_callbacks callbacks;
+        unique_lock lk{m_lock};
 
-        {
-            scoped_lock lk{m_lock};
+        if (m_stop_requested) {
+            // Already set, nothing to do.
+            return false;
+        }
 
-            if (m_stop_requested.load()) {
-                // Already set, nothing to do
-                return false;
-            }
+        // Mark stop requested.
+        m_stop_requested = true;
 
-            // Set as requested
-            m_stop_requested = true;
+        while (!m_callbacks.empty()) {
+            // Get an iterator to the first element.
+            const auto it = m_callbacks.begin();
 
-            // Copy callback list
-            callbacks = m_callbacks;
-        }
+            // Move the callback function out of the map.
+            function<void()> f;
+            swap(it->second, f);
+
+            // Erase the now-empty map element.
+            m_callbacks.erase(it);
 
-        for (auto callback : callbacks) {
-            callback();
+            // Run the callback.
+            if (f) {
+                f();
+            }
         }
 
         return true;
     }
 
     bool stop_requested() const {
-        return m_stop_requested.load();
+        unique_lock lk{m_lock};
+        return m_stop_requested;
     }
 
-    stop_state_callbacks::const_iterator insert_callback(function<void()> f) {
-        stop_state_callbacks::const_iterator ret{};
-        bool should_run{};
-
-        {
-            scoped_lock lk{m_lock};
-            should_run = m_stop_requested.load();
-            m_callbacks.push_front(f);
-            ret = m_callbacks.begin();
-        }
+    stop_state_callback insert_callback(function<void()> f) {
+        unique_lock lk{m_lock};
 
-        if (should_run) {
-            f();
+        if (m_stop_requested) {
+            // Stop already requested. Don't insert anything,
+            // just run the callback synchronously.
+            if (f) {
+                f();
+            }
+            return 0;
         }
 
+        // Insert the callback.
+        stop_state_callback ret = ++m_next_callback;
+        m_callbacks.emplace(ret, move(f));
         return ret;
     }
 
-    void remove_callback(stop_state_callbacks::const_iterator it) {
-        scoped_lock lk{m_lock};
-        m_callbacks.erase(it);
+    void remove_callback(stop_state_callback cb) {
+        unique_lock lk{m_lock};
+        m_callbacks.erase(cb);
     }
 
 private:
-    mutex m_lock;
-    atomic<bool> m_stop_requested;
-    stop_state_callbacks m_callbacks;
+    mutable recursive_mutex m_lock;
+    map<stop_state_callback, function<void()>> m_callbacks;
+    stop_state_callback m_next_callback{0};
+    bool m_stop_requested{false};
 };
 
 } // namespace polyfill
@@ -190,7 +213,7 @@ public:
     using callback_type = Callback;
 
     template <typename C>
-    requires constructible_from<Callback, C>
+        requires constructible_from<Callback, C>
     explicit stop_callback(const stop_token& st,
                            C&& cb) noexcept(is_nothrow_constructible_v<Callback, C>)
         : m_stop_state(st.m_stop_state) {
@@ -199,7 +222,7 @@ public:
         }
     }
     template <typename C>
-    requires constructible_from<Callback, C>
+        requires constructible_from<Callback, C>
     explicit stop_callback(stop_token&& st,
                            C&& cb) noexcept(is_nothrow_constructible_v<Callback, C>)
         : m_stop_state(move(st.m_stop_state)) {
@@ -209,7 +232,7 @@ public:
     }
     ~stop_callback() {
         if (m_stop_state && m_callback) {
-            m_stop_state->remove_callback(*m_callback);
+            m_stop_state->remove_callback(m_callback);
         }
     }
 
@@ -220,7 +243,7 @@ public:
 
 private:
     shared_ptr<polyfill::stop_state> m_stop_state;
-    optional<polyfill::stop_state_callbacks::const_iterator> m_callback;
+    polyfill::stop_state_callback m_callback;
 };
 
 template <typename Callback>
@@ -318,6 +341,28 @@ void CondvarWait(Condvar& cv, Lock& lock, std::stop_token token, Pred pred) {
     cv.wait(lock, [&] { return pred() || token.stop_requested(); });
 }
 
+template <typename Rep, typename Period>
+bool StoppableTimedWait(std::stop_token token, const std::chrono::duration<Rep, Period>& rel_time) {
+    if (token.stop_requested()) {
+        return false;
+    }
+
+    bool stop_requested = false;
+    std::condition_variable cv;
+    std::mutex m;
+
+    std::stop_callback cb(token, [&] {
+        // Wake up the waiting thread.
+        std::unique_lock lk{m};
+        stop_requested = true;
+        cv.notify_one();
+    });
+
+    // Perform the timed wait.
+    std::unique_lock lk{m};
+    return !cv.wait_for(lk, rel_time, [&] { return stop_requested; });
+}
+
 } // namespace Common
 
 #endif