gh-123271: Make builtin zip method safe under free-threading

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);
 

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()

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.

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);