Add utility for when_all that rethrows exceptions

Author: wence-

cpp/include/rapidsmpf/streaming/core/node.hpp

@@ -17,6 +17,20 @@ namespace rapidsmpf::streaming {
  */
 using Node = coro::task<void>;
 
+/**
+ * @brief Produce a new coroutine that waits for completion of all tasks.
+ *
+ * This function schedules each node and returns a new coroutine to
+ * await on after which all of them will have finished execution.
+ *
+ * @param nodes A vector of nodes to await.
+ * @throws If any of the underlying tasks throws an exception it is
+ * rethrown.
+ * @return Coroutine representing the completion of all of the tasks
+ * in `nodes`.
+ */
+Node when_all_or_throw(std::vector<Node>&& nodes);
+
 /**
  * @brief Runs a list of nodes concurrently and waits for all to complete.
  *

cpp/src/streaming/core/node.cpp

@@ -7,13 +7,17 @@
 
 namespace rapidsmpf::streaming {
 
-void run_streaming_pipeline(std::vector<Node> nodes) {
-    auto results = coro::sync_wait(coro::when_all(std::move(nodes)));
+Node when_all_or_throw(std::vector<Node>&& nodes) {
+    auto results = co_await coro::when_all(std::move(nodes));
+    // The node result itself is always `void` but we access it here to re-throw
+    // possible unhandled exceptions.
     for (auto& result : results) {
-        // The node result itself is always `void` but we access it here to re-throw
-        // possible unhandled exceptions.
         result.return_value();
     }
 }
 
+void run_streaming_pipeline(std::vector<Node> nodes) {
+    coro::sync_wait(when_all_or_throw(std::move(nodes)));
+}
+
 }  // namespace rapidsmpf::streaming

cpp/tests/streaming/test_leaf_node.cpp

@@ -78,10 +78,7 @@ Node shutdown(
     std::shared_ptr<Context> ctx, std::shared_ptr<Channel> ch, std::vector<Node>&& tasks
 ) {
     ShutdownAtExit c{ch};
-    auto results = co_await coro::when_all(std::move(tasks));
-    for (auto& r : results) {
-        r.return_value();
-    }
+    co_await when_all_or_throw(std::move(tasks));
     co_await ch->drain(ctx->executor());
 }