(Wangle) within should raise TimedOut()

Summary: I figured it out. This works. The two extra futures are a small overhead
(just two pointers). The `Core`s are allocated anyway, so this is pretty
much optimal.

A timeout now raises on the current Future, and a fulfilled promise
cancels the timeout Future.

Reviewed By: @yfeldblum

Differential Revision: D2232463

folly/futures/Future-inl.h

@@ -855,34 +855,41 @@ template <class E>
 Future<T> Future<T>::within(Duration dur, E e, Timekeeper* tk) {
 
   struct Context {
-    Context(E ex) : exception(std::move(ex)), promise() {}
+    Context(E ex, Future<Unit>&& f)
+        : exception(std::move(ex)), afterFuture(std::move(f)), promise() {}
     E exception;
+    Future<Unit> afterFuture;
+    Future<Unit> thisFuture;
     Promise<T> promise;
     std::atomic<bool> token {false};
   };
-  auto ctx = std::make_shared<Context>(std::move(e));
 
   if (!tk) {
     tk = folly::detail::getTimekeeperSingleton();
   }
 
-  tk->after(dur)
-    .then([ctx](Try<Unit> const& t) {
-      if (ctx->token.exchange(true) == false) {
-        if (t.hasException()) {
-          ctx->promise.setException(std::move(t.exception()));
-        } else {
-          ctx->promise.setException(std::move(ctx->exception));
-        }
-      }
-    });
+  auto ctx = std::make_shared<Context>(std::move(e), tk->after(dur));
 
-  this->then([ctx](Try<T>&& t) {
+  ctx->thisFuture = this->then([ctx](Try<T>&& t) mutable {
+    // "this" completed first, cancel "after"
+    ctx->afterFuture.raise(CancelTimer());
     if (ctx->token.exchange(true) == false) {
       ctx->promise.setTry(std::move(t));
     }
   });
 
+  ctx->afterFuture.then([ctx](Try<Unit> const& t) mutable {
+    // "after" completed first, cancel "this"
+    ctx->thisFuture.raise(TimedOut());
+    if (ctx->token.exchange(true) == false) {
+      if (t.hasException()) {
+        ctx->promise.setException(std::move(t.exception()));
+      } else {
+        ctx->promise.setException(std::move(ctx->exception));
+      }
+    }
+  });
+
   return ctx->promise.getFuture().via(getExecutor());
 }
 

folly/futures/FutureException.h

@@ -91,6 +91,11 @@ class TimedOut : public FutureException {
   TimedOut() : FutureException("Timed out") {}
 };
 
+class CancelTimer : public FutureException {
+ public:
+  CancelTimer() : FutureException("Timer should be cancelled") {}
+};
+
 class PredicateDoesNotObtain : public FutureException {
  public:
   PredicateDoesNotObtain() : FutureException("Predicate does not obtain") {}

folly/futures/test/InterruptTest.cpp

@@ -18,6 +18,7 @@
 
 #include <folly/futures/Future.h>
 #include <folly/futures/Promise.h>
+#include <folly/Baton.h>
 
 using namespace folly;
 
@@ -72,3 +73,44 @@ TEST(Interrupt, secondInterruptNoop) {
   f.cancel();
   EXPECT_EQ(1, count);
 }
+
+TEST(Interrupt, withinTimedOut) {
+  Promise<int> p;
+  Baton<> done;
+  p.setInterruptHandler([&](const exception_wrapper& e) { done.post(); });
+  p.getFuture().within(std::chrono::milliseconds(1));
+  // Give it 100ms to time out and call the interrupt handler
+  auto t = std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
+  EXPECT_TRUE(done.timed_wait(t));
+}
+
+class DummyTimeKeeper : public Timekeeper {
+ public:
+  explicit DummyTimeKeeper() : interrupted() {}
+
+  Future<Unit> after(Duration) override {
+    promise.setInterruptHandler(
+      [this](const exception_wrapper& e) {
+        EXPECT_THROW(e.throwException(), CancelTimer);
+        interrupted.post();
+      }
+    );
+    return promise.getFuture();
+  }
+
+  Baton<> interrupted;
+
+ private:
+  Promise<Unit> promise;
+};
+
+TEST(Interrupt, withinCancelTimer) {
+  DummyTimeKeeper tk;
+  Promise<int> p;
+  Baton<> done;
+  p.getFuture().within(std::chrono::milliseconds(10), TimedOut(), &tk);
+  p.setValue(1); // this should cancel the timer
+  // Give it 100ms to interrupt the timer Future
+  auto t = std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
+  EXPECT_TRUE(tk.interrupted.timed_wait(t));
+}