Set interrupt handler correctly on SharedPromise

Summary: If SharedPromise::getFuture() is invoked after a call to setInterruptHandler,
then the interrupt handler isn't set on the newly created promise. This diff
fixes that.

Reviewed By: yfeldblum

Differential Revision: D2610289

fb-gh-sync-id: bf8fce9e881b83ccac17d13c6788ec2afd0b0153

folly/futures/SharedPromise-inl.h

@@ -41,6 +41,7 @@ SharedPromise<T>& SharedPromise<T>::operator=(
   std::swap(size_, other.size_);
   std::swap(hasValue_, other.hasValue_);
   std::swap(try_, other.try_);
+  std::swap(interruptHandler_, other.interruptHandler_);
   std::swap(promises_, other.promises_);
 
   return *this;
@@ -60,6 +61,9 @@ Future<T> SharedPromise<T>::getFuture() {
     return makeFuture<T>(Try<T>(try_));
   } else {
     promises_.emplace_back();
+    if (interruptHandler_) {
+      promises_.back().setInterruptHandler(interruptHandler_);
+    }
     return promises_.back().getFuture();
   }
 }
@@ -88,6 +92,7 @@ void SharedPromise<T>::setInterruptHandler(
   if (hasValue_) {
     return;
   }
+  interruptHandler_ = fn;
   for (auto& p : promises_) {
     p.setInterruptHandler(fn);
   }

folly/futures/SharedPromise.h

@@ -113,6 +113,7 @@ private:
   bool hasValue_{false};
   Try<T> try_;
   std::vector<Promise<T>> promises_;
+  std::function<void(exception_wrapper const&)> interruptHandler_;
 };
 
 }

folly/futures/test/SharedPromiseTest.cpp

@@ -117,3 +117,12 @@ TEST(SharedPromise, isFulfilled) {
   p = std::move(p2);
   EXPECT_TRUE(p.isFulfilled());
 }
+
+TEST(SharedPromise, interruptHandler) {
+  SharedPromise<int> p;
+  bool flag = false;
+  p.setInterruptHandler([&](const exception_wrapper&) { flag = true; });
+  auto f = p.getFuture();
+  f.cancel();
+  EXPECT_TRUE(flag);
+}