Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions Lib/test/test_free_threading/test_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,43 @@ def read_set():
for t in threads:
t.join()

@threading_helper.reap_threads
def test_length_hint_with_mutating_set(self):
NUM_ITERS = 10
NUM_THREADS = 10
NUM_LOOPS = 2_000

for _ in range(NUM_ITERS):
s = set(range(2000))
it = iter(s)

def worker():
for i in range(NUM_LOOPS):
it.__length_hint__()
s.add(i)
s.discard(i - 1)

threading_helper.run_concurrently(worker, nthreads=NUM_THREADS)

@threading_helper.reap_threads
def test_length_hint_exhaust_race(self):
NUM_ITERS = 10
NUM_THREADS = 10

for _ in range(NUM_ITERS):
s = set(range(256))
it = iter(s)

def worker():
while True:
it.__length_hint__()
try:
next(it)
except StopIteration:
break

threading_helper.run_concurrently(worker, nthreads=NUM_THREADS)


@threading_helper.requires_working_threading()
class SmallSetTest(RaceTestBase, unittest.TestCase):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix potential races in set iterators (``__length_hint__`` and iteration) in free-threaded builds.
63 changes: 55 additions & 8 deletions Objects/setobject.c
Original file line number Diff line number Diff line change
Expand Up @@ -1056,8 +1056,22 @@ setiter_len(PyObject *op, PyObject *Py_UNUSED(ignored))
{
setiterobject *si = (setiterobject*)op;
Py_ssize_t len = 0;
if (si->si_set != NULL && si->si_used == si->si_set->used)

#ifdef Py_GIL_DISABLED
PySetObject *so = si->si_set;
assert(so != NULL);

Py_BEGIN_CRITICAL_SECTION2(op, so);
if (si->si_pos >= 0 && si->si_used == so->used) {
len = si->len;
}
Py_END_CRITICAL_SECTION2();
#else
if (si->si_set != NULL && si->si_used == si->si_set->used) {
len = si->len;
}
#endif

Comment on lines +1060 to +1074
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
#ifdef Py_GIL_DISABLED
PySetObject *so = si->si_set;
assert(so != NULL);
Py_BEGIN_CRITICAL_SECTION2(op, so);
if (si->si_pos >= 0 && si->si_used == so->used) {
len = si->len;
}
Py_END_CRITICAL_SECTION2();
#else
if (si->si_set != NULL && si->si_used == si->si_set->used) {
len = si->len;
}
#endif
PySetObject *so = si->si_set;
Py_BEGIN_CRITICAL_SECTION(op);
if (si->si_pos >= 0 && si->si_used == FT_ATOMIC_LOAD_SSIZE_RELAXED(so->used)) {
len = si->len;
}
Py_END_CRITICAL_SECTION();

This is less code and a bit more efficient (only a single lock):

From the set object so = si->si_setwe only need to read the fieldusedwhich we can do with an atomic load. In the FT buildsois never NULL. In the non-FT buildsocan only be NULL ifsi->si_set < 0`)

@hyongtao-code I did not test the above suggestion

return PyLong_FromSsize_t(len);
}

Expand Down Expand Up @@ -1089,17 +1103,23 @@ static PyMethodDef setiter_methods[] = {
{NULL, NULL} /* sentinel */
};

static PyObject *setiter_iternext(PyObject *self)
static PyObject *
setiter_iternext(PyObject *self)
{
setiterobject *si = (setiterobject*)self;
PyObject *key = NULL;
Py_ssize_t i, mask;
setentry *entry;
PySetObject *so = si->si_set;

if (so == NULL)
#ifdef Py_GIL_DISABLED
assert(so != NULL);
#else
if (so == NULL) {
return NULL;
assert (PyAnySet_Check(so));
}
#endif
assert(PyAnySet_Check(so));

Py_ssize_t so_used = FT_ATOMIC_LOAD_SSIZE_RELAXED(so->used);
Py_ssize_t si_used = FT_ATOMIC_LOAD_SSIZE_RELAXED(si->si_used);
Expand All @@ -1110,26 +1130,53 @@ static PyObject *setiter_iternext(PyObject *self)
return NULL;
}

#ifdef Py_GIL_DISABLED
Py_BEGIN_CRITICAL_SECTION2(self, so);
#else
Py_BEGIN_CRITICAL_SECTION(so);
#endif

i = si->si_pos;
assert(i>=0);
#ifdef Py_GIL_DISABLED
if (i < 0) {
/* iterator already exhausted */
goto done;
}
#endif

entry = so->table;
mask = so->mask;
while (i <= mask && (entry[i].key == NULL || entry[i].key == dummy)) {
i++;
}
if (i <= mask) {
key = Py_NewRef(entry[i].key);
si->si_pos = i + 1;
si->len--;
}
else {
/* exhausted */
si->si_pos = -1;
si->len = 0;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this needed?

#ifndef Py_GIL_DISABLED
si->si_set = NULL;
#endif
}

#ifdef Py_GIL_DISABLED
done:
Py_END_CRITICAL_SECTION2();
return key;
#else
Py_END_CRITICAL_SECTION();
si->si_pos = i+1;

if (key == NULL) {
si->si_set = NULL;
/* exhausted */
Py_DECREF(so);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you move the decref to the section

       /* exhausted */
        si->si_pos = -1;
        si->len = 0;

?

return NULL;
}
si->len--;
return key;
#endif
}

PyTypeObject PySetIter_Type = {
Expand Down
Loading