Make task_group::wait more thread safe (#1780)

Co-authored-by: Aleksei Fedotov <aleksei.fedotov@intel.com>
Co-authored-by: Konstantin Boyarinov <konstantin.boyarinov@intel.com>
Co-authored-by: Alexey Kukanov <alexey.kukanov@intel.com>

Author: 4 people

include/oneapi/tbb/task_group.h

@@ -58,6 +58,7 @@ template <bool>
 class context_guard_helper;
 struct task_arena_impl;
 class context_list;
+void handle_context_exception(d1::task_group_context& ctx, bool rethrow);
 
 TBB_EXPORT void __TBB_EXPORTED_FUNC execute(d1::task_arena_base&, d1::delegate_base&);
 TBB_EXPORT void __TBB_EXPORTED_FUNC isolate_within_arena(d1::delegate_base&, std::intptr_t);
@@ -448,6 +449,7 @@ class task_group_context : no_copy {
     friend struct r1::task_arena_impl;
     friend struct r1::task_group_context_impl;
     friend class d2::task_group_base;
+    friend void r1::handle_context_exception(d1::task_group_context&, bool rethrow);
 }; // class task_group_context
 
 static_assert(sizeof(task_group_context) == 128, "Wrong size of task_group_context");

src/tbb/arena.cpp

@@ -1,5 +1,6 @@
 /*
     Copyright (c) 2005-2025 Intel Corporation
+    Copyright (c) 2025 UXL Foundation Contributors
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -844,11 +845,7 @@ void task_arena_impl::execute(d1::task_arena_base& ta, d1::delegate_base& d) {
                 a->my_exit_monitors.notify_one(); // do not relax!
             }
             // process possible exception
-            auto exception = exec_context.my_exception.load(std::memory_order_acquire);
-            if (exception) {
-                __TBB_ASSERT(exec_context.is_group_execution_cancelled(), "The task group context with an exception should be canceled.");
-                exception->throw_self();
-            }
+            handle_context_exception(exec_context);
             __TBB_ASSERT(governor::is_thread_data_set(td), nullptr);
             return;
         } // if (index1 == arena::out_of_arena)

src/tbb/scheduler_common.h

@@ -336,13 +336,32 @@ class tbb_exception_ptr {
     /** Note that objects of this type can be created only by the allocate() method. **/
     void destroy() noexcept;
 
-    //! Throws the contained exception .
-    void throw_self();
+    //! Throws the contained exception and then destroys this object.
+    void rethrow_and_destroy();
 
 private:
     tbb_exception_ptr(const std::exception_ptr& src) : my_ptr(src) {}
 }; // class tbb_exception_ptr
 
+//! Utility function to atomically clear and handle exceptions from task_group_context
+/** This function implements thread-safe pattern for clearing exceptions
+    from a task_group_context and either destroying or throwing them. **/
+inline void handle_context_exception(d1::task_group_context& ctx, bool rethrow = true) {
+
+    tbb_exception_ptr* exception = ctx.my_exception.load(std::memory_order_acquire);
+    if (exception) {
+        if (ctx.my_exception.compare_exchange_strong(exception, nullptr, 
+                                                     std::memory_order_acq_rel)) {
+            // TODO: An exception should not be captured and then not rethrown.
+            //       Either add asserts or remove corner cases.
+            if (rethrow) {
+                exception->rethrow_and_destroy();
+            } else
+                exception->destroy();
+        }
+    }
+}
+
 //------------------------------------------------------------------------
 // Debugging support
 //------------------------------------------------------------------------

src/tbb/task_dispatcher.cpp

