diff --git a/Lib/test/test_free_threading/test_itertools.py b/Lib/test/test_free_threading/test_itertools.py index 7392bd739acd709..afcf77c8d3bad3f 100644 --- a/Lib/test/test_free_threading/test_itertools.py +++ b/Lib/test/test_free_threading/test_itertools.py @@ -4,6 +4,7 @@ batched, chain, combinations_with_replacement, + count, cycle, permutations, tee, @@ -142,5 +143,32 @@ def consume(): self.assertEqual(errors, [], msg=f"unexpected errors: {errors}") +class TestCountConcurrent(unittest.TestCase): + @staticmethod + def _spin_next(it, n=2000): + for _ in range(n): + next(it) + + @staticmethod + def _spin_repr(it, n=2000): + for _ in range(n): + repr(it) + + @threading_helper.reap_threads + def test_repr_racing_next_fast_mode(self): + for _ in range(10): + it = count() + workers = [self._spin_next] * 2 + [self._spin_repr] * 4 + threading_helper.run_concurrently(workers, args=(it,)) + + @threading_helper.reap_threads + def test_repr_racing_next_slow_mode(self): + for _ in range(10): + # Large count to trigger "slow mode" + it = count(10**18, 2) + workers = [self._spin_next] * 2 + [self._spin_repr] * 4 + threading_helper.run_concurrently(workers, args=(it,)) + + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/Library/2026-07-18-14-47-52.gh-issue-153981.yZcmKQ.rst b/Misc/NEWS.d/next/Library/2026-07-18-14-47-52.gh-issue-153981.yZcmKQ.rst new file mode 100644 index 000000000000000..3fc3c3cca437c4b --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-18-14-47-52.gh-issue-153981.yZcmKQ.rst @@ -0,0 +1 @@ +Fix use-after-free race condition when calling :func:`repr` on :class:`itertools.count` under the :term:`free-threaded build`. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index c0023c839ca7fe3..c77bae8a1eaea50 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3675,27 +3675,51 @@ static PyObject * count_repr(PyObject *op) { countobject *lz = countobject_CAST(op); - if (lz->long_cnt == NULL) { - Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(lz->cnt); - return PyUnicode_FromFormat("%s(%zd)", - _PyType_Name(Py_TYPE(lz)), cnt); + PyObject *long_cnt, *long_step; + + Py_BEGIN_CRITICAL_SECTION(lz); + long_cnt = Py_XNewRef(lz->long_cnt); + long_step = Py_NewRef(lz->long_step); + Py_END_CRITICAL_SECTION(); + + if (long_cnt == NULL) { + /* Fast mode: cnt is advanced by count_next()'s lock-free atomic CAS, + which never takes the critical section above, so read it atomically. + Fast mode also implies that long_step is the integer 1. */ + long_cnt = PyLong_FromSsize_t(FT_ATOMIC_LOAD_SSIZE_RELAXED(lz->cnt)); } - if (PyLong_Check(lz->long_step)) { - long step = PyLong_AsLong(lz->long_step); + /* Don't display step when it is an integer equal to 1 */ + int hide_step = 0; + if (PyLong_Check(long_step)) { + long step = PyLong_AsLong(long_step); if (step == -1 && PyErr_Occurred()) { PyErr_Clear(); } - if (step == 1) { - /* Don't display step when it is an integer equal to 1 */ - return PyUnicode_FromFormat("%s(%R)", - _PyType_Name(Py_TYPE(lz)), - lz->long_cnt); - } + hide_step = (step == 1); + } + + PyUnicodeWriter *writer = PyUnicodeWriter_Create(0); + int rc = (writer == NULL || long_cnt == NULL) ? -1 : 0; + if (rc == 0) { + rc = PyUnicodeWriter_Format(writer, "%s(%R", + _PyType_Name(Py_TYPE(lz)), long_cnt); + } + if (rc == 0 && !hide_step) { + rc = PyUnicodeWriter_Format(writer, ", %R", long_step); + } + if (rc == 0) { + rc = PyUnicodeWriter_WriteChar(writer, ')'); + } + + Py_XDECREF(long_cnt); + Py_DECREF(long_step); + + if (rc < 0) { + PyUnicodeWriter_Discard(writer); + return NULL; } - return PyUnicode_FromFormat("%s(%R, %R)", - _PyType_Name(Py_TYPE(lz)), - lz->long_cnt, lz->long_step); + return PyUnicodeWriter_Finish(writer); } static PyType_Slot count_slots[] = {