gh-157841: Synchronize functools.partial state access

Lib/test/test_free_threading/test_functools.py

@@ -1,7 +1,7 @@
 import random
 import unittest
 
-from functools import lru_cache
+from functools import Placeholder, lru_cache, partial
 from threading import Barrier, Thread
 
 from test.support import threading_helper
@@ -71,5 +71,144 @@ def test_reentrant_cache_clear_bounded(self):
         self._test_reentrant_cache_clear(maxsize=128)
 
 
+def partial_target(x, y=0, *extra, z=0):
+    return x, y, z, True
+
+
+class PartialCallable:
+    def __call__(self, x, y=0, *extra, z=0):
+        return x, y, z, False
+
+
+@threading_helper.requires_working_threading()
+class TestPartial(unittest.TestCase):
+    def _race_setstate(self, p, read, *, placeholders=False, readers=4):
+        iterations = 2000
+        barrier = Barrier(readers + 2)
+        errors = []
+        callable_object = PartialCallable()
+
+        def reader():
+            try:
+                barrier.wait()
+                for _ in range(iterations):
+                    read(p)
+            except Exception as exc:
+                errors.append(exc)
+
+        def writer(offset):
+            try:
+                barrier.wait()
+                for i in range(iterations):
+                    value = 10**12 + offset * iterations + i
+                    if placeholders and i % 2:
+                        args = (Placeholder, value)
+                    else:
+                        args = (value, value)
+                    func = partial_target if i % 2 else callable_object
+                    p.__setstate__((func, args, {"z": value}, {"tag": value}))
+            except Exception as exc:
+                errors.append(exc)
+
+        threads = [Thread(target=reader) for _ in range(readers)]
+        threads += [Thread(target=writer, args=(i,)) for i in range(2)]
+        with threading_helper.start_threads(threads):
+            pass
+        self.assertEqual(errors, [])
+
+    def _partial(self):
+        p = partial(PartialCallable(), 0, 0, z=0)
+        p.tag = 0
+        return p
+
+    def test_concurrent_setstate(self):
+        p = self._partial()
+        self._race_setstate(p, None, readers=0)
+        x, y, z, marker = p()
+        self.assertEqual((x, x), (y, z))
+        self.assertEqual(marker, bool(y % 2))
+
+    def test_concurrent_setstate_and_repr(self):
+        self._race_setstate(self._partial(), repr)
+
+    def test_concurrent_setstate_and_call(self):
+        def read(p):
+            x, y, z, marker = p()
+            self.assertEqual((x, x), (y, z))
+            self.assertEqual(marker, bool(y % 2))
+
+        self._race_setstate(self._partial(), read)
+
+    def test_concurrent_setstate_and_placeholders(self):
+        def read(p):
+            x, y, z, marker = p(99)
+            self.assertEqual(y, z)
+            self.assertIn(x, (99, y))
+            self.assertEqual(marker, bool(y % 2))
+
+        self._race_setstate(self._partial(), read, placeholders=True)
+
+    def test_concurrent_setstate_and_reduce(self):
+        def read(p):
+            constructor_args, state = p.__reduce__()[1:]
+            func, args, kw, instance_dict = state
+            self.assertIs(constructor_args[0], func)
+            self.assertEqual(func is partial_target, bool(args[0] % 2))
+            self.assertEqual((args[0], args[0]), (args[1], kw["z"]))
+            self.assertEqual(args[0], instance_dict["tag"])
+
+        self._race_setstate(self._partial(), read)
+
+    def test_concurrent_setstate_and_flatten(self):
+        p = partial(partial_target, 0, 0, z=0)
+        iterations = 2000
+        barrier = Barrier(5)
+        errors = []
+
+        def writer():
+            try:
+                barrier.wait()
+                for i in range(iterations):
+                    value = 10**12 + i
+                    args = (Placeholder, value) if i % 2 else (value, value)
+                    p.__setstate__((partial_target, args, {"z": value}, None))
+            except Exception as exc:
+                errors.append(exc)
+
+        def reader():
+            try:
+                barrier.wait()
+                for _ in range(iterations):
+                    q = partial(p, 99)
+                    x, y, z, marker = q()
+                    self.assertIs(q.func, partial_target)
+                    self.assertEqual(y, z)
+                    self.assertIn(x, (99, y))
+                    self.assertTrue(marker)
+            except Exception as exc:
+                errors.append(exc)
+
+        threads = [Thread(target=reader) for _ in range(4)]
+        threads.append(Thread(target=writer))
+        with threading_helper.start_threads(threads):
+            pass
+        self.assertEqual(errors, [])
+
+    def test_concurrent_setstate_and_members(self):
+        def read(p):
+            self.assertTrue(callable(p.func))
+            self.assertEqual(len(p.args), 2)
+            self.assertIn("z", p.keywords)
+
+        self._race_setstate(self._partial(), read)
+
+    def test_concurrent_setstate_and_instance_dict(self):
+        def read(p):
+            self.assertIsInstance(p.tag, int)
+            self.assertIn("tag", p.__dict__)
+
+        self._race_setstate(self._partial(), read)
+
+
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_functools.py

@@ -619,6 +619,46 @@ class TestPartialC(TestPartial, unittest.TestCase):
         module = c_functools
         partial = c_functools.partial
 
+    def test_call_snapshot_reentrancy(self):
+        PH = self.module.Placeholder
+        p = self.partial(capture, PH, 2, k=0)
+
+        class EvilKey(str):
+            armed = True
+            __hash__ = str.__hash__
+
+            def __eq__(self, other):
+                if self.armed:
+                    self.armed = False
+                    p.__setstate__((capture, (PH, PH, 3), {}, None))
+                return str.__eq__(self, other)
+
+        key = EvilKey("k")
+        # The slot wrapper exercises tp_call, including its keyword merge.
+        self.assertEqual(p.__call__(1, **{key: 1}), ((1, 2), {"k": 1}))
+        self.assertFalse(key.armed)
+        self.assertEqual(p(10, 20), ((10, 20, 3), {}))
+
+    def test_setstate_finalizer_reentrancy(self):
+        p = self.partial(capture)
+        seen = []
+
+        class OldCallable:
+            def __call__(self, *args, **kw):
+                pass
+
+            def __del__(self):
+                seen.append((p.func, p.args, p.keywords.copy(),
+                             p.__dict__.copy()))
+                p.__setstate__((capture, (3,), {"a": 3}, {"tag": 3}))
+
+        p.__setstate__((OldCallable(), (1,), {"a": 1}, {"tag": 1}))
+        p.__setstate__((capture, (2,), {"a": 2}, {"tag": 2}))
+        gc.collect()
+        self.assertEqual(seen, [(capture, (2,), {"a": 2}, {"tag": 2})])
+        self.assertEqual(p(), ((3,), {"a": 3}))
+        self.assertEqual(p.tag, 3)
+
     def test_attributes_unwritable(self):
         # attributes should not be writable
         p = self.partial(capture, 1, 2, a=10, b=20)