Merge pull request #1001 from dimacurrentai/fix_wa

C5T · May 1, 2024 · fec2e05 · fec2e05
2 parents 39d2edf + 6e9c160
commit fec2e05
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 2 deletions.
diff --git a/bricks/sync/test.cc b/bricks/sync/test.cc
@@ -349,3 +349,42 @@ TEST(WaitableAtomic, WaitFor) {
     EXPECT_FALSE(b2.WaitFor([](bool b) { return b; }, std::chrono::milliseconds(1)));
   }
 }
+
+TEST(WaitableAtomic, NotifiesOfChangesInWait) {
+  struct Obj final {
+    bool done = false;
+    int from = 0;
+    int into = 0;
+  };
+  current::WaitableAtomic<Obj> obj;
+  int extracted_into = 0;
+
+  std::thread t([&obj, &extracted_into]() {
+    bool done = false;
+    while (!done) {
+     obj.Wait(
+         [](const Obj& o) { return o.done || o.from != 0; },
+         [&done, &extracted_into](Obj& o) {
+           if (o.done) {
+             done = true;
+           }
+           if (o.from) {
+             // NOTE(dkorolev): This delay is 0.1ms, which is sufficient. I have tested the code with:
+             // ./.current/test --gtest_filter=WaitableAtomic.NotifiesOfChangesInWai --gtest_repeat=-1
+             std::this_thread::sleep_for(std::chrono::microseconds(100));
+             o.into = o.from;
+             o.from = 0;
+             extracted_into = o.into;
+           }
+         });
+    }
+  });
+
+  EXPECT_EQ(extracted_into, 0);
+  obj.MutableScopedAccessor()->from = 1;
+  obj.Wait([](const Obj& o) { return o.from == 0; });
+  EXPECT_EQ(extracted_into, 1);
+
+  obj.MutableScopedAccessor()->done = true;
+  t.join();
+}
diff --git a/bricks/sync/waitable_atomic.h b/bricks/sync/waitable_atomic.h
@@ -214,17 +214,30 @@ class WaitableAtomic {
   // TODO(dkorolev): The `.Wait()` above always returning `true` could use some TLC.
 
   template <typename F>
-  std::invoke_result_t<F, data_t&> Wait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
+  std::invoke_result_t<F, data_t&> DoWait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
     std::unique_lock<std::mutex> lock(data_mutex_);
     if (!wait_predicate(data_)) {
       const data_t& data = data_;
-      data_condition_variable_.wait(lock, [&wait_predicate, &data] { return wait_predicate(data); });
+      data_condition_variable_.wait(lock, [&wait_predicate, &data]() { return wait_predicate(data); });
       return retval_predicate(data_);
     } else {
       return retval_predicate(data_);
     }
   }
 
+  template <typename F, class = std::enable_if_t<std::is_same_v<std::invoke_result_t<F, data_t&>, void>>>
+  void Wait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
+    DoWait(wait_predicate, std::forward<F>(retval_predicate));
+    Notify();
+  }
+
+  template <typename F, class = std::enable_if_t<!std::is_same_v<std::invoke_result_t<F, data_t&>, void>>>
+  std::invoke_result_t<F, data_t&> Wait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
+    std::invoke_result_t<F, data_t&> retval = DoWait(wait_predicate, std::forward<F>(retval_predicate));
+    Notify();
+    return retval;
+  }
+
 #endif  // CURRENT_FOR_CPP14
 
   template <typename T>