Skip to content

Commit f3ac8e1

Browse files
akxclaude
andcommitted
gh-156310: Make the iter() sequence fallback iterator safe in free-threaded build
Sharing a single PySeqIter between threads could double-DECREF the underlying sequence and use it after free. Apply the same approach as listiter/tupleiter/reversed (gh-120608): use relaxed atomics for it_index with -1 as the exhaustion sentinel, and in the free-threaded build keep the reference to the sequence until the iterator is deallocated. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 53d2e14 commit f3ac8e1

3 files changed

Lines changed: 89 additions & 11 deletions

File tree

Lib/test/test_free_threading/test_iteration.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
1+
import sys
12
import threading
23
import unittest
34
from test import support
5+
from test.support import threading_helper
46

57
# The race conditions these tests were written for only happen every now and
68
# then, even with the current numbers. To find rare race conditions, bumping
@@ -112,6 +114,59 @@ def worker():
112114
self.assert_iterator_results(results, list(seq))
113115

114116

117+
class ContendedSeqIterExhaustionTest(unittest.TestCase):
118+
"""Test draining a shared iter() fallback iterator (PySeqIter_Type).
119+
120+
Sequences implementing __getitem__ but not __iter__ iterate through
121+
PySeqIter_Type. Unlike the other tests in this file, this uses a
122+
tiny sequence and many rounds so that many threads reach the racy
123+
exhaustion path simultaneously (see gh-156310, where this
124+
use-after-freed the sequence).
125+
"""
126+
127+
class Seq:
128+
def __init__(self, n):
129+
self.n = n
130+
131+
def __getitem__(self, i):
132+
if i >= self.n:
133+
raise IndexError(i)
134+
return i
135+
136+
def test_shared_iterator_exhaustion(self):
137+
nthreads = 8
138+
nrounds = 20 if support.check_sanitizer(thread=True) else 100
139+
seq = self.Seq(4)
140+
expected = set(range(seq.n))
141+
refcount_before = sys.getrefcount(seq)
142+
143+
def drain(it, barrier, results):
144+
items = []
145+
barrier.wait()
146+
for item in it:
147+
items.append(item)
148+
results.extend(items)
149+
150+
for _ in range(nrounds):
151+
it = iter(seq)
152+
barrier = threading.Barrier(nthreads)
153+
results = []
154+
threads = [
155+
threading.Thread(target=drain, args=(it, barrier, results))
156+
for _ in range(nthreads)
157+
]
158+
with threading_helper.start_threads(threads):
159+
pass
160+
del it
161+
# Threads may see duplicate or missing items, but never
162+
# invented ones.
163+
self.assertEqual(set(results) - expected, set())
164+
165+
# A double-DECREF of the sequence does not always crash; it
166+
# reliably shows up as a sagging reference count.
167+
self.assertEqual(sys.getrefcount(seq), refcount_before)
168+
169+
115170
class ContendedRangeIterationTest(ContendedTupleIterationTest):
116171
def make_testdata(self, n):
117172
return range(n)
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
Fix a use-after-free of the underlying sequence when a single :func:`iter`
2+
fallback iterator for objects implementing :meth:`~object.__getitem__`
3+
without :meth:`~object.__iter__` (``PySeqIter_Type``) was shared between
4+
threads in the free-threaded build. Such iterators are still not
5+
thread-safe in the sense that concurrent iteration may see duplicate or
6+
missing items, but they no longer corrupt the interpreter state.

Objects/iterobject.c

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,14 @@
66
#include "pycore_ceval.h" // _PyEval_GetBuiltin()
77
#include "pycore_genobject.h" // _PyCoro_GetAwaitableIter()
88
#include "pycore_object.h" // _PyObject_GC_TRACK()
9+
#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_SSIZE_RELAXED()
910

1011

