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..456f1150a209376 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3675,27 +3675,34 @@ static PyObject * count_repr(PyObject *op) { countobject *lz = countobject_CAST(op); - if (lz->long_cnt == NULL) { + PyObject *long_cnt; + + Py_BEGIN_CRITICAL_SECTION(lz); + long_cnt = Py_XNewRef(lz->long_cnt); + Py_END_CRITICAL_SECTION(); + + if (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); } + long step = 0; if (PyLong_Check(lz->long_step)) { - long step = PyLong_AsLong(lz->long_step); + step = PyLong_AsLong(lz->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); - } } - return PyUnicode_FromFormat("%s(%R, %R)", - _PyType_Name(Py_TYPE(lz)), - lz->long_cnt, lz->long_step); + + /* Don't display step when it is an integer equal to 1 */ + PyObject *msg = (step == 1) + ? PyUnicode_FromFormat("%s(%R)", _PyType_Name(Py_TYPE(lz)), long_cnt) + : PyUnicode_FromFormat("%s(%R, %R)", _PyType_Name(Py_TYPE(lz)), + long_cnt, lz->long_step); + + Py_DECREF(long_cnt); + return msg; } static PyType_Slot count_slots[] = {