From 7e38e6745d2f9ee235d934ab7f3c6b3085be2b70 Mon Sep 17 00:00:00 2001
From: Pieter Eendebak <pieter.eendebak@gmail.com>
Date: Tue, 27 Aug 2024 21:22:43 +0200
Subject: gh-123271: Make builtin zip method safe under free-threading
 (#123272)

The `zip_next` function uses a common optimization technique for methods
that generate tuples. The iterator maintains an internal reference to
the returned tuple. When the method is called again, it checks if the
internal tuple's reference count is 1. If so, the tuple can be reused.
However, this approach is not safe under the free-threading build:
after checking the reference count, another thread may perform the same
check and also reuse the tuple. This can result in a double decref on
the items of the replaced tuple and a double incref (memory leak) on
the items of the tuple being set.

This adds a function, `_PyObject_IsUniquelyReferenced` that
encapsulates the stricter logic necessary for the free-threaded build:
the internal tuple must be owned by the current thread, have a local
refcount of one, and a shared refcount of zero.
---
 Include/internal/pycore_object.h                   | 17 +++++++++
 Lib/test/test_free_threading/test_zip.py           | 41 ++++++++++++++++++++++
 .../2024-08-23-21-20-34.gh-issue-123271.xeVViR.rst |  1 +
 Python/bltinmodule.c                               |  3 +-
 4 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 Lib/test/test_free_threading/test_zip.py
 create mode 100644 Misc/NEWS.d/next/Core_and_Builtins/2024-08-23-21-20-34.gh-issue-123271.xeVViR.rst

diff --git a/Include/internal/pycore_object.h b/Include/internal/pycore_object.h
index 0f2de6f..2ef0aae 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 0000000..b4c5837
--- /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 0000000..51fdec4
--- /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 99ed069..f87f942 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);
-- 
cgit v0.12