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

diff --git a/folly/futures/SharedPromise-inl.h b/folly/futures/SharedPromise-inl.h
index 76bd0caec266b316f5e6af90e8882a0a55909606..8b3ac0f3efbadb69ce50349e63ad2b0927846144 100644 (file)
--- a/folly/futures/SharedPromise-inl.h
+++ b/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);
   }

diff --git a/folly/futures/SharedPromise.h b/folly/futures/SharedPromise.h
index 114193029cf4fe5ff15475e0227dbadccc3d2c36..cda8f40b7c1bf9fc11af3f2994a2833b584c2aff 100644 (file)
--- a/folly/futures/SharedPromise.h
+++ b/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_;
 };
 
 }

diff --git a/folly/futures/test/SharedPromiseTest.cpp b/folly/futures/test/SharedPromiseTest.cpp
index 1dd23b85bc46c5a9ee6b010fa55940c379e771b3..9ecef42320e88ed4c6bc0893678a81f4a9fdf0db 100644 (file)
--- a/folly/futures/test/SharedPromiseTest.cpp
+++ b/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);
+}