8000 gh-135557: use atomic stores in `heapq` operations in free-threading … · python/cpython@13cac83 · GitHub
[go: up one dir, main page]

Skip to content
8000

Commit 13cac83

Browse files
authored
gh-135557: use atomic stores in heapq operations in free-threading (#135601)
1 parent 8ca1e4d commit 13cac83

File tree

3 files changed

+51
-17
lines changed

3 files changed

+51
-17
lines changed

Lib/test/test_free_threading/test_heapq.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import heapq
44

55
from enum import Enum
6-
from threading import Thread, Barrier
6+
from threading import Thread, Barrier, Lock
77
from random import shuffle, randint
88

99
from test.support import threading_helper
@@ -178,6 +178,33 @@ def heapreplace_max_func(max_heap, replace_items):
178178
self.assertEqual(len(max_heap), OBJECT_COUNT)
179179
self.test_heapq.check_max_invariant(max_heap)
180180

181+
def test_lock_free_list_read(self):
182+
n, n_threads = 1_000, 10
183+
l = []
184+
barrier = Barrier(n_threads * 2)
185+
186+
count = 0
187+
lock = Lock()
188+
189+
def worker():
190+
with lock:
191+
nonlocal count
192+
x = count
193+
count += 1
194+
195+
barrier.wait()
196+
for i in range(n):
197+
if x % 2:
198+
heapq.heappush(l, 1)
199+
heapq.heappop(l)
200+
else:
201+
try:
202+
l[0]
203+
except IndexError:
204+
pass
205+
206+
self.run_concurrently(worker, (), n_threads * 2)
207+
181208
@staticmethod
182209
def is_sorted_ascending(lst):
183210
"""
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix races on :mod:`heapq` updates and :class:`list` reads on the :term:`free threaded <free threading>`
2+
build.

Modules/_heapqmodule.c

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ annotated by François Pinard, and converted to C by Raymond Hettinger.
1212

1313
#include "Python.h"
1414
#include "pycore_list.h" // _PyList_ITEMS(), _PyList_AppendTakeRef()
15+
#include "pycore_pyatomic_ft_wrappers.h"
1516

1617
#include "clinic/_heapqmodule.c.h"
1718

@@ -59,8 +60,8 @@ siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
5960
arr = _PyList_ITEMS(heap);
6061
parent = arr[parentpos];
6162
newitem = arr[pos];
62-
arr[parentpos] = newitem;
63-
arr[pos] = parent;
63+
FT_ATOMIC_STORE_PTR_RELAXED(arr[parentpos], newitem);
64+
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], parent);
6465
pos = parentpos;
6566
}
6667
return 0;
@@ -108,8 +109,8 @@ siftup(PyListObject *heap, Py_ssize_t pos)
108109
/* Move the smaller child up. */
109110
tmp1 = arr[childpos];
110111
tmp2 = arr[pos];
111-
arr[childpos] = tmp2;
112-
arr[pos] = tmp1;
112+
FT_ATOMIC_STORE_PTR_RELAXED(arr[childpos], tmp2);
113+
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], tmp1);
113114
pos = childpos;
114115
}
115116
/* 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))
172173
if (!n)
173174
return lastelt;
174175
returnitem = PyList_GET_ITEM(heap, 0);
175-
PyList_SET_ITEM(heap, 0, lastelt);
176-
if (siftup_func((PyListObject *)heap, 0)) {
176+
PyListObject *list = _PyList_CAST(heap);
177+
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], lastelt);
178+
if (siftup_func(list, 0)) {
177179
Py_DECREF(returnitem);
178180
return NULL;
179181
}
@@ -208,8 +210,9 @@ heapreplace_internal(PyObject *heap, PyObject *item, int siftup_func(PyListObjec
208210
}
209211

210212
returnitem = PyList_GET_ITEM(heap, 0);
211-
PyList_SET_ITEM(heap, 0, Py_NewRef(item));
212-
if (siftup_func((PyListObject *)heap, 0)) {
213+
PyListObject *list = _PyList_CAST(heap);
214+
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], Py_NewRef(item));
215+
if (siftup_func(list, 0)) {
213216
Py_DECREF(returnitem);
214217
return NULL;
215218
}
@@ -284,8 +287,9 @@ _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item)
284287
}
285288

286289
returnitem = PyList_GET_ITEM(heap, 0);
287-
PyList_SET_ITEM(heap, 0, Py_NewRef(item));
288-
if (siftup((PyListObject *)heap, 0)) {
290+
PyListObject *list = _PyList_CAST(heap);
291+
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], Py_NewRef(item));
292+
if (siftup(list, 0)) {
289293
Py_DECREF(returnitem);
290294
return NULL;
291295
}
@@ -437,8 +441,8 @@ siftdown_max(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos)
437441
arr = _PyList_ITEMS(heap);
438442
parent = arr[parentpos];
439443
newitem = arr[pos];
440-
arr[parentpos] = newitem;
441-
arr[pos] = parent;
444+
FT_ATOMIC_STORE_PTR_RELAXED(arr[parentpos], newitem);
445+
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], parent);
442446
pos = parentpos;
443447
}
444448
return 0;
@@ -486,8 +490,8 @@ siftup_max(PyListObject *heap, Py_ssize_t pos)
486490
/* Move the smaller child up. */
487491
tmp1 = arr[childpos];
488492
tmp2 = arr[pos];
489-
arr[childpos] = tmp2;
490-
arr[pos] = tmp1;
493+
FT_ATOMIC_STORE_PTR_RELAXED(arr[childpos], tmp2);
494+
FT_ATOMIC_STORE_PTR_RELAXED(arr[pos], tmp1);
491495
pos = childpos;
492496
}
493497
/* 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)
621625
}
622626

623627
returnitem = PyList_GET_ITEM(heap, 0);
624-
PyList_SET_ITEM(heap, 0, Py_NewRef(item));
625-
if (siftup_max((PyListObject *)heap, 0) < 0) {
628+
PyListObject *list = _PyList_CAST(heap);
629+
FT_ATOMIC_STORE_PTR_RELAXED(list->ob_item[0], Py_NewRef(item));
630+
if (siftup_max(list, 0) < 0) {
626631
Py_DECREF(returnitem);
627632
return NULL;
628633
}

0 commit comments

Comments
 (0)
0