Fix the allocator for Flow Graph critical tasks creating (#1596)

Co-authored-by: Aleksei Fedotov <[email protected]>
Co-authored-by: Mike Voss <[email protected]>

Author: 3 people

include/oneapi/tbb/detail/_flow_graph_impl.h

@@ -1,5 +1,5 @@
 /*
-    Copyright (c) 2005-2024 Intel Corporation
+    Copyright (c) 2005-2025 Intel Corporation
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -526,7 +526,10 @@ inline graph_task* prioritize_task(graph& g, graph_task& gt) {
     //! priority queue, and a new critical task is created to take and execute a work item with
     //! the highest known priority. The reference counting responsibility is transferred to
     //! the new task.
-    d1::task* critical_task = gt.my_allocator.new_object<priority_task_selector>(g.my_priority_queue, gt.my_allocator);
+    // A newly created small_object_allocator should be used to allocate the priority_task_selector
+    // instead of the allocator, associated with gt since gt can be allocated by another thread
+    d1::small_object_allocator allocator;
+    d1::task* critical_task = allocator.new_object<priority_task_selector>(g.my_priority_queue, allocator);
     __TBB_ASSERT( critical_task, "bad_alloc?" );
     g.my_priority_queue.push(&gt);
     using tbb::detail::d1::submit;

src/tbb/small_object_pool.cpp

@@ -1,5 +1,5 @@
 /*
-    Copyright (c) 2020-2021 Intel Corporation
+    Copyright (c) 2020-2025 Intel Corporation
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -45,6 +45,8 @@ void* __TBB_EXPORTED_FUNC allocate(d1::small_object_pool*& allocator, std::size_
 
 void* small_object_pool_impl::allocate_impl(d1::small_object_pool*& allocator, std::size_t number_of_bytes)
 {
+    __TBB_ASSERT(allocator == nullptr || allocator == this,
+                 "An attempt was made to allocate using another thread's small memory pool");
     small_object* obj{nullptr};
 
     if (number_of_bytes <= small_object_size) {

test/tbb/test_function_node.cpp

@@ -1,5 +1,5 @@
 /*
-    Copyright (c) 2005-2024 Intel Corporation
+    Copyright (c) 2005-2025 Intel Corporation
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -806,3 +806,25 @@ TEST_CASE("test function_node try_put_and_wait") {
     test_try_put_and_wait();
 }
 #endif
+
+// It was an issue when the critical task wrapper was allocated using the small object pool
+// of the task being wrapped. Since the original task creates under the aggregator, there is no
+// guarantee that the thread that requested the task creating is the same as actually created the task
+// Mismatch between memory pool caused internal assertion failure while deallocating the task
+//! \brief \ref regression
+TEST_CASE("test critical tasks memory pool correctness") {
+    using node_type = tbb::flow::function_node<int, tbb::flow::continue_msg>;
+    constexpr int num_iterations = 10000;
+    int num_calls = 0;
+    auto node_body = [&](int) { ++num_calls; };
+
+    tbb::flow::graph g;
+    node_type node(g, tbb::flow::serial, node_body, tbb::flow::node_priority_t{1});
+
+    for (int i = 0; i < num_iterations; ++i) {
+        node.try_put(i);
+    }
+
+    g.wait_for_all();
+    REQUIRE_MESSAGE(num_calls == num_iterations, "Incorrect number of body executions");
+}