gh-120161: Fix a Crash in the _datetime Modulegentoo-3.13.0b2_p1

Taken from https://github.com/python/cpython/pull/120182

5 files changed, 131 insertions, 71 deletions

diff --git a/Include/internal/pycore_typeobject.h b/Include/internal/pycore_typeobject.h
index 8664ae0e445..05f125f928c 100644
--- a/Include/internal/pycore_typeobject.h
+++ b/Include/internal/pycore_typeobject.h
@@ -17,11 +17,25 @@ extern "C" {
 #define _Py_TYPE_BASE_VERSION_TAG (2<<16)
 #define _Py_MAX_GLOBAL_TYPE_VERSION_TAG (_Py_TYPE_BASE_VERSION_TAG - 1)
 
+/* For now we hard-code this to a value for which we are confident
+   all the static builtin types will fit (for all builds). */
+#define _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES 200
+#define _Py_MAX_MANAGED_STATIC_EXT_TYPES 10
+#define _Py_MAX_MANAGED_STATIC_TYPES \
+    (_Py_MAX_MANAGED_STATIC_BUILTIN_TYPES + _Py_MAX_MANAGED_STATIC_EXT_TYPES)
+
 struct _types_runtime_state {
     /* Used to set PyTypeObject.tp_version_tag for core static types. */
     // bpo-42745: next_version_tag remains shared by all interpreters
     // because of static types.
     unsigned int next_version_tag;
+
+    struct {
+        struct {
+            PyTypeObject *type;
+            int64_t interp_count;
+        } types[_Py_MAX_MANAGED_STATIC_TYPES];
+    } managed_static;
 };
 
 
@@ -42,11 +56,6 @@ struct type_cache {
     struct type_cache_entry hashtable[1 << MCACHE_SIZE_EXP];
 };
 
-/* For now we hard-code this to a value for which we are confident
-   all the static builtin types will fit (for all builds). */
-#define _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES 200
-#define _Py_MAX_MANAGED_STATIC_EXT_TYPES 10
-
 typedef struct {
     PyTypeObject *type;
     int isbuiltin;
@@ -125,6 +134,7 @@ struct types_state {
 
 extern PyStatus _PyTypes_InitTypes(PyInterpreterState *);
 extern void _PyTypes_FiniTypes(PyInterpreterState *);
+extern void _PyTypes_FiniExtTypes(PyInterpreterState *interp);
 extern void _PyTypes_Fini(PyInterpreterState *);
 extern void _PyTypes_AfterFork(void);
 
@@ -163,10 +173,6 @@ extern managed_static_type_state * _PyStaticType_GetState(
 PyAPI_FUNC(int) _PyStaticType_InitForExtension(
     PyInterpreterState *interp,
      PyTypeObject *self);
-PyAPI_FUNC(void) _PyStaticType_FiniForExtension(
-    PyInterpreterState *interp,
-     PyTypeObject *self,
-     int final);
 
 
 /* Like PyType_GetModuleState, but skips verification
diff --git a/Lib/test/datetimetester.py b/Lib/test/datetimetester.py
index db6502b3471..195db92c04b 100644
--- a/Lib/test/datetimetester.py
+++ b/Lib/test/datetimetester.py
@@ -13,6 +13,7 @@ import random
 import re
 import struct
 import sys
+import textwrap
 import unittest
 import warnings
 
@@ -22,7 +23,7 @@ from operator import lt, le, gt, ge, eq, ne, truediv, floordiv, mod
 
 from test import support
 from test.support import is_resource_enabled, ALWAYS_EQ, LARGEST, SMALLEST
-from test.support import warnings_helper
+from test.support import script_helper, warnings_helper
 
 import datetime as datetime_module
 from datetime import MINYEAR, MAXYEAR
@@ -6794,6 +6795,47 @@ class CapiTest(unittest.TestCase):
                     self.assertEqual(dt_orig, dt_rt)
 
 
+class ExtensionModuleTests(unittest.TestCase):
+
+    def setUp(self):
+        if self.__class__.__name__.endswith('Pure'):
+            self.skipTest('Not relevant in pure Python')
+
+    def test_gh_120161(self):
+        with self.subTest('simple'):
+            script = textwrap.dedent("""
+                import datetime
+                from _ast import Tuple
+                f = lambda: None
+                Tuple.dims = property(f, f)
+
+                class tzutc(datetime.tzinfo):
+                    pass
+                """)
+            script_helper.assert_python_ok('-c', script)
+
+        with self.subTest('complex'):
+            script = textwrap.dedent("""
+                import asyncio
+                import datetime
+                from typing import Type
+
+                class tzutc(datetime.tzinfo):
+                    pass
+                _EPOCHTZ = datetime.datetime(1970, 1, 1, tzinfo=tzutc())
+
+                class FakeDateMeta(type):
+                    def __instancecheck__(self, obj):
+                        return True
+                class FakeDate(datetime.date, metaclass=FakeDateMeta):
+                    pass
+                def pickle_fake_date(datetime_) -> Type[FakeDate]:
+                    # A pickle function for FakeDate
+                    return FakeDate
+                """)
+            script_helper.assert_python_ok('-c', script)
+
+
 def load_tests(loader, standard_tests, pattern):
     standard_tests.addTest(ZoneInfoCompleteTest())
     return standard_tests
diff --git a/Modules/_datetimemodule.c b/Modules/_datetimemodule.c
index d6fa273c75e..770f637d31a 100644
--- a/Modules/_datetimemodule.c
+++ b/Modules/_datetimemodule.c
@@ -7126,37 +7126,6 @@ clear_state(datetime_state *st)
 }
 
 
-/* ---------------------------------------------------------------------------
- * Global module state.
- */
-
-// If we make _PyStaticType_*ForExtension() public
-// then all this should be managed by the runtime.
-
-static struct {
-    PyMutex mutex;
-    int64_t interp_count;
-} _globals = {0};
-
-static void
-callback_for_interp_exit(void *Py_UNUSED(data))
-{
-    PyInterpreterState *interp = PyInterpreterState_Get();
-
-    assert(_globals.interp_count > 0);
-    PyMutex_Lock(&_globals.mutex);
-    _globals.interp_count -= 1;
-    int final = !_globals.interp_count;
-    PyMutex_Unlock(&_globals.mutex);
-
-    /* They must be done in reverse order so subclasses are finalized
-     * before base classes. */
-    for (size_t i = Py_ARRAY_LENGTH(capi_types); i > 0; i--) {
-        PyTypeObject *type = capi_types[i-1];
-        _PyStaticType_FiniForExtension(interp, type, final);
-    }
-}
-
 static int
 init_static_types(PyInterpreterState *interp, int reloading)
 {
@@ -7179,19 +7148,6 @@ init_static_types(PyInterpreterState *interp, int reloading)
         }
     }
 
-    PyMutex_Lock(&_globals.mutex);
-    assert(_globals.interp_count >= 0);
-    _globals.interp_count += 1;
-    PyMutex_Unlock(&_globals.mutex);
-
-    /* It could make sense to add a separate callback
-     * for each of the types.  However, for now we can take the simpler
-     * approach of a single callback. */
-    if (PyUnstable_AtExit(interp, callback_for_interp_exit, NULL) < 0) {
-        callback_for_interp_exit(NULL);
-        return -1;
-    }
-
     return 0;
 }
 
@@ -7376,8 +7332,8 @@ module_clear(PyObject *mod)
     PyInterpreterState *interp = PyInterpreterState_Get();
     clear_current_module(interp, mod);
 
-    // We take care of the static types via an interpreter atexit hook.
-    // See callback_for_interp_exit() above.
+    // The runtime takes care of the static types for us.
+    // See _PyTypes_FiniExtTypes()..
 
     return 0;
 }
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 18b7dfce3d8..f55b77a6e54 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -159,18 +159,28 @@ managed_static_type_index_clear(PyTypeObject *self)
     self->tp_subclasses = NULL;
 }
 
