From 7c9da242ef59824697e500ab0bab80420f516c22 Mon Sep 17 00:00:00 2001
From: Matt Page <mpage@meta.com>
Date: Fri, 27 Sep 2024 15:29:00 -0700
Subject: [PATCH] Exclude tlbc from refleak counts

TLBC is intended to be reused across threads with different lifetimes,
so may appear as a "leak" depending on the order in which threads execute
code objects. They are freed when the code object is destroyed, which
typically occurs when the runtime is finalized.
---
 Lib/test/libregrtest/refleak.py |  6 ++++++
 Modules/_testinternalcapi.c     | 28 ++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/Lib/test/libregrtest/refleak.py b/Lib/test/libregrtest/refleak.py
index fa447a4336a399..2e495ee2a1540f 100644
--- a/Lib/test/libregrtest/refleak.py
+++ b/Lib/test/libregrtest/refleak.py
@@ -145,6 +145,12 @@ def get_pooled_int(value):
             # Use an internal-only keyword argument that mypy doesn't know yet
             _only_immortal=True)  # type: ignore[call-arg]
         alloc_after = getallocatedblocks() - interned_immortal_after
+        if support.Py_GIL_DISABLED:
+            # Ignore any thread-local bytecode that was allocated. These will be
+            # released when the code object is destroyed, typically at runtime
+            # shutdown
+            import _testinternalcapi
+            alloc_after -= _testinternalcapi.get_tlbc_blocks()
         rc_after = gettotalrefcount()
         fd_after = fd_count()
 
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index 6170d127385152..cb8429db692c95 100644
--- a/Modules/_testinternalcapi.c
+++ b/Modules/_testinternalcapi.c
@@ -2006,6 +2006,33 @@ get_tlbc_id(PyObject *Py_UNUSED(module), PyObject *obj)
     }
     return PyLong_FromVoidPtr(bc);
 }
+
+static int
+count_tlbc_blocks(PyObject *obj, Py_ssize_t *count)
+{
+    if (PyCode_Check(obj)) {
+        _PyCodeArray *tlbc = ((PyCodeObject *)obj)->co_tlbc;
+        // First entry always points to the bytecode at the end of the code
+        // object. Exclude it from the count as it is allocated as part of
+        // creating the code object.
+        for (Py_ssize_t i = 1; i < tlbc->size; i++) {
+            if (tlbc->entries[i] != NULL) {
+                (*count)++;
+            }
+        }
+    }
+    return 1;
+}
+
+// Return the total number of thread-local bytecode copies, excluding the
+// copies that are embedded in the code object.
+static PyObject *
+get_tlbc_blocks(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(ignored))
+{
+    Py_ssize_t count = 0;
+    PyUnstable_GC_VisitObjects((gcvisitobjects_t) count_tlbc_blocks, &count);
+    return PyLong_FromSsize_t(count);
+}
 #endif
 
 static PyObject *
@@ -2180,6 +2207,7 @@ static PyMethodDef module_functions[] = {
     {"py_thread_id", get_py_thread_id, METH_NOARGS},
     {"get_tlbc", get_tlbc, METH_O, NULL},
     {"get_tlbc_id", get_tlbc_id, METH_O, NULL},
+    {"get_tlbc_blocks", get_tlbc_blocks, METH_NOARGS},
 #endif
     {"suppress_immortalization", suppress_immortalization, METH_O},
     {"get_immortalize_deferred", get_immortalize_deferred, METH_NOARGS},