Hold current module until interp-end

Author: neonene

Include/internal/pycore_interp_structs.h

@@ -947,6 +947,8 @@ struct _is {
     _Py_hashtable_t *closed_stackrefs_table;
 #  endif
 #endif
+
+    void *datetime_module_state;
 };
 
 

Modules/_datetimemodule.c

@@ -49,6 +49,10 @@ typedef struct {
 
     /* The interned Unix epoch datetime instance */
     PyObject *epoch;
+
+    /* Reference to the interpreter's dict where the current module will be
+     * reserved even after the referent dict becomes NULL at shutdown. */
+    PyObject *interp_dict;
 } datetime_state;
 
 /* The module has a fixed number of static objects, due to being exposed
@@ -133,18 +137,13 @@ get_current_module(PyInterpreterState *interp, int *p_reloading)
     if (dict == NULL) {
         goto error;
     }
-    PyObject *ref = NULL;
-    if (PyDict_GetItemRef(dict, INTERP_KEY, &ref) < 0) {
+    if (PyDict_GetItemRef(dict, INTERP_KEY, &mod) < 0) {
         goto error;
     }
-    if (ref != NULL) {
+    if (mod != NULL) {
         reloading = 1;
-        if (ref != Py_None) {
-            (void)PyWeakref_GetRef(ref, &mod);
-            if (mod == Py_None) {
-                Py_CLEAR(mod);
-            }
-            Py_DECREF(ref);
+        if (mod == Py_None) {
+            mod = NULL;
         }
     }
     if (p_reloading != NULL) {
@@ -159,124 +158,59 @@ get_current_module(PyInterpreterState *interp, int *p_reloading)
 
 static PyModuleDef datetimemodule;
 
-static datetime_state *
-_get_current_state(PyObject **p_mod)
+static inline datetime_state *
+get_current_state()
 {
     PyInterpreterState *interp = PyInterpreterState_Get();
-    PyObject *mod = get_current_module(interp, NULL);
-    if (mod == NULL) {
-        assert(!PyErr_Occurred());
-        if (PyErr_Occurred()) {
-            return NULL;
-        }
-        /* The static types can outlive the module,
-         * so we must re-import the module. */
-        mod = PyImport_ImportModule("_datetime");
-        if (mod == NULL) {
-            PyErr_Clear();
-            /* Create a module at shutdown */
-            PyObject *dict = PyInterpreterState_GetDict(interp);
-            if (dict == NULL) {
-                return NULL;
-            }
-            PyObject *spec;
-            if (PyDict_GetItemStringRef(dict, "datetime_module_spec", &spec) != 1) {
-                return NULL;
-            }
-            mod = PyModule_FromDefAndSpec(&datetimemodule, spec);
-            if (mod == NULL) {
-                Py_DECREF(spec);
-                return NULL;
-            }
-            Py_DECREF(spec);
-
-            /* The module will be held by heaptypes. Prefer
-             * it not to be stored in the interpreter's dict. */
-            if (PyModule_ExecDef(mod, &datetimemodule) < 0) {
-                return NULL;
-            }
-        }
-    }
-    datetime_state *st = get_module_state(mod);
-    *p_mod = mod;
-    return st;
+    void *state = interp->datetime_module_state;
+    assert(state != NULL);
+    return (datetime_state *)state;
 }
 
-#define GET_CURRENT_STATE(MOD_VAR)  \
-    _get_current_state(&MOD_VAR)
-#define RELEASE_CURRENT_STATE(ST_VAR, MOD_VAR)  \
-    Py_DECREF(MOD_VAR)
-
 static int
