gh-144356: Avoid races when computing set_iterator.__length_hint__ under no-gil

[hyongtao-code]

Long log:

setiter_len() was reading so->used without atomic access while concurrent
mutations update it atomically under Py_GIL_DISABLED.

In free-threaded builds, setiter_len() could race with concurrent set
mutation and iterator exhaustion.

Use an atomic load for so->used to avoid a data race. This preserves the
existing semantics of __length_hint__ while making the access thread-safe.

Signed-off-by: Yongtao Huang yongtaoh2022@gmail.com

• Issue: Data race in set iterator length_hint under no-gil #144356

[hyongtao-code]

Thanks for the review. I’ve decided to address both issues in this PR. I also added a corresponding test case for the issue you pointed out.

[eendebakpt]

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()

[hyongtao-code]

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.

[kumaraditya303]

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.

[colesbury]

I think we should follow the pattern that we use in other iterators: don't clear si->si_set when the iterator is exhausted in the free-threaded build.

That will keep other things simpler.

[eendebakpt]

I left some more review comments. I think we can get this right, but a simpler approach here would be to put a critical section on the set iterator itself.

[eendebakpt]

Suggested change

	#ifdef Py_GIL_DISABLED
	FT_ATOMIC_STORE_SSIZE_RELAXED(si->si_pos, i + 1);
	#else
	si->si_pos = i + 1;
	#endif
	FT_ATOMIC_STORE_SSIZE_RELAXED(si->si_pos, i + 1);

On the normal build the macro will expand to si->si_pos = i + 1;

[eendebakpt]

This (and some other places) should also be atomic?

[eendebakpt]

The value -1 written here could be overwritten by a concurrent thread (at line 1155). Which means that over exhaustion of the set iterator it is restored back to life. This does not lead to overflows or other issues (afaic), but is a bit odd behaviour.

[eendebakpt]

I think in the normal build you still have to do si->si_set = NULL;, otherwise the si->si_set is decref'ed again in setiter_dealloc.

[hyongtao-code]

put a critical section on the set iterator

Thanks for the suggestion. I simplified the implementation by protecting both the iterator and the set with Py_BEGIN_CRITICAL_SECTION2(self, so), so concurrent use of the same iterator is properly serialized. This keeps the logic aligned with other iterators and avoids the races.

It looks like there is an unrelated flaky test case.

[eendebakpt]

@hyongtao-code Thanks for your work on this! The core idea of the PR has shifted a bit, which is why I have some more review comments, but we are getting closer I believe.

[eendebakpt]

Since we are in a critical section on so the atomic load here is not required.

[eendebakpt]

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.

[eendebakpt]

I suspect the reader here will be finished much faster than the writer. (the length hint seems way cheaper than the add and discard). There are several ways to make this more balanced. Here is a suggestion, but there might be better ones.

Suggested change

	i += 1
	def reader():
	barrier.wait()
	while s:
	it.__length_hint__()
	def writer():
	barrier.wait()
	i = 0
	for _ in range(NUM_LOOPS):
	s.add(i)
	s.discard(i - 1)
	i += 1
	s.clear()

[hyongtao-code]

Sorry for the delayed response. I will check them one by one. Really, thanks for your guidance.