Propagate exceptions thrown by parallel_scheduler bulk functions

include/stdexec/__detail/__parallel_scheduler.hpp

@@ -27,6 +27,7 @@ import stdexec;
 #  include "__execution_fwd.hpp"
 
 #  include "../functional.hpp"
+#  include "__atomic.hpp"
 #  include "__bulk.hpp"
 #  include "__domain.hpp"
 #  include "__manual_lifetime.hpp"
@@ -39,6 +40,7 @@ import stdexec;
 #  include "__variant.hpp"
 
 #  if !STDEXEC_USE_MODULES()
+#    include <exception>
 #    include <optional>
 #    include <utility>
 #  endif
@@ -374,6 +376,13 @@ namespace STDEXEC
       using __rcvr_t  = _BulkState::__rcvr_t;
       using __base_t  = __forward_args_receiver<_Previous, env_of_t<__rcvr_t>>;
       using __tuple_t = __decayed_tuple<_As...>;
+      using __fun_t   = _BulkState::__fun_t;
+
+      //! Whether calling the bulk functor can throw.
+      static constexpr bool __nothrow_fun =
+        _BulkState::__is_unchunked
+          ? __nothrow_callable<__fun_t&, size_t, __decay_t<_As>&...>
+          : __nothrow_callable<__fun_t&, size_t, size_t, __decay_t<_As>&...>;
 
       //! Stores `__as` in the base class storage, with the right types.
       explicit __typed_forward_args_receiver(_As&&... __as)
