From 5c4a81a5cabc10b1970642cf0610e93c7b2a52b2 Mon Sep 17 00:00:00 2001
From: Sam Gross <colesbury@gmail.com>
Date: Fri, 12 Apr 2024 17:43:10 +0000
Subject: [PATCH] gh-117376: Make code objects use deferred reference counting

We want code objects to use deferred reference counting in the
free-threaded build. This requires them to be tracked by the GC, so we
set `Py_TPFLAGS_HAVE_GC` in the free-threaded build, but not the default
build.
---
 Lib/test/test_capi/test_watchers.py |  4 +++-
 Lib/test/test_gc.py                 |  4 +++-
 Objects/codeobject.c                | 33 ++++++++++++++++++++++++++++-
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/Lib/test/test_capi/test_watchers.py b/Lib/test/test_capi/test_watchers.py
index ae062b1bda26b7c..8e84d0077c7573a 100644
--- a/Lib/test/test_capi/test_watchers.py
+++ b/Lib/test/test_capi/test_watchers.py
@@ -1,7 +1,7 @@
 import unittest
 
 from contextlib import contextmanager, ExitStack
-from test.support import catch_unraisable_exception, import_helper
+from test.support import catch_unraisable_exception, import_helper, gc_collect
 
 
 # Skip this test if the _testcapi module isn't available.
@@ -372,6 +372,7 @@ def code_watcher(self, which_watcher):
 
     def assert_event_counts(self, exp_created_0, exp_destroyed_0,
                             exp_created_1, exp_destroyed_1):
+        gc_collect()  # code objects are collected by GC in free-threaded build
         self.assertEqual(
             exp_created_0, _testcapi.get_code_watcher_num_created_events(0))
         self.assertEqual(
@@ -432,6 +433,7 @@ def test_dealloc_error(self):
         with self.code_watcher(2):
             with catch_unraisable_exception() as cm:
                 del co
+                gc_collect()
 
                 self.assertEqual(str(cm.unraisable.exc_value), "boom!")
 
diff --git a/Lib/test/test_gc.py b/Lib/test/test_gc.py
index 71c7fb0edebaa50..52681dc18cfb867 100644
--- a/Lib/test/test_gc.py
+++ b/Lib/test/test_gc.py
@@ -226,7 +226,9 @@ def test_function(self):
         exec("def f(): pass\n", d)
         gc.collect()
         del d
-        self.assertEqual(gc.collect(), 2)
+        # In the free-threaded build, the count returned by `gc.collect()`
+        # is 3 because it includes f's code object.
+        self.assertIn(gc.collect(), (2, 3))
 
     def test_function_tp_clear_leaves_consistent_state(self):
         # https://github.com/python/cpython/issues/91636
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index 825f521b439235d..014632962bfcf36 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -6,6 +6,7 @@
 #include "pycore_code.h"          // _PyCodeConstructor
 #include "pycore_frame.h"         // FRAME_SPECIALS_SIZE
 #include "pycore_interp.h"        // PyInterpreterState.co_extra_freefuncs
+#include "pycore_object.h"        // _PyObject_SetDeferredRefcount
 #include "pycore_opcode_metadata.h" // _PyOpcode_Deopt, _PyOpcode_Caches
 #include "pycore_opcode_utils.h"  // RESUME_AT_FUNC_START
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
@@ -557,13 +558,22 @@ _PyCode_New(struct _PyCodeConstructor *con)
     }
 
     Py_ssize_t size = PyBytes_GET_SIZE(con->code) / sizeof(_Py_CODEUNIT);
-    PyCodeObject *co = PyObject_NewVar(PyCodeObject, &PyCode_Type, size);
+    PyCodeObject *co;
+#ifdef Py_GIL_DISABLED
+    co = PyObject_GC_NewVar(PyCodeObject, &PyCode_Type, size);
+#else
+    co = PyObject_NewVar(PyCodeObject, &PyCode_Type, size);
+#endif
     if (co == NULL) {
         Py_XDECREF(replacement_locations);
         PyErr_NoMemory();
         return NULL;
     }
     init_code(co, con);
+#ifdef Py_GIL_DISABLED
+    _PyObject_SetDeferredRefcount((PyObject *)co);
+    _PyObject_GC_TRACK(co);
+#endif
     Py_XDECREF(replacement_locations);
     return co;
 }
@@ -1710,6 +1720,10 @@ code_dealloc(PyCodeObject *co)
     }
     Py_SET_REFCNT(co, 0);
 
+#ifdef Py_GIL_DISABLED
+    PyObject_GC_UnTrack(co);
+#endif
+
     _PyFunction_ClearCodeByVersion(co->co_version);
     if (co->co_extra != NULL) {
         PyInterpreterState *interp = _PyInterpreterState_GET();
@@ -1752,6 +1766,15 @@ code_dealloc(PyCodeObject *co)
     PyObject_Free(co);
 }
 
+#ifdef Py_GIL_DISABLED
+static int
+code_traverse(PyCodeObject *co, visitproc visit, void *arg)
+{
+    Py_VISIT(co->co_consts);
+    return 0;
+}
+#endif
+
 static PyObject *
 code_repr(PyCodeObject *co)
 {
@@ -2196,9 +2219,17 @@ PyTypeObject PyCode_Type = {
     PyObject_GenericGetAttr,            /* tp_getattro */
     0,                                  /* tp_setattro */
     0,                                  /* tp_as_buffer */
+#ifdef Py_GIL_DISABLED
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
+#else
     Py_TPFLAGS_DEFAULT,                 /* tp_flags */
+#endif
     code_new__doc__,                    /* tp_doc */
+#ifdef Py_GIL_DISABLED
+    (traverseproc)code_traverse,        /* tp_traverse */
+#else
     0,                                  /* tp_traverse */
+#endif
     0,                                  /* tp_clear */
     code_richcompare,                   /* tp_richcompare */
     offsetof(PyCodeObject, co_weakreflist),     /* tp_weaklistoffset */