Let Future check before adding a continuation

Summary:
[Folly] Let Future check before adding a continuation whether a continuation has already been added.

Requires making wait behave more like timed-wait; alternatively, it would be possible to add a reverse state transition, but there seems not to be a need.

Reviewed By: marshallcline

Differential Revision: D8344923

fbshipit-source-id: f077fb4b83b92c96af9a83e1b1479c4cb41b048a

folly/futures/Future-inl.h

@@ -244,6 +244,13 @@ void FutureBase<T>::throwIfInvalid() const {
   }
 }
 
+template <class T>
+void FutureBase<T>::throwIfContinued() const {
+  if (!core_ || core_->hasCallback()) {
+    throw_exception<FutureAlreadyContinued>();
+  }
+}
+
 template <class T>
 Optional<Try<T>> FutureBase<T>::poll() {
   auto& core = getCore();
@@ -259,6 +266,7 @@ void FutureBase<T>::raise(exception_wrapper exception) {
 template <class T>
 template <class F>
 void FutureBase<T>::setCallback_(F&& func) {
+  throwIfContinued();
   getCore().setCallback(std::forward<F>(func));
 }
 
@@ -1807,9 +1815,15 @@ void waitImpl(FutureType& f) {
     return;
   }
 
-  FutureBatonType baton;
-  f.setCallback_([&](const Try<T>& /* t */) { baton.post(); });
-  baton.wait();
+  Promise<T> promise;
+  auto ret = promise.getSemiFuture();
+  auto baton = std::make_shared<FutureBatonType>();
+  f.setCallback_([baton, promise = std::move(promise)](Try<T>&& t) mutable {
+    promise.setTry(std::move(t));
+    baton->post();
+  });
+  convertFuture(std::move(ret), f);
+  baton->wait();
   assert(f.isReady());
 }
 

folly/futures/Future.h

@@ -57,6 +57,11 @@ class FOLLY_EXPORT FutureInvalid : public FutureException {
   FutureInvalid() : FutureException("Future invalid") {}
 };
 
+class FOLLY_EXPORT FutureAlreadyContinued : public FutureException {
+ public:
+  FutureAlreadyContinued() : FutureException("Future already continued") {}
+};
+
 class FOLLY_EXPORT FutureNotReady : public FutureException {
  public:
   FutureNotReady() : FutureException("Future not ready") {}
@@ -386,6 +391,7 @@ class FutureBase {
   void detach();
 
   void throwIfInvalid() const;
+  void throwIfContinued() const;
 
   void assign(FutureBase<T>&& other) noexcept;
 
@@ -1789,6 +1795,7 @@ class Future : private futures::detail::FutureBase<T> {
   friend class FutureSplitter;
 
   using Base::setExecutor;
+  using Base::throwIfContinued;
   using Base::throwIfInvalid;
   using typename Base::corePtr;
 

folly/futures/detail/Core.h

@@ -209,6 +209,14 @@ class Core final {
   Core(Core&&) noexcept = delete;
   Core& operator=(Core&&) = delete;
 
+  /// May call from any thread
+  bool hasCallback() const noexcept {
+    constexpr auto allowed = State::OnlyCallback | State::Done;
+    auto const ans = State() != (fsm_.getState() & allowed);
+    assert(!ans || !!callback_); // callback_ must exist in this state
+    return ans;
+  }
+
   /// May call from any thread
   ///
   /// True if state is OnlyResult or Done.

folly/futures/test/WaitTest.cpp

@@ -352,7 +352,10 @@ TEST(Wait, WaitPlusThen) {
 
     EXPECT_FALSE(f.isReady()); // deterministically passes in practice
     f.wait();
-    EXPECT_THROW(f.then([](auto&&) {}), std::logic_error);
+    EXPECT_TRUE(f.isReady());
+    auto continuation = 0;
+    EXPECT_NO_THROW(f.then([&](auto&& v) { continuation = v; }));
+    EXPECT_EQ(continuation, 42);
     t.join();
   }
 
@@ -436,9 +439,12 @@ TEST(Wait, WaitPlusThen) {
 
     EXPECT_FALSE(f.isReady()); // deterministically passes in practice
     f.wait();
+    EXPECT_TRUE(f.isReady());
+    auto continuation = 0;
     InlineExecutor e;
     auto f2 = std::move(f).via(&e);
-    EXPECT_THROW(f2.then([](auto&&) {}), std::logic_error);
+    EXPECT_NO_THROW(f2.then([&](auto&& v) { continuation = v; }));
+    EXPECT_EQ(continuation, 42);
     t.join();
   }