Make await exception safe

Summary: This fixes fibers::await to correctly handle exception thrown by user-passed lambda. Await still always waits for the promise to be fulfilled (if the promise was not moved out, it will be destroyed and thus auto-fulfilled with "promise not fulfilled" exception). However if user-passed lambda throws, promise result is ignored (even if exception) and exception thrown by lambda is re-thrown.

Reviewed By: pavlo-fb

Differential Revision: D3303393

fbshipit-source-id: c4ba01fde0e156cc88e5d07aaf763e3abe121d11

folly/experimental/fibers/FiberManager-inl.h

@@ -547,12 +547,23 @@ typename FirstArgOf<F>::type::value_type inline await(F&& func) {
   typedef typename FirstArgOf<F>::type::value_type Result;
 
   folly::Try<Result> result;
+  std::exception_ptr funcException;
 
   Baton baton;
-  baton.wait([&func, &result, &baton]() mutable {
-    func(Promise<Result>(result, baton));
+  baton.wait([&func, &result, &baton, &funcException]() mutable {
+    try {
+      func(Promise<Result>(result, baton));
+    } catch (...) {
+      // Save the exception, but still wait for baton to be posted by user code
+      // or promise destructor.
+      funcException = std::current_exception();
+    }
   });
 
+  if (UNLIKELY(funcException != nullptr)) {
+    std::rethrow_exception(funcException);
+  }
+
   return folly::moveFromTry(result);
 }
 }

folly/experimental/fibers/test/FibersTest.cpp

@@ -319,6 +319,31 @@ TEST(FiberManager, addTasksNoncopyable) {
   loopController.loop(std::move(loopFunc));
 }
 
+TEST(FiberManager, awaitThrow) {
+  folly::EventBase evb;
+  struct ExpectedException {};
+  getFiberManager(evb)
+      .addTaskFuture([&] {
+        EXPECT_THROW(
+          await([](Promise<int> p) {
+              p.setValue(42);
+              throw ExpectedException();
+            }),
+          ExpectedException
+        );
+
+        EXPECT_THROW(
+          await([&](Promise<int> p) {
+              evb.runInEventBaseThread([p = std::move(p)]() mutable {
+                  p.setValue(42);
+                });
+              throw ExpectedException();
+            }),
+          ExpectedException);
+      })
+      .waitVia(&evb);
+}
+
 TEST(FiberManager, addTasksThrow) {
   std::vector<Promise<int>> pendingFibers;
   bool taskAdded = false;