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)
 {