@@ -387,7 +396,15 @@ namespace STDEXEC
       void set_value() noexcept override
       {
         auto __state = reinterpret_cast<_BulkState*>(this);
-        auto __args  = std::move(__base_t::__arguments_.__get());
+        if constexpr (!__nothrow_fun)
+        {
+          if (__state->__exception_)
+          {
+            __typed_forward_args_receiver::set_error(std::move(__state->__exception_));
+            return;
+          }
+        }
+        auto __args = std::move(__base_t::__arguments_.__get());
         __base_t::__arguments_.__destroy();
         std::destroy_at(this);
         auto const __fun = STDEXEC::__flip(std::bind_front(__apply, STDEXEC::set_value));
@@ -416,6 +433,32 @@ namespace STDEXEC
       //! Calls the bulk functor passing `__index` and the values from the
       //! previous sender.
       void execute(size_t __begin, size_t __end) noexcept override
+      {
+        if constexpr (__nothrow_fun)
+        {
+          __execute(__begin, __end);
+        }
+        else
+        {
+          STDEXEC_TRY
+          {
+            __execute(__begin, __end);
+          }
+          STDEXEC_CATCH_ALL
+          {
+            // Only the first exception is kept; it is reported from
+            // `set_value()` once all the bulk work has completed.
+            auto __state = reinterpret_cast<_BulkState*>(this);
+            if (!__state->__has_exception_.exchange(true, __std::memory_order_relaxed))
+            {
+              __state->__exception_ = std::current_exception();
+            }
+          }
+        }
+      }
+
+     private:
+      void __execute(size_t __begin, size_t __end) noexcept(__nothrow_fun)
       {
         auto __state = reinterpret_cast<_BulkState*>(this);
         if constexpr (_BulkState::__is_unchunked)
@@ -487,6 +530,7 @@ namespace STDEXEC
     struct __bulk_state_base
     {
       using __rcvr_t                       = _Rcvr;
+      using __fun_t                        = _Fn;
       using __forward_args_helper_t        = __forward_args_receiver<_Previous, env_of_t<__rcvr_t>>;
       static constexpr bool __is_unchunked = _IsUnchunked;
       static constexpr bool __parallelize  = _Parallelize;
@@ -512,6 +556,11 @@ namespace STDEXEC
       //! The size of the bulk operation.
       _Size __size_;
 
+      //! Set by the first bulk item call that throws.
+      __std::atomic<bool> __has_exception_{false};
+      //! The exception thrown by the first bulk item call that throws.
+      std::exception_ptr __exception_;
+
       __bulk_state_base(_Fn&& __fun, _Rcvr&& __rcvr, _Size __size)
         : __fun_{std::move(__fun)}
         , __rcvr_{std::move(__rcvr)}

test/stdexec/schedulers/test_parallel_scheduler.cpp

@@ -688,6 +688,79 @@ TEST_CASE("bulk on parallel_scheduler destroys stored predecessor values after e
 
   CHECK(live->load(std::memory_order_relaxed) == 0);
 }
+
+TEST_CASE("bulk on parallel_scheduler propagates exceptions thrown by the function",
+          "[scheduler][parallel_scheduler]")
+{
+  auto live = std::make_shared<std::atomic<int>>(0);
+
+  {
+    auto fn = [](std::size_t i, destructor_tracked_value&)
+    {
+      if (i == 7)
+        throw std::runtime_error{"bulk"};
+    };
+    STDEXEC::parallel_scheduler sched = STDEXEC::get_parallel_scheduler();
+    auto                        snd = tracked_value_sender{sched, live} | ex::bulk(ex::par, 16, fn);
+
+    CHECK_THROWS_AS(ex::sync_wait(std::move(snd)), std::runtime_error);
+  }
+
+  CHECK(live->load(std::memory_order_relaxed) == 0);
+}
+
+TEST_CASE("bulk_chunked on parallel_scheduler propagates exceptions thrown by the function",
+          "[scheduler][parallel_scheduler]")
+{
+  STDEXEC::parallel_scheduler sched = STDEXEC::get_parallel_scheduler();
+  std::atomic<int>            calls = 0;
+
+  auto par_snd = ex::bulk_chunked(ex::schedule(sched),
+                                  ex::par,
+                                  1'000,
+                                  [&](std::size_t, std::size_t)
+                                  {
+                                    ++calls;
+                                    throw std::runtime_error{"bulk_chunked"};
+                                  });
+  CHECK_THROWS_AS(ex::sync_wait(std::move(par_snd)), std::runtime_error);
+  CHECK(calls.load() >= 1);
+
+  auto seq_snd = ex::bulk_chunked(ex::schedule(sched),
+                                  ex::seq,
+                                  1'000,
+                                  [](std::size_t, std::size_t)
+                                  { throw std::runtime_error{"bulk_chunked"}; });
+  CHECK_THROWS_AS(ex::sync_wait(std::move(seq_snd)), std::runtime_error);
+}
+
+TEST_CASE("bulk_unchunked on parallel_scheduler propagates exceptions thrown by the function",
+          "[scheduler][parallel_scheduler]")
+{
+  STDEXEC::parallel_scheduler sched = STDEXEC::get_parallel_scheduler();
+
+  auto par_snd = ex::bulk_unchunked(ex::schedule(sched),
+                                    ex::par,
+                                    64,
+                                    [](std::size_t i)
+                                    {
+                                      if (i % 2 == 0)
+                                        throw std::runtime_error{"bulk_unchunked"};
+                                    });
+  CHECK_THROWS_AS(ex::sync_wait(std::move(par_snd)), std::runtime_error);
+
+  // With `seq`, the iterations after the throwing one are not executed.
+  std::size_t executed = 0;
+  auto        seq_fn   = [&](std::size_t i)
+  {
+    ++executed;
+    if (i == 3)
+      throw std::runtime_error{"bulk_unchunked"};
+  };
+  auto seq_snd = ex::bulk_unchunked(ex::schedule(sched), ex::seq, 64, seq_fn);
+  CHECK_THROWS_AS(ex::sync_wait(std::move(seq_snd)), std::runtime_error);
+  CHECK(executed == 4);
+}
 #endif
 
 TEST_CASE("bulk on parallel_scheduler destroys stored predecessor values after stopped",