-static inline managed_static_type_state *
-static_builtin_state_get(PyInterpreterState *interp, PyTypeObject *self)
+static PyTypeObject *
+static_ext_type_lookup(PyInterpreterState *interp, size_t index,
+                       int64_t *p_interp_count)
 {
-    return &(interp->types.builtins.initialized[
-                        managed_static_type_index_get(self)]);
-}
+    assert(interp->runtime == &_PyRuntime);
+    assert(index < _Py_MAX_MANAGED_STATIC_EXT_TYPES);
 
-static inline managed_static_type_state *
-static_ext_type_state_get(PyInterpreterState *interp, PyTypeObject *self)
-{
-    return &(interp->types.for_extensions.initialized[
-                        managed_static_type_index_get(self)]);
+    size_t full_index = index + _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES;
+    int64_t interp_count =
+            _PyRuntime.types.managed_static.types[full_index].interp_count;
+    assert((interp_count == 0) ==
+            (_PyRuntime.types.managed_static.types[full_index].type == NULL));
+    *p_interp_count = interp_count;
+
+    PyTypeObject *type = interp->types.for_extensions.initialized[index].type;
+    if (type == NULL) {
+        return NULL;
+    }
+    assert(!interp->types.for_extensions.initialized[index].isbuiltin);
+    assert(type == _PyRuntime.types.managed_static.types[full_index].type);
+    assert(managed_static_type_index_is_set(type));
+    return type;
 }
 
 static managed_static_type_state *
