Add fix for memory pools mismatch for critical tasks

include/oneapi/tbb/detail/_flow_graph_impl.h

@@ -526,7 +526,8 @@ 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);
+    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

@@ -40,6 +40,7 @@ void* __TBB_EXPORTED_FUNC allocate(d1::small_object_pool*& allocator, std::size_
     // TODO: optimize if the allocator contains a valid pool.
     auto tls = governor::get_thread_data();
     auto pool = tls->my_small_object_pool;
+    __TBB_ASSERT(allocator == nullptr || pool == allocator, "An attempt was made to allocate using another thread's small memory pool");
     return pool->allocate_impl(allocator, number_of_bytes);
 }
 

test/tbb/test_function_node.cpp

@@ -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 pull 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");
+}