From f8f6e199e97aa224fe6e4916fe3385cccd67fb8e Mon Sep 17 00:00:00 2001
From: Hans Fugal <fugalh@fb.com>
Date: Fri, 2 Jan 2015 14:36:42 -0800
Subject: [PATCH] Future::within

Summary: For when you have a future that you want to complete within a duration, else raise a `TimedOut` exception.

Test Plan: new unit tests

Reviewed By: jsedgwick@fb.com

Subscribers: trunkagent, fugalh, exa, folly-diffs@

FB internal diff: D1756580

Tasks: 4548494

Signature: t1:1756580:1420215704:862f68816fc3a9d05a77077c439bec002aa29cf3
---
 folly/wangle/futures/Future-inl.h            | 59 ++++++++++++++++----
 folly/wangle/futures/Future.h                | 10 ++++
 folly/wangle/futures/test/TimekeeperTest.cpp | 37 ++++++++++++
 3 files changed, 96 insertions(+), 10 deletions(-)

diff --git a/folly/wangle/futures/Future-inl.h b/folly/wangle/futures/Future-inl.h
index 5b887bf8..2c2f0aa6 100644
--- a/folly/wangle/futures/Future-inl.h
+++ b/folly/wangle/futures/Future-inl.h
@@ -725,17 +725,14 @@ namespace {
 
     folly::wangle::detail::getTimekeeperSingleton()->after(dur)
       .then([&,token](Try<void> const& t) {
-        try {
-          t.value();
-          if (token->exchange(true) == false) {
+        if (token->exchange(true) == false) {
+          try {
+            t.value();
             p.setException(TimedOut());
-            baton.post();
-          }
-        } catch (std::exception const& e) {
-          if (token->exchange(true) == false) {
+          } catch (std::exception const& e) {
             p.setException(std::current_exception());
-            baton.post();
           }
+          baton.post();
         }
       });
 
@@ -766,6 +763,7 @@ T Future<T>::get() {
 template <>
 inline void Future<void>::get() {
   getWaitHelper(this);
+  value();
 }
 
 template <class T>
@@ -779,8 +777,49 @@ inline void Future<void>::get(Duration dur) {
 }
 
 template <class T>
-Future<T> Future<T>::delayed(Duration dur, Timekeeper* tk)
-{
+Future<T> Future<T>::within(Duration dur, Timekeeper* tk) {
+  return within(dur, TimedOut(), tk);
+}
+
+template <class T>
+template <class E>
+Future<T> Future<T>::within(Duration dur, E e, Timekeeper* tk) {
+
+  struct Context {
+    Context(E ex) : exception(std::move(ex)), promise(), token(false) {}
+    E exception;
+    Promise<T> promise;
+    std::atomic<bool> token;
+  };
+  auto ctx = std::make_shared<Context>(std::move(e));
+
+  if (!tk) {
+    tk = folly::wangle::detail::getTimekeeperSingleton();
+  }
+
+  tk->after(dur)
+    .then([ctx](Try<void> const& t) {
+      if (ctx->token.exchange(true) == false) {
+        try {
+          t.throwIfFailed();
+          ctx->promise.setException(std::move(ctx->exception));
+        } catch (std::exception const&) {
+          ctx->promise.setException(std::current_exception());
+        }
+      }
+    });
+
+  this->then([ctx](Try<T>&& t) {
+    if (ctx->token.exchange(true) == false) {
+      ctx->promise.fulfilTry(std::move(t));
+    }
+  });
+
+  return ctx->promise.getFuture();
+}
+
+template <class T>
+Future<T> Future<T>::delayed(Duration dur, Timekeeper* tk) {
   return whenAll(*this, futures::sleep(dur, tk))
     .then([](Try<std::tuple<Try<T>, Try<void>>>&& tup) {
       Try<T>& t = std::get<0>(tup.value());
diff --git a/folly/wangle/futures/Future.h b/folly/wangle/futures/Future.h
index a51c2949..066b7465 100644
--- a/folly/wangle/futures/Future.h
+++ b/folly/wangle/futures/Future.h
@@ -462,6 +462,16 @@ class Future {
     raise(FutureCancellation());
   }
 
+  /// Throw TimedOut if this Future does not complete within the given
+  /// duration from now. The optional Timeekeeper is as with futures::sleep().
+  Future<T> within(Duration, Timekeeper* = nullptr);
+
+  /// Throw the given exception if this Future does not complete within the
+  /// given duration from now. The optional Timeekeeper is as with
+  /// futures::sleep().
+  template <class E>
+  Future<T> within(Duration, E exception, Timekeeper* = nullptr);
+
   /// Delay the completion of this Future for at least this duration from
   /// now. The optional Timekeeper is as with futures::sleep().
   Future<T> delayed(Duration, Timekeeper* = nullptr);
diff --git a/folly/wangle/futures/test/TimekeeperTest.cpp b/folly/wangle/futures/test/TimekeeperTest.cpp
index 5d2a359d..6edbb91d 100644
--- a/folly/wangle/futures/test/TimekeeperTest.cpp
+++ b/folly/wangle/futures/test/TimekeeperTest.cpp
@@ -88,3 +88,40 @@ TEST(Timekeeper, futureDelayed) {
 
   EXPECT_GE(dur, one_ms);
 }
+
+TEST(Timekeeper, futureWithinThrows) {
+  Promise<int> p;
+  auto f = p.getFuture()
+    .within(one_ms)
+    .onError([](TimedOut&) { return -1; });
+
+  EXPECT_EQ(-1, f.get());
+}
+
+TEST(Timekeeper, futureWithinAlreadyComplete) {
+  auto f = makeFuture(42)
+    .within(one_ms)
+    .onError([&](TimedOut&){ return -1; });
+
+  EXPECT_EQ(42, f.get());
+}
+
+TEST(Timekeeper, futureWithinFinishesInTime) {
+  Promise<int> p;
+  auto f = p.getFuture()
+    .within(std::chrono::minutes(1))
+    .onError([&](TimedOut&){ return -1; });
+  p.setValue(42);
+
+  EXPECT_EQ(42, f.get());
+}
+
+TEST(Timekeeper, futureWithinVoidSpecialization) {
+  makeFuture().within(one_ms);
+}
+
+TEST(Timekeeper, futureWithinException) {
+  Promise<void> p;
+  auto f = p.getFuture().within(awhile, std::runtime_error("expected"));
+  EXPECT_THROW(f.get(), std::runtime_error);
+}
-- 
2.34.1