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
55 changes: 55 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,61 @@ def read_set():
for t in threads:
t.join()

@threading_helper.reap_threads
def test_length_hint_used_race(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 = 50
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.reap_threads
def test_iternext_concurrent_exhaust_race(self):
Copy link
Contributor

Choose a reason for hiding this comment

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

The test_iternext_concurrent_exhaust_race and test_length_hint_exhaust_race test slightly difference things, but I think they can be combined (to reduce pressure on the CI)

NUM_ITERS = 200
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
NUM_ITERS = 200
NUM_ITERS = 10

Reduce this to keep computation time low, same for some of the other parameters. TSAN will catch most threading issues even with a low number of iterations/loops.

NUM_THREADS = 10

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

def worker():
while True:
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.
65 changes: 56 additions & 9 deletions Objects/setobject.c
Original file line number Diff line number Diff line change
Expand Up @@ -1056,8 +1056,23 @@ 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
Copy link
Contributor

Choose a reason for hiding this comment

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

This might work for setiter_len, but setiter_iternext itself is not yet thread safe (also because of setting si->si_set to zero).

For several other iterations the approach is to keep the reference si->si_set , but use another attribute to signal exhaustion of the iterator. For example for itertools.cycle or the reversed operator.

Note: I tried creating a minimal example where concurrent iteration fails, but I have succeeded yet (the example does not crash, although I have not run thread sanitizer on it yet)

Test for concurrent iteration on set iterator
import unittest
from threading import Thread, Barrier


class TestSetIter(unittest.TestCase):
    def test_set_iter(self):
        """Test concurrent iteration over a set"""

        NUM_LOOPS = 10_000
        NUM_THREADS = 4
        

        for ii in range(NUM_LOOPS):
            if ii % 1000 ==0:
                print(f'test_set_iter {ii}')
            barrier = Barrier(NUM_THREADS)
            
            # make sure the underlying set is unique referenced by the iterator
            iterator = iter(set((1,2,))) 
            
            def worker():
                barrier.wait()
                while True:
                    iterator.__length_hint__()
                    try:
                        next(iterator)
                    except StopIteration:
                        break

                
            threads = [Thread(target=worker) for _ in range(NUM_THREADS)]
            for t in threads:
                t.start()
            for t in threads:
                t.join()
                
            assert iterator.__length_hint__()==0

if __name__ == "__main__":
    unittest.main()

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thank you. I think your points make a lot of sense, and I really appreciate the two links you shared—they helped me get a more complete picture of the iterator-related data race.
I’ll try to construct the case you mentioned under a TSan environment.
If it turns out to be appropriate, we can address it fully in this PR, that would be great. Of course, this will take some time.

Copy link
Contributor

Choose a reason for hiding this comment

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

Yes, we should fix this like we have fixed others and as Sam suggested only clear the associated set in non-free-threading builds. The current code is incorrect because it uses try incref which can fail spuriously if the set object is not marked to enable try incref.

Copy link
Contributor

Choose a reason for hiding this comment

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

If we also set si->si_pos to -1 in the non-free threading build, then we can keep the code for the FT and non-FT the same I think. For the non-FT build we then set both si->si)set to zero and si->si_pos to -1 at exhaustion, but the code will be simpler.

Also the two code paths in setiter_iternext will be more similar then.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thanks — using a flag makes the code clearer.

PySetObject *so = si->si_set;
if (so != NULL) {
Copy link
Contributor

Choose a reason for hiding this comment

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

In the FT build so cannot be NULL. Remove the if (maybe replace with an assert).

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

return PyLong_FromSsize_t(len);
}

Expand Down Expand Up @@ -1089,17 +1104,22 @@ 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;
#ifndef Py_GIL_DISABLED
int decref_so = 0;
#endif

if (so == NULL)
if (so == NULL) {
return NULL;
assert (PyAnySet_Check(so));
}
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--;
}
Py_END_CRITICAL_SECTION();
si->si_pos = i+1;
if (key == NULL) {
else {
/* exhausted */
si->si_pos = -1;
si->len = 0;
#ifndef Py_GIL_DISABLED
si->si_set = NULL;
decref_so = 1;
Copy link
Contributor

Choose a reason for hiding this comment

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

We can do the Py_DECREF(so); here and then remove the decref_so variable. (we are exhausted, so key is still NULL).

#endif
}

#ifdef Py_GIL_DISABLED
done:
Py_END_CRITICAL_SECTION2();
return key;
#else
Py_END_CRITICAL_SECTION();

if (decref_so) {
Py_DECREF(so);
return NULL;
}
si->len--;
return key;
#endif
}

PyTypeObject PySetIter_Type = {
Expand Down
Loading