1112
typedef struct {
1213
PyObject_HEAD
13-
Py_ssize_t it_index;
14-
PyObject *it_seq; /* Set to NULL when iterator is exhausted */
14+
Py_ssize_t it_index; /* -1 when iterator is exhausted */
15+
PyObject *it_seq; /* Set to NULL when iterator is exhausted
16+
(in the default build) */
1517
} seqiterobject;
1618

1719
PyObject *
@@ -58,26 +60,34 @@ iter_iternext(PyObject *iterator)
5860

5961
assert(PySeqIter_Check(iterator));
6062
it = (seqiterobject *)iterator;
63+
Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(it->it_index);
64+
if (index < 0)
65+
return NULL;
6166
seq = it->it_seq;
67+
#ifndef Py_GIL_DISABLED
6268
if (seq == NULL)
6369
return NULL;
64-
if (it->it_index == PY_SSIZE_T_MAX) {
70+
#endif
71+
if (index == PY_SSIZE_T_MAX) {
6572
PyErr_SetString(PyExc_OverflowError,
6673
"iter index too large");
6774
return NULL;
6875
}
6976

70-
result = PySequence_GetItem(seq, it->it_index);
77+
result = PySequence_GetItem(seq, index);
7178
if (result != NULL) {
72-
it->it_index++;
79+
FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, index + 1);
7380
return result;
7481
}
7582
if (PyErr_ExceptionMatches(PyExc_IndexError) ||
7683
PyErr_ExceptionMatches(PyExc_StopIteration))
7784
{
7885
PyErr_Clear();
86+
FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, -1);
87+
#ifndef Py_GIL_DISABLED
7988
it->it_seq = NULL;
8089
Py_DECREF(seq);
90+
#endif
8191
}
8292
return NULL;
8393
}
@@ -88,7 +98,8 @@ iter_len(PyObject *op, PyObject *Py_UNUSED(ignored))
8898
seqiterobject *it = (seqiterobject*)op;
8999
Py_ssize_t seqsize, len;
90100

91-
if (it->it_seq) {
101+
Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(it->it_index);
102+
if (index >= 0 && it->it_seq != NULL) {
92103
if (_PyObject_HasLen(it->it_seq)) {
93104
seqsize = PySequence_Size(it->it_seq);
94105
if (seqsize == -1)
@@ -97,7 +108,7 @@ iter_len(PyObject *op, PyObject *Py_UNUSED(ignored))
97108
else {
98109
Py_RETURN_NOTIMPLEMENTED;
99110
}
100-
len = seqsize - it->it_index;
111+
len = seqsize - index;
101112
if (len >= 0)
102113
return PyLong_FromSsize_t(len);
103114
}
@@ -116,8 +127,9 @@ iter_reduce(PyObject *op, PyObject *Py_UNUSED(ignored))
116127
* call must be before access of iterator pointers.
117128
* see issue #101765 */
118129

119-
if (it->it_seq != NULL)
120-
return Py_BuildValue("N(O)n", iter, it->it_seq, it->it_index);
130+
Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(it->it_index);
131+
if (index >= 0 && it->it_seq != NULL)
132+
return Py_BuildValue("N(O)n", iter, it->it_seq, index);
121133
else
122134
return Py_BuildValue("N(())", iter);
123135
}
@@ -131,10 +143,15 @@ iter_setstate(PyObject *op, PyObject *state)
131143
Py_ssize_t index = PyLong_AsSsize_t(state);
132144
if (index == -1 && PyErr_Occurred())
133145
return NULL;
134-
if (it->it_seq != NULL) {
146+
/* An exhausted iterator keeps its reference to the sequence in the
147+
* free-threaded build, but must not be revived, matching the
148+
* default build where the reference is already gone. See gh-120971. */
149+
if (it->it_seq != NULL
150+
&& FT_ATOMIC_LOAD_SSIZE_RELAXED(it->it_index) >= 0)
151+
{
135152
if (index < 0)
136153
index = 0;
137-
it->it_index = index;
154+
FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, index);
138155
}
139156
Py_RETURN_NONE;
140157
}

0 commit comments

Comments
 (0)