[3.13] gh-117657: Fix itertools.count thread safety (GH-119268) by colesbury · Pull Request #120007 · python/cpython · GitHub
Skip to content

[3.13] gh-117657: Fix itertools.count thread safety (GH-119268) #120007

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
colesbury merged 2 commits into from
Jun 3, 2024
+54 −11
Merged

[3.13] gh-117657: Fix itertools.count thread safety (GH-119268) #120007

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5576e64
[3.13] gh-117657: Fix itertools.count thread safety (GH-119268)
wiggin15 May 21, 2024
9170561
Merge branch '3.13' into backport-87939bd-3.13
colesbury Jun 3, 2024
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
24 changes: 23 additions & 1 deletion Lib/test/test_itertools.py
View file
Open in desktop
40 changes: 31 additions & 9 deletions Modules/itertoolsmodule.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
#include "Python.h"
#include "pycore_call.h" // _PyObject_CallNoArgs()
#include "pycore_ceval.h" // _PyEval_GetBuiltin()
#include "pycore_long.h" // _PyLong_GetZero()
#include "pycore_moduleobject.h" // _PyModule_GetState()
#include "pycore_typeobject.h" // _PyType_GetModuleState()
#include "pycore_object.h" // _PyObject_GC_TRACK()
#include "pycore_tuple.h" // _PyTuple_ITEMS()
#include "pycore_call.h" // _PyObject_CallNoArgs()
#include "pycore_ceval.h" // _PyEval_GetBuiltin()
#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
#include "pycore_long.h" // _PyLong_GetZero()
#include "pycore_moduleobject.h" // _PyModule_GetState()
#include "pycore_typeobject.h" // _PyType_GetModuleState()
#include "pycore_object.h" // _PyObject_GC_TRACK()
#include "pycore_tuple.h" // _PyTuple_ITEMS()

#include <stddef.h> // offsetof()
#include <stddef.h> // offsetof()

/* Itertools module written and maintained
by Raymond D. Hettinger <python@rcn.com>
Expand Down Expand Up @@ -4037,7 +4038,7 @@ fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.

assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Advances with: cnt += 1
When count hits Y_SSIZE_T_MAX, switch to slow_mode.
When count hits PY_SSIZE_T_MAX, switch to slow_mode.

slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.

Expand Down Expand Up @@ -4186,9 +4187,30 @@ count_nextlong(countobject *lz)
static PyObject *
count_next(countobject *lz)
{
#ifndef Py_GIL_DISABLED
if (lz->cnt == PY_SSIZE_T_MAX)
return count_nextlong(lz);
return PyLong_FromSsize_t(lz->cnt++);
#else
// free-threading version
// fast mode uses compare-exchange loop
// slow mode uses a critical section
PyObject *returned;
Py_ssize_t cnt;

cnt = _Py_atomic_load_ssize_relaxed(&lz->cnt);
for (;;) {
if (cnt == PY_SSIZE_T_MAX) {
Py_BEGIN_CRITICAL_SECTION(lz);
returned = count_nextlong(lz);
Py_END_CRITICAL_SECTION();
return returned;
}
if (_Py_atomic_compare_exchange_ssize(&lz->cnt, &cnt, cnt + 1)) {
return PyLong_FromSsize_t(cnt);
}
}
#endif
}

static PyObject *
Expand Down
1 change: 0 additions & 1 deletion Tools/tsan/suppressions_free_threading.txt
View file
Open in desktop