gh-123619: Add an unstable C API function for enabling deferred reference counting (GH-123635)

Co-authored-by: Sam Gross <colesbury@gmail.com>

1 files changed, 46 insertions, 0 deletions

diff --git a/Lib/test/test_capi/test_object.py b/Lib/test/test_capi/test_object.py
index cc9c9b6..a38b203 100644
--- a/Lib/test/test_capi/test_object.py
+++ b/Lib/test/test_capi/test_object.py
@@ -1,10 +1,13 @@
 import enum
 import unittest
+from test import support
 from test.support import import_helper
 from test.support import os_helper
+from test.support import threading_helper
 
 _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
 _testcapi = import_helper.import_module('_testcapi')
+_testinternalcapi = import_helper.import_module('_testinternalcapi')
 
 
 class Constant(enum.IntEnum):
@@ -131,5 +134,48 @@ class ClearWeakRefsNoCallbacksTest(unittest.TestCase):
         _testcapi.pyobject_clear_weakrefs_no_callbacks(obj)
 
 
+class EnableDeferredRefcountingTest(unittest.TestCase):
+    """Test PyUnstable_Object_EnableDeferredRefcount"""
+    @support.requires_resource("cpu")
+    def test_enable_deferred_refcount(self):
+        from threading import Thread
+
+        self.assertEqual(_testcapi.pyobject_enable_deferred_refcount("not tracked"), 0)
+        foo = []
+        self.assertEqual(_testcapi.pyobject_enable_deferred_refcount(foo), int(support.Py_GIL_DISABLED))
+
+        # Make sure reference counting works on foo now
+        self.assertEqual(foo, [])
+        if support.Py_GIL_DISABLED:
+            self.assertTrue(_testinternalcapi.has_deferred_refcount(foo))
+
+        # Make sure that PyUnstable_Object_EnableDeferredRefcount is thread safe
+        def silly_func(obj):
+            self.assertIn(
+                _testcapi.pyobject_enable_deferred_refcount(obj),
+                (0, 1)
+            )
+
+        silly_list = [1, 2, 3]
+        threads = [
+            Thread(target=silly_func, args=(silly_list,)) for _ in range(5)
+        ]
+
+        with threading_helper.catch_threading_exception() as cm:
+            for t in threads:
+                t.start()
+
+            for i in range(10):
+                silly_list.append(i)
+
+            for t in threads:
+                t.join()
+
+            self.assertIsNone(cm.exc_value)
+
+        if support.Py_GIL_DISABLED:
+            self.assertTrue(_testinternalcapi.has_deferred_refcount(silly_list))
+
+
 if __name__ == "__main__":
     unittest.main()