Extend MutexGuarded::when() with support for timeouts.

c++/src/kj/mutex-test.c++

@@ -19,6 +19,13 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
+#if _WIN32
+#define WIN32_LEAN_AND_MEAN 1  // lolz
+#define WINVER 0x0600
+#define _WIN32_WINNT 0x0600
+#define NOGDI  // NOGDI is needed to make EXPECT_EQ(123u, *lock) compile for some reason
+#endif
+
 #include "mutex.h"
 #include "debug.h"
 #include "thread.h"
@@ -26,12 +33,12 @@
 #include <stdlib.h>
 
 #if _WIN32
-#define NOGDI  // NOGDI is needed to make EXPECT_EQ(123u, *lock) compile for some reason
 #include <windows.h>
 #undef NOGDI
 #else
 #include <pthread.h>
 #include <unistd.h>
+#include <time.h>
 #endif
 
 namespace kj {
@@ -39,8 +46,16 @@ namespace {
 
 #if _WIN32
 inline void delay() { Sleep(10); }
+TimePoint now() {
+  return kj::origin<TimePoint>() + GetTickCount64() * kj::MILLISECONDS;
+}
 #else
 inline void delay() { usleep(10000); }
+TimePoint now() {
+  struct timespec now;
+  KJ_SYSCALL(clock_gettime(CLOCK_MONOTONIC, &now));
+  return kj::origin<TimePoint>() + now.tv_sec * kj::SECONDS + now.tv_nsec * kj::NANOSECONDS;
+}
 #endif
 
 TEST(Mutex, MutexGuarded) {
@@ -170,6 +185,83 @@ TEST(Mutex, When) {
   }
 }
 
+TEST(Mutex, WhenWithTimeout) {
+  MutexGuarded<uint> value(123);
+
+  // A timeout that won't expire.
+  static constexpr Duration LONG_TIMEOUT = 10 * kj::SECONDS;
+
+  {
+    uint m = value.when([](uint n) { return n < 200; }, [](uint& n) {
+      ++n;
+      return n + 2;
+    }, LONG_TIMEOUT);
+    KJ_EXPECT(m == 126);
+
+    KJ_EXPECT(*value.lockShared() == 124);
+  }
+
+  {
+    kj::Thread thread([&]() {
+      delay();
+      *value.lockExclusive() = 321;
+    });
+
+    uint m = value.when([](uint n) { return n > 200; }, [](uint& n) {
+      ++n;
+      return n + 2;
+    }, LONG_TIMEOUT);
+    KJ_EXPECT(m == 324);
+
+    KJ_EXPECT(*value.lockShared() == 322);
+  }
+
+  {
+    // Stress test. 100 threads each wait for a value and then set the next value.
+    *value.lockExclusive() = 0;
+
+    auto threads = kj::heapArrayBuilder<kj::Own<kj::Thread>>(100);
+    for (auto i: kj::zeroTo(100)) {
+      threads.add(kj::heap<kj::Thread>([i,&value]() {
+        if (i % 2 == 0) delay();
+        uint m = value.when([i](const uint& n) { return n == i; },
+            [](uint& n) { return n++; }, LONG_TIMEOUT);
+        KJ_ASSERT(m == i);
+      }));
+    }
+
+    uint m = value.when([](uint n) { return n == 100; }, [](uint& n) {
+      return n++;
+    }, LONG_TIMEOUT);
+    KJ_EXPECT(m == 100);
+
+    KJ_EXPECT(*value.lockShared() == 101);
+  }
+
+  {
+    auto start = now();
+    uint m = value.when([](uint n) { return n == 0; }, [&](uint& n) {
+      KJ_ASSERT(n == 101);
+      KJ_EXPECT(now() - start >= 10 * kj::MILLISECONDS);
+      return 12;
+    }, 10 * kj::MILLISECONDS);
+    KJ_EXPECT(m == 12);
+
+    m = value.when([](uint n) { return n == 0; }, [&](uint& n) {
+      KJ_ASSERT(n == 101);
+      KJ_EXPECT(now() - start >= 20 * kj::MILLISECONDS);
+      return 34;
+    }, 10 * kj::MILLISECONDS);
+    KJ_EXPECT(m == 34);
+
+    m = value.when([](uint n) { return n > 0; }, [&](uint& n) {
+      KJ_ASSERT(n == 101);
+      return 56;
+    }, LONG_TIMEOUT);
+    KJ_EXPECT(m == 56);
+  }
+}
+
 TEST(Mutex, Lazy) {
   Lazy<uint> lazy;
   volatile bool initStarted = false;

c++/src/kj/mutex.h

@@ -27,6 +27,7 @@
 
 #include "memory.h"
 #include <inttypes.h>
+#include "time.h"
 
 #if __linux__ && !defined(KJ_USE_FUTEX)
 #define KJ_USE_FUTEX 1
@@ -71,8 +72,10 @@ public:
     virtual bool check() = 0;
   };
 
-  void lockWhen(Predicate& predicate);
-  // Lock (exclusively) when predicate.check() returns true.
+  void lockWhen(Predicate& predicate, Maybe<Duration> timeout = nullptr);
+  // Lock (exclusively) when predicate.check() returns true, or when the timeout (if any) expires.
+  // The mutex is always locked when this returns regardless of whether the timeout expired (and
+  // always unlocked if it throws).
 
 private:
 #if KJ_USE_FUTEX
@@ -100,6 +103,7 @@ private:
     Predicate& predicate;
 #if KJ_USE_FUTEX
     uint futex;
+    bool hasTimeout;
 #elif _WIN32
     uintptr_t condvar;
     // Actually CONDITION_VARIABLE, but don't want to #include <windows.h> in header.
@@ -284,7 +288,8 @@ public:
   // Like `getWithoutLock()`, but asserts that the lock is already held by the calling thread.
 
   template <typename Cond, typename Func>
-  auto when(Cond&& condition, Func&& callback) const -> decltype(callback(instance<T&>())) {
+  auto when(Cond&& condition, Func&& callback, Maybe<Duration> timeout = nullptr) const
+      -> decltype(callback(instance<T&>())) {
     // Waits until condition(state) returns true, then calls callback(state) under lock.
     //
     // `condition`, when called, receives as its parameter a const reference to the state, which is
@@ -295,6 +300,10 @@ public:
     // condition to become true. It may even return true once, but then be called more times.
     // It is guaranteed, though, that at the time `callback()` is finally called, `condition()`
     // would currently return true (assuming it is a pure function of the guarded data).
+    //
+    // If `timeout` is specified, then after the given amount of time, the callback will be called
+    // regardless of whether the condition is true. In this case, when `callback()` is called,
+    // `condition()` may in fact evaluate false, but *only* if the timeout was reached.
 
     struct PredicateImpl final: public _::Mutex::Predicate {
       bool check() override {
@@ -309,7 +318,7 @@ public:
     };
 
     PredicateImpl impl(kj::fwd<Cond>(condition), value);
-    mutex.lockWhen(impl);
+    mutex.lockWhen(impl, timeout);
     KJ_DEFER(mutex.unlock(_::Mutex::EXCLUSIVE));
     return callback(value);
   }