22
33#include < algorithm>
44#include < cstddef>
5+ #include < exception>
6+ #include < mutex>
7+ #include < stdexcept>
58#include < thread>
69#include < utility>
710#include < vector>
@@ -27,34 +30,66 @@ inline std::size_t normalize_thread_count(
2730
2831// Split [0, number_of_work_items) into n_threads contiguous chunks and invoke
2932// `chunk(thread_id, begin, end)` on each chunk. Thread 0 runs on the calling
30- // thread; threads 1..n_threads-1 run in parallel std::threads . The caller is
33+ // thread; threads 1..n_threads-1 run in parallel std::jthreads . The caller is
3134// responsible for picking n_threads via normalize_thread_count and for the
3235// thread safety of `chunk`.
36+ //
37+ // Exceptions thrown by `chunk` are captured rather than allowed to escape a
38+ // worker thread (which would call std::terminate). The first exception on any
39+ // thread is stored and rethrown on the calling thread after every worker has
40+ // been joined, so nanobind can translate it into a Python exception.
3341template <class Chunk >
3442void parallel_for_chunks (
3543 const std::size_t n_threads,
3644 const std::size_t number_of_work_items,
3745 Chunk &&chunk
3846) {
47+ if (n_threads == 0 ) {
48+ throw std::invalid_argument (" parallel_for_chunks requires n_threads >= 1" );
49+ }
3950 const auto bounds = [&](const std::size_t thread_id) {
4051 const auto begin = thread_id * number_of_work_items / n_threads;
4152 const auto end = (thread_id + 1 ) * number_of_work_items / n_threads;
4253 return std::pair<std::size_t , std::size_t >{begin, end};
4354 };
4455
45- std::vector<std::thread> threads;
56+ std::exception_ptr first_exception;
57+ std::mutex exception_mutex;
58+ const auto guarded_chunk = [&] (
59+ const std::size_t thread_id,
60+ const std::size_t begin,
61+ const std::size_t end
62+ ) {
63+ try {
64+ chunk (thread_id, begin, end);
65+ } catch (...) {
66+ const std::lock_guard<std::mutex> lock (exception_mutex);
67+ if (!first_exception) {
68+ first_exception = std::current_exception ();
69+ }
70+ }
71+ };
72+
73+ // jthread's destructor joins, so a failure while creating a later worker
74+ // cannot destroy an earlier still-joinable thread and terminate the
75+ // process. We still join explicitly below to rethrow only after all work
76+ // has completed.
77+ std::vector<std::jthread> threads;
4678 threads.reserve (n_threads > 0 ? n_threads - 1 : 0 );
4779 for (std::size_t thread_id = 1 ; thread_id < n_threads; ++thread_id) {
4880 const auto [begin, end] = bounds (thread_id);
49- threads.emplace_back ([thread_id, begin, end, &chunk ]() {
50- chunk (thread_id, begin, end);
81+ threads.emplace_back ([thread_id, begin, end, &guarded_chunk ]() {
82+ guarded_chunk (thread_id, begin, end);
5183 });
5284 }
5385 const auto [begin, end] = bounds (0 );
54- chunk (0 , begin, end);
86+ guarded_chunk (0 , begin, end);
5587 for (auto &thread : threads) {
5688 thread.join ();
5789 }
90+ if (first_exception) {
91+ std::rethrow_exception (first_exception);
92+ }
5893}
5994
6095} // namespace bioimage_cpp::detail
0 commit comments