map_reduce: prevent mapper or reducer exception from poisoning state

In map_reduce, if the mapper throws, we can end up with poisoned
state. This is because the compiler can choose to move s->result
before evaluating f.get(), so s->result gets broken (if the mapped
type becomes invalid after move).

Fix by checking for exceptions from the mapper before calling the reducer,
and exceptions from the reducer after calling the reducer, and storing them
in the state. If we see an exception, don't bother calling the reducer on
failed state. Since the reducer is a pure function (at least in expected
usage), and wouldn't have been called on mapper exceptions, we don't lose
anything by calling it.

A unit test is added for mapper exceptions and for reducer exceptions.
The tests do fail with clang before the patch, though if the arguments
are evaluated in reverse order, it can pass with a different compiler).

Author: avikivity

include/seastar/core/map_reduce.hh

@@ -185,23 +185,30 @@ map_reduce(Iterator begin, Iterator end, Mapper&& mapper, Initial initial, Reduc
         Mapper mapper;
         Initial result;
         Reduce reduce;
+        std::exception_ptr ex;
     };
-    auto s = make_lw_shared(state{std::forward<Mapper>(mapper), std::move(initial), std::move(reduce)});
+    auto s = make_lw_shared(state{std::forward<Mapper>(mapper), std::move(initial), std::move(reduce), std::exception_ptr()});
     future<> ret = make_ready_future<>();
     while (begin != end) {
         ret = futurize_invoke(s->mapper, *begin++).then_wrapped([s = s.get(), ret = std::move(ret)] (auto f) mutable {
-            try {
-                s->result = s->reduce(std::move(s->result), f.get());
-                return std::move(ret);
-            } catch (...) {
-                return std::move(ret).then_wrapped([ex = std::current_exception()] (auto f) {
-                    f.ignore_ready_future();
-                    return make_exception_future<>(ex);
-                });
+            if (s->ex) {
+                f.ignore_ready_future();  // We only report the first exception
+            } else if (f.failed()) {
+                s->ex = f.get_exception();
+            } else {
+                try {
+                    s->result = s->reduce(std::move(s->result), f.get());
+                } catch (...) {
+                    s->ex = std::current_exception();
+                }
             }
+            return std::move(ret);
         });
     }
     return ret.then([s] {
+        if (s->ex) {
+            return make_exception_future<Initial>(std::move(s->ex));
+        }
         return make_ready_future<Initial>(std::move(s->result));
     });
 }

tests/unit/futures_test.cc

@@ -862,6 +862,60 @@ SEASTAR_TEST_CASE(test_map_reduce1_lifetime) {
     });
 }
 
+SEASTAR_TEST_CASE(map_reduce_with_throwing_mapper) {
+    try {
+        auto vec = std::vector<int>{1, 2, 3, 4, 5, 6, 7};
+        auto ret = co_await map_reduce(
+            vec,
+            // Mapper: identity function, but throws
+            [] (int x) -> future<int> {
+                if (x == 5) {
+                    throw std::runtime_error("test");
+                }
+                co_return x;
+            },
+            // Initial value (and accumulator): move-only type
+            std::make_unique<int>(0),
+            // Reducer: test that it does not act on a moved-from value
+            [] (std::unique_ptr<int> acc, int x) -> std::unique_ptr<int> {
+                BOOST_REQUIRE(bool(acc));
+                *acc += x;
+                return acc;
+            }
+        );
+        BOOST_FAIL("should have thrown");
+    } catch (...) {
+        // Exception is expected and uninteresting
+    }
+}
+
+SEASTAR_TEST_CASE(map_reduce_with_throwing_reducer) {
+    try {
+        auto vec = std::vector<int>{1, 2, 3, 4, 5, 6, 7};
+        auto ret = co_await map_reduce(
+            vec,
+            // Mapper: square function
+            [] (int x) -> future<int> {
+                co_return x * x;
+            },
+            // Initial value (and accumulator): move-only type
+            std::make_unique<int>(0),
+            // Reducer: simple sum, but randomly throws
+            [] (std::unique_ptr<int> acc, int x) -> std::unique_ptr<int> {
+                BOOST_REQUIRE(bool(acc));
+                if (*acc > 14) {
+                    throw std::runtime_error("accumulator overflow, launch missiles");
+                }
+                *acc += x;
+                return acc;
+            }
+        );
+        BOOST_FAIL("should have thrown");
+    } catch (...) {
+        // Exception is expected and uninteresting
+    }
+}
+
 // This test doesn't actually test anything - it just waits for the future
 // returned by sleep to complete. However, a bug we had in sleep() caused
 // this test to fail the sanitizer in the debug build, so this is a useful