-set_current_module(PyInterpreterState *interp, PyObject *mod)
+set_current_module(datetime_state *st,
+                   PyInterpreterState *interp, PyObject *mod)
 {
     assert(mod != NULL);
-    PyObject *dict = PyInterpreterState_GetDict(interp);
+    PyObject *dict = st->interp_dict;
     if (dict == NULL) {
         return -1;
     }
-    PyObject *ref = PyWeakref_NewRef(mod, NULL);
-    if (ref == NULL) {
-        return -1;
-    }
-    if (PyDict_SetItem(dict, INTERP_KEY, ref) < 0) {
-        Py_DECREF(ref);
-        return -1;
-    }
-    Py_DECREF(ref);
-
-    /* Make the module spec remain in the interpreter's dict. Not required,
-     * but reserve the new one for memory efficiency. */
-    PyObject *mod_dict = PyModule_GetDict(mod);
-    if (mod_dict == NULL) {
-        return -1;
-    }
-    PyObject *spec;
-    if (PyDict_GetItemRef(mod_dict, &_Py_ID(__spec__), &spec) != 1) {
-        return -1;
-    }
-    if (PyDict_SetItemString(dict, "datetime_module_spec", spec) < 0) {
-        Py_DECREF(spec);
+    if (PyDict_SetItem(dict, INTERP_KEY, mod) < 0) {
         return -1;
     }
-    Py_DECREF(spec);
+    interp->datetime_module_state = st;
     return 0;
 }
 
 static void
-clear_current_module(PyInterpreterState *interp, PyObject *expected)
+clear_current_module(datetime_state *st,
+                     PyInterpreterState *interp, PyObject *expected)
 {
-    PyObject *exc = PyErr_GetRaisedException();
-
-    PyObject *dict = PyInterpreterState_GetDict(interp);
+    PyObject *dict = st->interp_dict;
     if (dict == NULL) {
-        goto error;
+        return;  /* Already cleared */
     }
 
+    PyObject *exc = PyErr_GetRaisedException();
+
     if (expected != NULL) {
-        PyObject *ref = NULL;
-        if (PyDict_GetItemRef(dict, INTERP_KEY, &ref) < 0) {
+        PyObject *current;
+        if (PyDict_GetItemRef(dict, INTERP_KEY, &current) < 0) {
             goto error;
         }
-        if (ref != NULL) {
-            PyObject *current = NULL;
-            int rc = PyWeakref_GetRef(ref, &current);
-            /* We only need "current" for pointer comparison. */
-            Py_XDECREF(current);
-            Py_DECREF(ref);
-            if (rc < 0) {
-                goto error;
-            }
-            if (current != expected) {
-                goto finally;
-            }
+        /* We only need "current" for pointer comparison. */
+        Py_XDECREF(current);
+        if (current != expected) {
+            goto finally;
         }
     }
 
     /* We use None to identify that the module was previously loaded. */
     if (PyDict_SetItem(dict, INTERP_KEY, Py_None) < 0) {
         goto error;
     }
+    interp->datetime_module_state = NULL;
 
     goto finally;
 
@@ -2147,8 +2081,7 @@ delta_to_microseconds(PyDateTime_Delta *self)
     PyObject *x3 = NULL;
     PyObject *result = NULL;
 
-    PyObject *current_mod = NULL;
-    datetime_state *st = GET_CURRENT_STATE(current_mod);
+    datetime_state *st = get_current_state();
 
     x1 = PyLong_FromLong(GET_TD_DAYS(self));
     if (x1 == NULL)
@@ -2186,7 +2119,6 @@ delta_to_microseconds(PyDateTime_Delta *self)
     Py_XDECREF(x1);
     Py_XDECREF(x2);
     Py_XDECREF(x3);
-    RELEASE_CURRENT_STATE(st, current_mod);
     return result;
 }
 
@@ -2226,8 +2158,7 @@ microseconds_to_delta_ex(PyObject *pyus, PyTypeObject *type)
     PyObject *num = NULL;
     PyObject *result = NULL;
 
-    PyObject *current_mod = NULL;
-    datetime_state *st = GET_CURRENT_STATE(current_mod);
+    datetime_state *st = get_current_state();
 
     tuple = checked_divmod(pyus, CONST_US_PER_SECOND(st));
     if (tuple == NULL) {
@@ -2272,7 +2203,6 @@ microseconds_to_delta_ex(PyObject *pyus, PyTypeObject *type)
 Done:
     Py_XDECREF(tuple);
     Py_XDECREF(num);
-    RELEASE_CURRENT_STATE(st, current_mod);
     return result;
 
 BadDivmod:
@@ -2811,8 +2741,7 @@ delta_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 {
     PyObject *self = NULL;
 
-    PyObject *current_mod = NULL;
-    datetime_state *st = GET_CURRENT_STATE(current_mod);
+    datetime_state *st = get_current_state();
 
     /* Argument objects. */
     PyObject *day = NULL;
@@ -2917,7 +2846,6 @@ delta_new(PyTypeObject *type, PyObject *args, PyObject *kw)
     Py_DECREF(x);
 
 Done:
-    RELEASE_CURRENT_STATE(st, current_mod);
     return self;
 
 #undef CLEANUP
@@ -3030,12 +2958,10 @@ delta_total_seconds(PyObject *op, PyObject *Py_UNUSED(dummy))
     if (total_microseconds == NULL)
         return NULL;
 
-    PyObject *current_mod = NULL;
-    datetime_state *st = GET_CURRENT_STATE(current_mod);
+    datetime_state *st = get_current_state();
 
     total_seconds = PyNumber_TrueDivide(total_microseconds, CONST_US_PER_SECOND(st));
 
-    RELEASE_CURRENT_STATE(st, current_mod);
     Py_DECREF(total_microseconds);
     return total_seconds;
 }
@@ -3813,12 +3739,10 @@ date_isocalendar(PyObject *self, PyObject *Py_UNUSED(dummy))
         week = 0;
     }
 
-    PyObject *current_mod = NULL;
-    datetime_state *st = GET_CURRENT_STATE(current_mod);
+    datetime_state *st = get_current_state();
 
     PyObject *v = iso_calendar_date_new_impl(ISOCALENDAR_DATE_TYPE(st),
                                              year, week + 1, day + 1);
-    RELEASE_CURRENT_STATE(st, current_mod);
     if (v == NULL) {
         return NULL;
     }
@@ -6638,11 +6562,9 @@ local_timezone(PyDateTime_DateTime *utc_time)
     PyObject *one_second;
     PyObject *seconds;
 
-    PyObject *current_mod = NULL;
-    datetime_state *st = GET_CURRENT_STATE(current_mod);
+    datetime_state *st = get_current_state();
 
     delta = datetime_subtract((PyObject *)utc_time, CONST_EPOCH(st));
-    RELEASE_CURRENT_STATE(st, current_mod);
     if (delta == NULL)
         return NULL;
 
@@ -6882,12 +6804,10 @@ datetime_timestamp(PyObject *op, PyObject *Py_UNUSED(dummy))
     PyObject *result;
 
     if (HASTZINFO(self) && self->tzinfo != Py_None) {
-        PyObject *current_mod = NULL;
-        datetime_state *st = GET_CURRENT_STATE(current_mod);
+        datetime_state *st = get_current_state();
 
         PyObject *delta;
         delta = datetime_subtract(op, CONST_EPOCH(st));
-        RELEASE_CURRENT_STATE(st, current_mod);
         if (delta == NULL)
             return NULL;
         result = delta_total_seconds(delta, NULL);
@@ -7252,8 +7172,15 @@ create_timezone_from_delta(int days, int sec, int ms, int normalize)
  */
 
 static int
-init_state(datetime_state *st, PyObject *module, PyObject *old_module)
+init_state(datetime_state *st,
+           PyInterpreterState *interp, PyObject *module, PyObject *old_module)
 {
+    PyObject *dict = PyInterpreterState_GetDict(interp);
+    if (dict == NULL) {
+        return -1;
+    }
+    st->interp_dict = Py_NewRef(dict);
+
     /* Each module gets its own heap types. */
 #define ADD_TYPE(FIELD, SPEC, BASE)                 \
     do {                                            \
@@ -7272,6 +7199,7 @@ init_state(datetime_state *st, PyObject *module, PyObject *old_module)
         assert(old_module != module);
         datetime_state *st_old = get_module_state(old_module);
         *st = (datetime_state){
+            .interp_dict = st->interp_dict,
             .isocalendar_date_type = st->isocalendar_date_type,
             .us_per_ms = Py_NewRef(st_old->us_per_ms),
             .us_per_second = Py_NewRef(st_old->us_per_second),
@@ -7331,9 +7259,8 @@ init_state(datetime_state *st, PyObject *module, PyObject *old_module)
 static int
 traverse_state(datetime_state *st, visitproc visit, void *arg)
 {
-    /* heap types */
     Py_VISIT(st->isocalendar_date_type);
-
+    Py_VISIT(st->interp_dict);
     return 0;
 }
 
@@ -7349,6 +7276,7 @@ clear_state(datetime_state *st)
     Py_CLEAR(st->us_per_week);
     Py_CLEAR(st->seconds_per_day);
     Py_CLEAR(st->epoch);
+    Py_CLEAR(st->interp_dict);
     return 0;
 }
 
@@ -7416,7 +7344,7 @@ _datetime_exec(PyObject *module)
         }
     }
 
-    if (init_state(st, module, old_module) < 0) {
+    if (init_state(st, interp, module, old_module) < 0) {
         goto error;
     }
 
@@ -7522,7 +7450,7 @@ _datetime_exec(PyObject *module)
     static_assert(DI100Y == 25 * DI4Y - 1, "DI100Y");
     assert(DI100Y == days_before_year(100+1));
 
-    if (set_current_module(interp, module) < 0) {
+    if (set_current_module(st, interp, module) < 0) {
         goto error;
     }
 
@@ -7556,10 +7484,9 @@ static int
 module_clear(PyObject *mod)
 {
     datetime_state *st = get_module_state(mod);
-    clear_state(st);
-
     PyInterpreterState *interp = PyInterpreterState_Get();
-    clear_current_module(interp, mod);
+    clear_current_module(st, interp, mod);
+    clear_state(st);
 
     // The runtime takes care of the static types for us.
     // See _PyTypes_FiniExtTypes()..