gh-127020: Make `PyCode_GetCode` thread-safe for free threading (#127043)

Some fields in PyCodeObject are lazily initialized. Use atomics and
critical sections to make their initializations and accesses thread-safe.

3 files changed, 85 insertions, 27 deletions

diff --git a/Lib/test/test_free_threading/test_code.py b/Lib/test/test_free_threading/test_code.py
new file mode 100644
index 0000000..a5136a3
--- /dev/null
+++ b/Lib/test/test_free_threading/test_code.py
@@ -0,0 +1,30 @@
+import unittest
+
+from threading import Thread
+from unittest import TestCase
+
+from test.support import threading_helper
+
+@threading_helper.requires_working_threading()
+class TestCode(TestCase):
+    def test_code_attrs(self):
+        """Test concurrent accesses to lazily initialized code attributes"""
+        code_objects = []
+        for _ in range(1000):
+            code_objects.append(compile("a + b", "<string>", "eval"))
+
+        def run_in_thread():
+            for code in code_objects:
+                self.assertIsInstance(code.co_code, bytes)
+                self.assertIsInstance(code.co_freevars, tuple)
+                self.assertIsInstance(code.co_varnames, tuple)
+
+        threads = [Thread(target=run_in_thread) for _ in range(2)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-11-19-21-49-58.gh-issue-127020.5vvI17.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-11-19-21-49-58.gh-issue-127020.5vvI17.rst
new file mode 100644
index 0000000..a8fd927
--- /dev/null
+++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-11-19-21-49-58.gh-issue-127020.5vvI17.rst
@@ -0,0 +1,4 @@
+Fix a crash in the free threading build when :c:func:`PyCode_GetCode`,
+:c:func:`PyCode_GetVarnames`, :c:func:`PyCode_GetCellvars`, or
+:c:func:`PyCode_GetFreevars` were called from multiple threads at the same
+time.
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index dba43d5..ec5d8a5 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -302,21 +302,32 @@ validate_and_copy_tuple(PyObject *tup)
 }
 
 static int
-init_co_cached(PyCodeObject *self) {
-    if (self->_co_cached == NULL) {
-        self->_co_cached = PyMem_New(_PyCoCached, 1);
-        if (self->_co_cached == NULL) {
+init_co_cached(PyCodeObject *self)
+{
+    _PyCoCached *cached = FT_ATOMIC_LOAD_PTR(self->_co_cached);
+    if (cached != NULL) {
+        return 0;
+    }
+
+    Py_BEGIN_CRITICAL_SECTION(self);
+    cached = self->_co_cached;
+    if (cached == NULL) {
+        cached = PyMem_New(_PyCoCached, 1);
+        if (cached == NULL) {
             PyErr_NoMemory();
-            return -1;
         }
-        self->_co_cached->_co_code = NULL;
-        self->_co_cached->_co_cellvars = NULL;
-        self->_co_cached->_co_freevars = NULL;
-        self->_co_cached->_co_varnames = NULL;
+        else {
+            cached->_co_code = NULL;
+            cached->_co_cellvars = NULL;
+            cached->_co_freevars = NULL;
+            cached->_co_varnames = NULL;
+            FT_ATOMIC_STORE_PTR(self->_co_cached, cached);
+        }
     }
-    return 0;
-
+    Py_END_CRITICAL_SECTION();
+    return cached != NULL ? 0 : -1;
 }
+
 /******************
  * _PyCode_New()
  ******************/
@@ -1571,16 +1582,21 @@ get_cached_locals(PyCodeObject *co, PyObject **cached_field,
 {
     assert(cached_field != NULL);
     assert(co->_co_cached != NULL);
-    if (*cached_field != NULL) {
-        return Py_NewRef(*cached_field);
+    PyObject *varnames = FT_ATOMIC_LOAD_PTR(*cached_field);
+    if (varnames != NULL) {
+        return Py_NewRef(varnames);
     }
-    assert(*cached_field == NULL);
-    PyObject *varnames = get_localsplus_names(co, kind, num);
+
+    Py_BEGIN_CRITICAL_SECTION(co);
+    varnames = *cached_field;
     if (varnames == NULL) {
-        return NULL;
+        varnames = get_localsplus_names(co, kind, num);
+        if (varnames != NULL) {
+            FT_ATOMIC_STORE_PTR(*cached_field, varnames);
+        }
     }
-    *cached_field = Py_NewRef(varnames);
-    return varnames;
+    Py_END_CRITICAL_SECTION();
+    return Py_XNewRef(varnames);
 }
 
 PyObject *
@@ -1674,18 +1690,26 @@ _PyCode_GetCode(PyCodeObject *co)
     if (init_co_cached(co)) {
         return NULL;
     }
-    if (co->_co_cached->_co_code != NULL) {
-        return Py_NewRef(co->_co_cached->_co_code);
+
+    _PyCoCached *cached = co->_co_cached;
+    PyObject *code = FT_ATOMIC_LOAD_PTR(cached->_co_code);
+    if (code != NULL) {
+        return Py_NewRef(code);
     }
-    PyObject *code = PyBytes_FromStringAndSize((const char *)_PyCode_CODE(co),
-                                               _PyCode_NBYTES(co));
+
+    Py_BEGIN_CRITICAL_SECTION(co);
+    code = cached->_co_code;
     if (code == NULL) {
-        return NULL;
+        code = PyBytes_FromStringAndSize((const char *)_PyCode_CODE(co),
+                                         _PyCode_NBYTES(co));
+        if (code != NULL) {
+            deopt_code(co, (_Py_CODEUNIT *)PyBytes_AS_STRING(code));
+            assert(cached->_co_code == NULL);
+            FT_ATOMIC_STORE_PTR(cached->_co_code, code);
+        }
     }
-    deopt_code(co, (_Py_CODEUNIT *)PyBytes_AS_STRING(code));
-    assert(co->_co_cached->_co_code == NULL);
-    co->_co_cached->_co_code = Py_NewRef(code);
-    return code;
+    Py_END_CRITICAL_SECTION();
+    return Py_XNewRef(code);
 }
 
 PyObject *