@@ -202,6 +212,8 @@ static void
 managed_static_type_state_init(PyInterpreterState *interp, PyTypeObject *self,
                                int isbuiltin, int initial)
 {
+    assert(interp->runtime == &_PyRuntime);
+
     size_t index;
     if (initial) {
         assert(!managed_static_type_index_is_set(self));
@@ -228,6 +240,21 @@ managed_static_type_state_init(PyInterpreterState *interp, PyTypeObject *self,
             assert(index < _Py_MAX_MANAGED_STATIC_EXT_TYPES);
         }
     }
+    size_t full_index = isbuiltin
+        ? index
+        : index + _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES;
+
+    assert((initial == 1) ==
+            (_PyRuntime.types.managed_static.types[full_index].interp_count == 0));
+    _PyRuntime.types.managed_static.types[full_index].interp_count += 1;
+
+    if (initial) {
+        assert(_PyRuntime.types.managed_static.types[full_index].type == NULL);
+        _PyRuntime.types.managed_static.types[full_index].type = self;
+    }
+    else {
+        assert(_PyRuntime.types.managed_static.types[full_index].type == self);
+    }
 
     managed_static_type_state *state = isbuiltin
         ? &(interp->types.builtins.initialized[index])
@@ -256,15 +283,28 @@ static void
 managed_static_type_state_clear(PyInterpreterState *interp, PyTypeObject *self,
                                 int isbuiltin, int final)
 {
+    size_t index = managed_static_type_index_get(self);
+    size_t full_index = isbuiltin
+        ? index
+        : index + _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES;
+
     managed_static_type_state *state = isbuiltin
-        ? static_builtin_state_get(interp, self)
-        : static_ext_type_state_get(interp, self);
+        ? &(interp->types.builtins.initialized[index])
+        : &(interp->types.for_extensions.initialized[index]);
+    assert(state != NULL);
+
+    assert(_PyRuntime.types.managed_static.types[full_index].interp_count > 0);
+    assert(_PyRuntime.types.managed_static.types[full_index].type == state->type);
 
     assert(state->type != NULL);
     state->type = NULL;
     assert(state->tp_weaklist == NULL);  // It was already cleared out.
 
+    _PyRuntime.types.managed_static.types[full_index].interp_count -= 1;
     if (final) {
+        assert(!_PyRuntime.types.managed_static.types[full_index].interp_count);
+        _PyRuntime.types.managed_static.types[full_index].type = NULL;
+
         managed_static_type_index_clear(self);
     }
 
@@ -840,8 +880,12 @@ _PyTypes_Fini(PyInterpreterState *interp)
     struct type_cache *cache = &interp->types.type_cache;
     type_cache_clear(cache, NULL);
 
+    // All the managed static types should have been finalized already.
+    assert(interp->types.for_extensions.num_initialized == 0);
+    for (size_t i = 0; i < _Py_MAX_MANAGED_STATIC_EXT_TYPES; i++) {
+        assert(interp->types.for_extensions.initialized[i].type == NULL);
+    }
     assert(interp->types.builtins.num_initialized == 0);
-    // All the static builtin types should have been finalized already.
     for (size_t i = 0; i < _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES; i++) {
         assert(interp->types.builtins.initialized[i].type == NULL);
     }
@@ -5674,9 +5718,20 @@ fini_static_type(PyInterpreterState *interp, PyTypeObject *type,
 }
 
 void
-_PyStaticType_FiniForExtension(PyInterpreterState *interp, PyTypeObject *type, int final)
+_PyTypes_FiniExtTypes(PyInterpreterState *interp)
 {
-    fini_static_type(interp, type, 0, final);
+    for (size_t i = _Py_MAX_MANAGED_STATIC_EXT_TYPES; i > 0; i--) {
+        if (interp->types.for_extensions.num_initialized == 0) {
+            break;
+        }
+        int64_t count = 0;
+        PyTypeObject *type = static_ext_type_lookup(interp, i-1, &count);
+        if (type == NULL) {
+            continue;
+        }
+        int final = (count == 1);
+        fini_static_type(interp, type, 0, final);
+    }
 }
 
 void
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index cbdf5c1b771..3639cf67120 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -1818,6 +1818,7 @@ flush_std_files(void)
 static void
 finalize_interp_types(PyInterpreterState *interp)
 {
+    _PyTypes_FiniExtTypes(interp);
     _PyUnicode_FiniTypes(interp);
     _PySys_FiniTypes(interp);
     _PyXI_FiniTypes(interp);