diff --git a/Include/internal/pycore_object.h b/Include/internal/pycore_object.h index 0f2de6fd28f56e..2ef0aaee0f82b3 100644 --- a/Include/internal/pycore_object.h +++ b/Include/internal/pycore_object.h @@ -159,6 +159,23 @@ static inline void _Py_RefcntAdd(PyObject* op, Py_ssize_t n) } #define _Py_RefcntAdd(op, n) _Py_RefcntAdd(_PyObject_CAST(op), n) +// Checks if an object has a single, unique reference. If the caller holds a +// unique reference, it may be able to safely modify the object in-place. +static inline int +_PyObject_IsUniquelyReferenced(PyObject *ob) +{ +#if !defined(Py_GIL_DISABLED) + return Py_REFCNT(ob) == 1; +#else + // NOTE: the entire ob_ref_shared field must be zero, including flags, to + // ensure that other threads cannot concurrently create new references to + // this object. + return (_Py_IsOwnedByCurrentThread(ob) && + _Py_atomic_load_uint32_relaxed(&ob->ob_ref_local) == 1 && + _Py_atomic_load_ssize_relaxed(&ob->ob_ref_shared) == 0); +#endif +} + PyAPI_FUNC(void) _Py_SetImmortal(PyObject *op); PyAPI_FUNC(void) _Py_SetImmortalUntracked(PyObject *op); diff --git a/Lib/test/test_free_threading/test_zip.py b/Lib/test/test_free_threading/test_zip.py new file mode 100644 index 00000000000000..b4c5837b1d8936 --- /dev/null +++ b/Lib/test/test_free_threading/test_zip.py @@ -0,0 +1,41 @@ +import unittest +from threading import Thread + +from test.support import threading_helper + + +class ZipThreading(unittest.TestCase): + @staticmethod + def work(enum): + while True: + try: + next(enum) + except StopIteration: + break + + @threading_helper.reap_threads + @threading_helper.requires_working_threading() + def test_threading(self): + number_of_threads = 8 + number_of_iterations = 8 + n = 40_000 + enum = zip(range(n), range(n)) + for _ in range(number_of_iterations): + worker_threads = [] + for ii in range(number_of_threads): + worker_threads.append( + Thread( + target=self.work, + args=[ + enum, + ], + ) + ) + for t in worker_threads: + t.start() + for t in worker_threads: + t.join() + + +if __name__ == "__main__": + unittest.main() diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-08-23-21-20-34.gh-issue-123271.xeVViR.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-08-23-21-20-34.gh-issue-123271.xeVViR.rst new file mode 100644 index 00000000000000..51fdec452c1d41 --- /dev/null +++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-08-23-21-20-34.gh-issue-123271.xeVViR.rst @@ -0,0 +1 @@ +Make concurrent iterations over the same :func:`zip` iterator safe under free-threading. diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c index 99ed06972be98e..f87f942cc76258 100644 --- a/Python/bltinmodule.c +++ b/Python/bltinmodule.c @@ -2971,7 +2971,8 @@ zip_next(zipobject *lz) if (tuplesize == 0) return NULL; - if (Py_REFCNT(result) == 1) { + + if (_PyObject_IsUniquelyReferenced(result)) { Py_INCREF(result); for (i=0 ; i < tuplesize ; i++) { it = PyTuple_GET_ITEM(lz->ittuple, i);