gh-135557: use atomic stores in heapq operations in free-threading (#135601)

This commit is contained in:
Xuanteng Huang 2025-06-21 16:43:15 +08:00 committed by GitHub
parent 8ca1e4d846
commit 13cac83347
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 51 additions and 17 deletions

View file

@ -3,7 +3,7 @@ import unittest
import heapq
from enum import Enum
from threading import Thread, Barrier
from threading import Thread, Barrier, Lock
from random import shuffle, randint
from test.support import threading_helper
@ -178,6 +178,33 @@ class TestHeapq(unittest.TestCase):
self.assertEqual(len(max_heap), OBJECT_COUNT)
self.test_heapq.check_max_invariant(max_heap)
def test_lock_free_list_read(self):
n, n_threads = 1_000, 10
l = []
barrier = Barrier(n_threads * 2)
count = 0
lock = Lock()
def worker():
with lock:
nonlocal count
x = count
count += 1
barrier.wait()
for i in range(n):
if x % 2:
heapq.heappush(l, 1)
heapq.heappop(l)
else:
try:
l[0]
except IndexError:
pass
self.run_concurrently(worker, (), n_threads * 2)
@staticmethod
def is_sorted_ascending(lst):
"""

View file

@ -0,0 +1,2 @@
Fix races on :mod:`heapq` updates and :class:`list` reads on the :term:`free threaded <free threading>`
build.

View file

@ -12,6 +12,7 @@ annotated by François Pinard, and converted to C by Raymond Hettinger.
#include "Python.h"
#include "pycore_list.h" // _PyList_ITEMS(), _PyList_AppendTakeRef()
#include "pycore_pyatomic_ft_wrappers.h"
#include "clinic/_heapqmodule.c.h"
@ -59,8 +60,8 @@ siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
arr = _PyList_ITEMS(heap);
parent = arr[parentpos];
newitem = arr[pos];
arr[parentpos] = newitem;
arr[pos] = parent;
FT_ATOMIC_STORE_PTR_RELAXED(arr[parentpos], newitem);
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], parent);
pos = parentpos;
}
return 0;
@ -108,8 +109,8 @@ siftup(PyListObject *heap, Py_ssize_t pos)
/* Move the smaller child up. */
tmp1 = arr[childpos];
tmp2 = arr[pos];
arr[childpos] = tmp2;
arr[pos] = tmp1;
FT_ATOMIC_STORE_PTR_RELAXED(arr[childpos], tmp2);
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], tmp1);
pos = childpos;
}
/* Bubble it up to its final resting place (by sifting its parents down). */
@ -172,8 +173,9 @@ heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
if (!n)
return lastelt;
returnitem = PyList_GET_ITEM(heap, 0);
PyList_SET_ITEM(heap, 0, lastelt);
if (siftup_func((PyListObject *)heap, 0)) {
PyListObject *list = _PyList_CAST(heap);
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], lastelt);
if (siftup_func(list, 0)) {
Py_DECREF(returnitem);
return NULL;
}
@ -208,8 +210,9 @@ heapreplace_internal(PyObject *heap, PyObject *item, int siftup_func(PyListObjec
}
returnitem = PyList_GET_ITEM(heap, 0);
PyList_SET_ITEM(heap, 0, Py_NewRef(item));
if (siftup_func((PyListObject *)heap, 0)) {
PyListObject *list = _PyList_CAST(heap);
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], Py_NewRef(item));
if (siftup_func(list, 0)) {
Py_DECREF(returnitem);
return NULL;
}
@ -284,8 +287,9 @@ _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item)
}
returnitem = PyList_GET_ITEM(heap, 0);
PyList_SET_ITEM(heap, 0, Py_NewRef(item));
if (siftup((PyListObject *)heap, 0)) {
PyListObject *list = _PyList_CAST(heap);
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], Py_NewRef(item));
if (siftup(list, 0)) {
Py_DECREF(returnitem);
return NULL;
}
@ -437,8 +441,8 @@ siftdown_max(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
arr = _PyList_ITEMS(heap);
parent = arr[parentpos];
newitem = arr[pos];
arr[parentpos] = newitem;
arr[pos] = parent;
FT_ATOMIC_STORE_PTR_RELAXED(arr[parentpos], newitem);
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], parent);
pos = parentpos;
}
return 0;
@ -486,8 +490,8 @@ siftup_max(PyListObject *heap, Py_ssize_t pos)
/* Move the smaller child up. */
tmp1 = arr[childpos];
tmp2 = arr[pos];
arr[childpos] = tmp2;
arr[pos] = tmp1;
FT_ATOMIC_STORE_PTR_RELAXED(arr[childpos], tmp2);
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], tmp1);
pos = childpos;
}
/* Bubble it up to its final resting place (by sifting its parents down). */
@ -621,8 +625,9 @@ _heapq_heappushpop_max_impl(PyObject *module, PyObject *heap, PyObject *item)
}
returnitem = PyList_GET_ITEM(heap, 0);
PyList_SET_ITEM(heap, 0, Py_NewRef(item));
if (siftup_max((PyListObject *)heap, 0) < 0) {
PyListObject *list = _PyList_CAST(heap);
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], Py_NewRef(item));
if (siftup_max(list, 0) < 0) {
Py_DECREF(returnitem);
return NULL;
}