@@ -1,5 +1,6 @@
 /*
     Copyright (c) 2020-2025 Intel Corporation
+    Copyright (c) 2025 UXL Foundation Contributors
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -173,11 +174,7 @@ void task_dispatcher::execute_and_wait(d1::task* t, d1::wait_context& wait_ctx,
         local_td.m_thread_data->my_inbox.set_is_idle(false);
     }
 
-    auto exception = w_ctx.my_exception.load(std::memory_order_acquire);
-    if (exception) {
-        __TBB_ASSERT(w_ctx.is_group_execution_cancelled(), "The task group context with an exception should be canceled.");
-        exception->throw_self();
-    }
+    handle_context_exception(w_ctx);
 }
 
 #if __TBB_RESUMABLE_TASKS

src/tbb/task_dispatcher.h

@@ -380,9 +380,15 @@ d1::task* task_dispatcher::local_wait_for_all(d1::task* t, Waiter& waiter ) {
             if (global_control::active_value(global_control::terminate_on_exception) == 1) {
                 do_throw_noexcept([] { throw; });
             }
-            if (ed.context->cancel_group_execution()) {
-                /* We are the first to signal cancellation, so store the exception that caused it. */
-                ed.context->my_exception.store(tbb_exception_ptr::allocate(), std::memory_order_release);
+
+            ed.context->cancel_group_execution();
+            tbb_exception_ptr* exception = ed.context->my_exception.load(std::memory_order_acquire);
+            if (!exception) {
+                auto e = tbb_exception_ptr::allocate();
+                if (!ed.context->my_exception.compare_exchange_strong(exception, e,
+                                                                      std::memory_order_acq_rel)) {
+                    e->destroy();
+                }
             }
         }
     } // Infinite exception loop

src/tbb/task_group_context.cpp

@@ -1,5 +1,6 @@
 /*
-    Copyright (c) 2005-2023 Intel Corporation
+    Copyright (c) 2005-2025 Intel Corporation
+    Copyright (c) 2025 UXL Foundation Contributors
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -42,9 +43,11 @@ void tbb_exception_ptr::destroy() noexcept {
     deallocate_memory(this);
 }
 
-void tbb_exception_ptr::throw_self() {
+void tbb_exception_ptr::rethrow_and_destroy() {
+    auto temp_ptr = my_ptr;
+    destroy();
     if (governor::rethrow_exception_broken()) fix_broken_rethrow();
-    std::rethrow_exception(my_ptr);
+    std::rethrow_exception(temp_ptr);
 }
 
 //------------------------------------------------------------------------
@@ -65,10 +68,7 @@ void task_group_context_impl::destroy(d1::task_group_context& ctx) {
 #endif
     ctl->~cpu_ctl_env();
 
-    auto exception = ctx.my_exception.load(std::memory_order_relaxed);
-    if (exception) {
-        exception->destroy();
-    }
+    handle_context_exception(ctx, /* rethrow = */ false);
     ITT_STACK_DESTROY(ctx.my_itt_caller);
 
     poison_pointer(ctx.my_parent);
@@ -240,19 +240,14 @@ bool task_group_context_impl::is_group_execution_cancelled(const d1::task_group_
     return ctx.my_cancellation_requested.load(std::memory_order_relaxed) != 0;
 }
 
-// IMPORTANT: It is assumed that this method is not used concurrently!
+// IMPORTANT: If used while tasks are in the context, the cancellation signal can be lost
 void task_group_context_impl::reset(d1::task_group_context& ctx) {
     __TBB_ASSERT(!is_poisoned(ctx.my_context_list), nullptr);
     //! TODO: Add assertion that this context does not have children
-    // No fences are necessary since this context can be accessed from another thread
-    // only after stealing happened (which means necessary fences were used).
+    
+    handle_context_exception(ctx, /* rethrow = */ false);
 
-    auto exception = ctx.my_exception.load(std::memory_order_relaxed);
-    if (exception) {
-        exception->destroy();
-        ctx.my_exception.store(nullptr, std::memory_order_relaxed);
-    }
-    ctx.my_cancellation_requested = 0;
+    ctx.my_cancellation_requested.store(0, std::memory_order_relaxed);
 }
 
 // IMPORTANT: It is assumed that this method is not used concurrently!