From 2c0ece4a278ae289abd4b08e01d1613bf326c1e7 Mon Sep 17 00:00:00 2001
From: Shubhanshu Agrawal <shubhanshu@fb.com>
Date: Fri, 14 Oct 2016 06:43:42 -0700
Subject: [PATCH] BatchDispatcher exception handling

Summary: error handling for batch dispatcher, for the case where dispatch function throws.

Reviewed By: yfeldblum

Differential Revision: D4016246

fbshipit-source-id: ac61b873a425128be7839a119591fca10608210a
---
 folly/fibers/BatchDispatcher.h   | 26 +++++++++++++++++++-------
 folly/fibers/test/FibersTest.cpp | 25 +++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/folly/fibers/BatchDispatcher.h b/folly/fibers/BatchDispatcher.h
index dada4275..65940ea6 100644
--- a/folly/fibers/BatchDispatcher.h
+++ b/folly/fibers/BatchDispatcher.h
@@ -105,13 +105,25 @@ class BatchDispatcher {
     state.values.swap(values);
     state.promises.swap(promises);
 
-    auto results = state.dispatchFunc(std::move(values));
-    if (results.size() != promises.size()) {
-      throw std::logic_error(
-          "Unexpected number of results returned from dispatch function");
-    }
-    for (size_t i = 0; i < promises.size(); i++) {
-      promises[i].setValue(std::move(results[i]));
+    try {
+      auto results = state.dispatchFunc(std::move(values));
+      if (results.size() != promises.size()) {
+        throw std::logic_error(
+            "Unexpected number of results returned from dispatch function");
+      }
+
+      for (size_t i = 0; i < promises.size(); i++) {
+        promises[i].setValue(std::move(results[i]));
+      }
+    } catch (const std::exception& ex) {
+      for (size_t i = 0; i < promises.size(); i++) {
+        promises[i].setException(
+            exception_wrapper(std::current_exception(), ex));
+      }
+    } catch (...) {
+      for (size_t i = 0; i < promises.size(); i++) {
+        promises[i].setException(exception_wrapper(std::current_exception()));
+      }
     }
   }
 
diff --git a/folly/fibers/test/FibersTest.cpp b/folly/fibers/test/FibersTest.cpp
index f5218dd2..eeee0b2e 100644
--- a/folly/fibers/test/FibersTest.cpp
+++ b/folly/fibers/test/FibersTest.cpp
@@ -1728,6 +1728,31 @@ TEST(FiberManager, doubleBatchDispatchTest) {
   evb.loop();
 }
 
+template <typename ExecutorT>
+void batchDispatchExceptionHandling(ExecutorT& executor, int i) {
+  thread_local BatchDispatcher<int, int, ExecutorT> batchDispatcher(
+      executor, [=, &executor](std::vector<int> &&) -> std::vector<int> {
+        throw std::runtime_error("Surprise!!");
+      });
+
+  EXPECT_THROW(batchDispatcher.add(i).get(), std::runtime_error);
+}
+
+TEST(FiberManager, batchDispatchExceptionHandlingTest) {
+  folly::EventBase evb;
+  auto& executor = getFiberManager(evb);
+
+  // Launch multiple fibers with a single id.
+  executor.add([&]() {
+    int totalNumberOfElements = 5;
+    for (int i = 0; i < totalNumberOfElements; i++) {
+      executor.add(
+          [=, &executor]() { batchDispatchExceptionHandling(executor, i); });
+    }
+  });
+  evb.loop();
+}
+
 /**
  * Test that we can properly track fiber stack usage.
  *
-- 
2.34.1