gh-84436: Improve Immortalization for Builtin Types

Include/internal/pycore_global_objects.h

@@ -20,6 +20,24 @@ extern "C" {
 #define _PY_NSMALLNEGINTS           5
 
 
+struct _Py_immortalized_object {
+    PyObject *obj;
+    int final_refcnt;
+    struct _Py_immortalized_object *next;
+};
+
+struct _Py_immortalized_objects {
+    Py_ssize_t count;
+    struct _Py_immortalized_object *head;
+    struct _Py_immortalized_object *last;
+#define _Py_IMMORTALIZED_ARRAY_SIZE 100
+    struct _Py_immortalized_object _objects[_Py_IMMORTALIZED_ARRAY_SIZE];
+};
+
+extern void _Py_EnsureImmortal(PyObject *);
+extern void _Py_ImmortalObjectsFini(void);
+
+
 // Only immutable objects should be considered runtime-global.
 // All others must be per-interpreter.
 

Include/internal/pycore_runtime.h

@@ -152,6 +152,8 @@ typedef struct pyruntimestate {
     struct _Py_unicode_runtime_state unicode_state;
     struct _types_runtime_state types;
 
+    /* All non-static immortal objects (need to be cleaned up during fini). */
+    struct _Py_immortalized_objects immortalized_objects;
     /* All the objects that are shared by the runtime's interpreters. */
     struct _Py_static_objects static_objects;
 

Objects/structseq.c

@@ -512,20 +512,34 @@ _PyStructSequence_InitBuiltinWithFlags(PyTypeObject *type,
     PyMemberDef *members;
     Py_ssize_t n_members, n_unnamed_members;
 
+    assert(type->tp_name == NULL);
+    assert(type->tp_members == NULL);
     members = initialize_members(desc, &n_members, &n_unnamed_members);
     if (members == NULL) {
         return -1;
     }
     initialize_static_fields(type, desc, members, tp_flags);
-    if (_PyStaticType_InitBuiltin(type) < 0) {
+
+    PyObject *dict = PyDict_New();
+    if (dict == NULL) {
         PyMem_Free(members);
-        PyErr_Format(PyExc_RuntimeError,
-                     "Can't initialize builtin type %s",
-                     desc->name);
         return -1;
     }
+    type->tp_dict = dict;
+
     if (initialize_static_type(type, desc, n_members, n_unnamed_members) < 0) {
         PyMem_Free(members);
+        Py_CLEAR(type->tp_dict);
+        return -1;
+    }
+
+    _Py_SetImmortal(type);
+    if (_PyStaticType_InitBuiltin(type) < 0) {
+        PyMem_Free(members);
+        Py_CLEAR(type->tp_dict);
+        PyErr_Format(PyExc_RuntimeError,
+                     "Can't initialize builtin type %s",
+                     desc->name);
         return -1;
     }
     return 0;
@@ -570,35 +584,33 @@ PyStructSequence_InitType(PyTypeObject *type, PyStructSequence_Desc *desc)
 }
 
 
+/* This is exposed in the internal API, not the public API.
+   It is only called on builtin static types, which are all
+   initialized via _PyStructSequence_InitBuiltinWithFlags(). */
+
 void
 _PyStructSequence_FiniType(PyTypeObject *type)
 {
     // Ensure that the type is initialized
     assert(type->tp_name != NULL);
     assert(type->tp_base == &PyTuple_Type);
+    assert((type->tp_flags & _Py_TPFLAGS_STATIC_BUILTIN));
+    assert(_Py_IsImmortal(type));
 
     // Cannot delete a type if it still has subclasses
     if (_PyType_HasSubclasses(type)) {
+        // XXX Shouldn't this be an error?
         return;
     }
 
-    // Undo PyStructSequence_NewType()
-    type->tp_name = NULL;
+    // Undo _PyStructSequence_InitBuiltinWithFlags()
     PyMem_Free(type->tp_members);
+    type->tp_members = NULL;
 
     _PyStaticType_Dealloc(type);
-    assert(Py_REFCNT(type) == 1);
-    // Undo Py_INCREF(type) of _PyStructSequence_InitType().
-    // Don't use Py_DECREF(): static type must not be deallocated
-    Py_SET_REFCNT(type, 0);
-#ifdef Py_REF_DEBUG
-    _Py_DecRefTotal(_PyInterpreterState_GET());
-#endif
 
-    // Make sure that _PyStructSequence_InitType() will initialize
-    // the type again
-    assert(Py_REFCNT(type) == 0);
-    assert(type->tp_name == NULL);
+    // Mark the type as un-initialized.
+    type->tp_name = NULL;
 }
 
 

Objects/typeobject.c

@@ -4498,6 +4498,17 @@ _PyStaticType_Dealloc(PyTypeObject *type)
 
     type_dealloc_common(type);
 
+    if (type->tp_flags & _Py_TPFLAGS_STATIC_BUILTIN) {
+        if (type->tp_dict != NULL) {
+            type->tp_dict->ob_refcnt = 1;
+        }
+        if (type->tp_bases != NULL && PyTuple_GET_SIZE(type->tp_bases) > 0) {
+            type->tp_bases->ob_refcnt = 1;
+        }
+        if (type->tp_mro != NULL && PyTuple_GET_SIZE(type->tp_mro) > 0) {
+            type->tp_mro->ob_refcnt = 1;
+        }
+    }
     Py_CLEAR(type->tp_dict);
     Py_CLEAR(type->tp_bases);
     Py_CLEAR(type->tp_mro);
@@ -7030,6 +7041,7 @@ PyType_Ready(PyTypeObject *type)
 int
 _PyStaticType_InitBuiltin(PyTypeObject *self)
 {
+    assert(_Py_IsImmortal((PyObject *)self));
     self->tp_flags |= _Py_TPFLAGS_STATIC_BUILTIN;
 
     assert(NEXT_GLOBAL_VERSION_TAG <= _Py_MAX_GLOBAL_TYPE_VERSION_TAG);
@@ -7042,6 +7054,11 @@ _PyStaticType_InitBuiltin(PyTypeObject *self)
     if (res < 0) {
         static_builtin_state_clear(self);
     }
+
+    _Py_EnsureImmortal(self->tp_dict);
+    _Py_EnsureImmortal(self->tp_bases);
+    _Py_EnsureImmortal(self->tp_mro);
+
     return res;
 }
 

Python/pylifecycle.c

@@ -1733,6 +1733,7 @@ finalize_interp_clear(PyThreadState *tstate)
         _Py_ClearFileSystemEncoding();
         _Py_Deepfreeze_Fini();
         _PyPerfTrampoline_Fini();
+        _Py_ImmortalObjectsFini();
     }
 
     finalize_interp_types(tstate->interp);

Python/pystate.c

@@ -2758,6 +2758,99 @@ _register_builtins_for_crossinterpreter_data(struct _xidregistry *xidregistry)
 }
 
 
+/******************/
+/* global objects */
+/******************/
+
+static void
+immortalized_add(struct _Py_immortalized_objects *state, PyObject *obj)
+{
+    struct _Py_immortalized_object *entry;
+    if (state->count < _Py_IMMORTALIZED_ARRAY_SIZE) {
+        entry = &state->_objects[state->count];
+    }
+    else {
+        entry = PyMem_RawMalloc(sizeof(*entry));
+        if (entry == NULL) {
+            Py_FatalError("could not allocate immortalized list entry");
+        }
+    }
+
+    entry->obj = obj;
+    entry->final_refcnt = Py_REFCNT(obj);
+    entry->next = NULL;
+
+    if (state->head == NULL) {
+        assert(state->count == 0);
+        assert(state->last == NULL);
+        state->head = entry;
+    }
+    else {
+        state->last->next = entry;
+    }
+    state->count += 1;
+    state->last = entry;
+}
+
+static void
+immortalized_fini(struct _Py_immortalized_objects *state)
+{
+    struct _Py_immortalized_object *next = state->head;
+    state->head = NULL;
+    state->last = NULL;
+    for (int i = 0; i < _Py_IMMORTALIZED_ARRAY_SIZE && next != NULL; i++) {
+        struct _Py_immortalized_object *entry = next;
+        next = entry->next;
+        entry->obj->ob_refcnt = entry->final_refcnt;
+    }
+    while (next != NULL) {
+        struct _Py_immortalized_object *entry = next;
+        next = next->next;
+        entry->obj->ob_refcnt = entry->final_refcnt;
+        PyMem_RawFree(entry);
+    }
+}
+
+void
+_Py_EnsureImmortal(PyObject *obj)
+{
+    assert(_Py_IsMainInterpreter(_PyInterpreterState_GET()));
+    if (_Py_IsImmortal(obj)) {
+        return;
+    }
+
+    _Py_SetImmortal(obj);
+    immortalized_add(&_PyRuntime.immortalized_objects, obj);
+
+    if (Py_TYPE(obj) == &PyDict_Type) {
+        Py_ssize_t i = 0;
+        PyObject *key, *value;  // borrowed ref
+        while (PyDict_Next(obj, &i, &key, &value)) {
+            _Py_EnsureImmortal(key);
+            _Py_EnsureImmortal(value);
+        }
+    }
+    else if (Py_TYPE(obj) == &PyTuple_Type) {
+        assert(PyTuple_GET_SIZE(obj) > 0);
+        Py_ssize_t size = PyTuple_GET_SIZE(obj);
+        assert(size > 0);
+        for (Py_ssize_t i = 0; i < size; i++) {
+            _Py_EnsureImmortal(PyTuple_GET_ITEM(obj, i));
+        }
+    }
+}
+
+void
+_Py_ImmortalObjectsFini(void)
+{
+    immortalized_fini(&_PyRuntime.immortalized_objects);
+}
+
+
+/*************/
+/* other API */
+/*************/
+
 _PyFrameEvalFunction
 _PyInterpreterState_GetEvalFrameFunc(PyInterpreterState *interp)
 {