pythongh-116738: Make _heapq module thread-safe

Author: yoney

Lib/test/test_free_threading/test_heapq.py

@@ -0,0 +1,115 @@
+import unittest
+
+import heapq
+
+from threading import Thread, Barrier
+from random import shuffle, randint
+from unittest import TestCase
+
+from test.support import threading_helper
+
+
+NTHREAD: int = 10
+OBJECT_COUNT: int = 5_000
+
+
+@threading_helper.requires_working_threading()
+class TestHeapq(TestCase):
+    def test_racing_heapify(self):
+        heap = list(range(OBJECT_COUNT))
+        shuffle(heap)
+
+        barrier = Barrier(NTHREAD)
+
+        def heapify_func(heap: list[int]):
+            barrier.wait()
+            heapq.heapify(heap)
+
+        self.run_parallel(heapify_func, (heap,), NTHREAD)
+        self.is_heap_property_satisfied(heap)
+
+    def test_racing_heappush(self):
+        heap = []
+
+        barrier = Barrier(NTHREAD)
+
+        def heappush_func(heap):
+            barrier.wait()
+            for item in reversed(range(OBJECT_COUNT)):
+                heapq.heappush(heap, item)
+
+        self.run_parallel(heappush_func, (heap,), NTHREAD)
+        self.is_heap_property_satisfied(heap)
+
+    def test_racing_heappop(self):
+        heap = list(range(OBJECT_COUNT))
+        shuffle(heap)
+        heapq.heapify(heap)
+
+        barrier = Barrier(NTHREAD)
+
+        # Each thread pops (OBJECT_COUNT / NTHREAD) items
+        self.assertEqual(0, OBJECT_COUNT % NTHREAD)
+        per_thread_pop_count = OBJECT_COUNT // NTHREAD
+
+        def heappop_func(heap, pop_count):
+            barrier.wait()
+            local_list = []
+            for _ in range(pop_count):
+                item = heapq.heappop(heap)
+                local_list.append(item)
+
+            # Each local list should be sorted
+            self.is_sorted(local_list)
+
+        self.run_parallel(heappop_func, (heap, per_thread_pop_count), NTHREAD)
+        self.assertEqual(0, len(heap))
+
+    def test_racing_heappushpop(self):
+        heap = list(range(OBJECT_COUNT))
+        shuffle(heap)
+
+        heapq.heapify(heap)
+
+        barrier = Barrier(NTHREAD)
+
+        pushpop_items = [randint(-OBJECT_COUNT, OBJECT_COUNT) for _ in range(OBJECT_COUNT)]
+        def heappushpop_func(heap, pushpop_items):
+            barrier.wait()
+            local_list = []
+            for item in pushpop_items:
+                popped_item = heapq.heappushpop(heap, item)
+                self.assertTrue(popped_item <= item)
+
+        self.run_parallel(heappushpop_func, (heap, pushpop_items), NTHREAD)
+        self.assertEqual(OBJECT_COUNT, len(heap))
+        self.is_heap_property_satisfied(heap)
+
+    def is_heap_property_satisfied(self, heap: list[object]):
+        # The value of a parent is always less than or equal to
+        # the value of its children.
+        # position 0 has no parent
+        for pos in range(1, len(heap)):
+            parent_pos = (pos - 1) >> 1
+            if heap[parent_pos] > heap[pos]:
+                return False
+
+        return True
+
+    def is_sorted(self, lst):
+        return all(lst[i - 1] <= lst[i] for i in range(1, len(lst)))
+
+    @staticmethod
+    def run_parallel(worker_func, args, nthreads):
+        workers = []
+        for _ in range(nthreads):
+            worker = Thread(target=worker_func, args=args)
+            workers.append(worker)
+            worker.start()
+
+        for worker in workers:
+            worker.join()
+
+
+if __name__ == "__main__":
+    unittest.main()

Modules/_heapqmodule.c

@@ -117,6 +117,7 @@ siftup(PyListObject *heap, Py_ssize_t pos)
 }
 
 /*[clinic input]
+@critical_section heap
 _heapq.heappush
 
     heap: object(subclass_of='&PyList_Type')
@@ -128,7 +129,7 @@ Push item onto heap, maintaining the heap invariant.
 
 static PyObject *
 _heapq_heappush_impl(PyObject *module, PyObject *heap, PyObject *item)
-/*[clinic end generated code: output=912c094f47663935 input=7c69611f3698aceb]*/
+/*[clinic end generated code: output=912c094f47663935 input=f7a4f03ef8d52e67]*/
 {
     if (PyList_Append(heap, item))
         return NULL;
@@ -171,6 +172,7 @@ heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
 }
 
 /*[clinic input]
+@critical_section heap
 _heapq.heappop
 
     heap: object(subclass_of='&PyList_Type')
@@ -181,7 +183,7 @@ Pop the smallest item off the heap, maintaining the heap invariant.
 
 static PyObject *
 _heapq_heappop_impl(PyObject *module, PyObject *heap)
-/*[clinic end generated code: output=96dfe82d37d9af76 input=91487987a583c856]*/
+/*[clinic end generated code: output=96dfe82d37d9af76 input=ed396461b153dd51]*/
 {
     return heappop_internal(heap, siftup);
 }
@@ -232,6 +234,7 @@ _heapq_heapreplace_impl(PyObject *module, PyObject *heap, PyObject *item)
 }
 
 /*[clinic input]
+@critical_section heap
 _heapq.heappushpop
 
     heap: object(subclass_of='&PyList_Type')
@@ -246,7 +249,7 @@ a separate call to heappop().
 
 static PyObject *
 _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item)
-/*[clinic end generated code: output=67231dc98ed5774f input=5dc701f1eb4a4aa7]*/
+/*[clinic end generated code: output=67231dc98ed5774f input=db05c81b1dd92c44]*/
 {
     PyObject *returnitem;
     int cmp;
@@ -371,6 +374,7 @@ heapify_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
 }
 
 /*[clinic input]
+@critical_section heap
 _heapq.heapify
 
     heap: object(subclass_of='&PyList_Type')
@@ -381,7 +385,7 @@ Transform list into a heap, in-place, in O(len(heap)) time.
 
 static PyObject *
 _heapq_heapify_impl(PyObject *module, PyObject *heap)
-/*[clinic end generated code: output=e63a636fcf83d6d0 input=53bb7a2166febb73]*/
+/*[clinic end generated code: output=e63a636fcf83d6d0 input=aaaaa028b9b6af08]*/
 {
     return heapify_internal(heap, siftup);
 }