Fix leaks and LOAD_ATTR specialization

Include/internal/pycore_dict.h

@@ -38,6 +38,7 @@ extern PyObject* _PyDict_GetItemIdWithError(PyObject *dp,
 extern int _PyDict_ContainsId(PyObject *, _Py_Identifier *);
 extern int _PyDict_SetItemId(PyObject *dp, _Py_Identifier *key, PyObject *item);
 extern int _PyDict_DelItemId(PyObject *mp, _Py_Identifier *key);
+extern void _PyDict_ClearKeysVersion(PyObject *mp);
 
 extern int _PyDict_Next(
     PyObject *mp, Py_ssize_t *pos, PyObject **key, PyObject **value, Py_hash_t *hash);

Include/internal/pycore_moduleobject.h

@@ -58,6 +58,8 @@ extern Py_ssize_t _PyModule_GetFilenameUTF8(
 PyObject* _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress);
 PyObject* _Py_module_getattro(PyObject *m, PyObject *name);
 
+PyAPI_FUNC(int) _PyModule_ReplaceLazyValue(PyObject *dict, PyObject *name, PyObject *value);
+
 #ifdef __cplusplus
 }
 #endif

Objects/dictobject.c

@@ -4719,6 +4719,14 @@ sizeof_lock_held(PyDictObject *mp)
     return (Py_ssize_t)res;
 }
 
+void
+_PyDict_ClearKeysVersion(PyObject *mp)
+{
+    ASSERT_DICT_LOCKED(mp);
+
+    FT_ATOMIC_STORE_UINT32_RELAXED(((PyDictObject *)mp)->ma_keys->dk_version, 0);
+}
+
 Py_ssize_t
 _PyDict_SizeOf(PyDictObject *mp)
 {

Objects/moduleobject.c

@@ -1036,6 +1036,23 @@ _PyModule_IsPossiblyShadowing(PyObject *origin)
     return result;
 }
 
+int
+_PyModule_ReplaceLazyValue(PyObject *dict, PyObject *name, PyObject *value)
+{
+    // The adaptive interpreter uses the dictionary version to return the slot at
+    // a given index from the module. When replacing a value the version number doesn't
+    // change, so we need to atomically clear the version before replacing so that it
+    // doesn't return a lazy value.
+    int err;
+    Py_BEGIN_CRITICAL_SECTION(dict);
+
+    _PyDict_ClearKeysVersion(dict);
+    err = _PyDict_SetItem_LockHeld((PyDictObject *)dict, name, value);
+
+    Py_END_CRITICAL_SECTION();
+    return err;
+}
+
 PyObject*
 _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress)
 {
@@ -1052,13 +1069,13 @@ _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress)
                     // instead treated as if the attribute doesn't exist.
                     PyErr_Clear();
                 }
-
-                return NULL;
-            } else if (PyDict_SetItem(m->md_dict, name, new_value) < 0) {
-                Py_DECREF(new_value);
                 Py_DECREF(attr);
                 return NULL;
             }
+
+            if (_PyModule_ReplaceLazyValue(m->md_dict, name, new_value) < 0) {
+                Py_CLEAR(new_value);
+            }
             Py_DECREF(attr);
             return new_value;
         }
@@ -1490,7 +1507,7 @@ module_get_dict(PyObject *mod, void *Py_UNUSED(ignored))
                     if (new_value == NULL) {
                         return NULL;
                     }
-                    if (PyDict_SetItem(self->md_dict, key, new_value) < 0) {
+                    if (_PyModule_ReplaceLazyValue(self->md_dict, key, new_value) < 0) {
                         Py_DECREF(new_value);
                         return NULL;
                     }

Python/bytecodes.c

@@ -1799,12 +1799,13 @@ dummy_func(
             ERROR_IF(v_o == NULL);
             if (PyLazyImport_CheckExact(v_o)) {
                 PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, v_o);
-                if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
+                Py_DECREF(v_o);
+                ERROR_IF(l_v == NULL);
+                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                if (err < 0) {
                     JUMP_TO_LABEL(error);
                 }
-                Py_DECREF(v_o);
                 v_o = l_v;
-                ERROR_IF(v_o == NULL);
             }
 
             v = PyStackRef_FromPyObjectSteal(v_o);
@@ -1838,13 +1839,14 @@ dummy_func(
             PyObject *res_o = PyStackRef_AsPyObjectBorrow(*res);
             if (PyLazyImport_CheckExact(res_o)) {
                 PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, res_o);
-                if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
+                PyStackRef_CLOSE(res[0]);
+                ERROR_IF(l_v == NULL);
+                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                if (err < 0) {
+                    Py_DECREF(l_v);
                     JUMP_TO_LABEL(error);
                 }
-                res_o = l_v;
+                *res = PyStackRef_FromPyObjectSteal(l_v);
-                PyStackRef_CLOSE(res[0]);
-                ERROR_IF(res_o == NULL);
-                *res = PyStackRef_FromPyObjectSteal(res_o);
             }
         }
 

Python/ceval.c

@@ -3038,6 +3038,9 @@ check_lazy_import_comatibility(PyThreadState *tstate, PyObject *globals,
      // Check if this module should be imported lazily due to the compatbility mode support via
     // __lazy_modules__.
     PyObject *lazy_modules = NULL;
+    PyObject *abs_name = NULL;
+    int res = -1;
+
     if (globals != NULL &&
         PyMapping_GetOptionalItem(globals, &_Py_ID(__lazy_modules__), &lazy_modules) < 0) {
         return -1;
@@ -3047,15 +3050,19 @@ check_lazy_import_comatibility(PyThreadState *tstate, PyObject *globals,
 
     int ilevel = PyLong_AsInt(level);
     if (ilevel == -1 && _PyErr_Occurred(tstate)) {
-        return -1;
+        goto error;
     }
 
-    PyObject *abs_name = _PyImport_GetAbsName(tstate, name, globals, ilevel);
+    abs_name = _PyImport_GetAbsName(tstate, name, globals, ilevel);
     if (abs_name == NULL) {
-        return -1;
+        goto error;
     }
 
-    return PySequence_Contains(lazy_modules, abs_name);
+    res = PySequence_Contains(lazy_modules, abs_name);
+error:
+    Py_XDECREF(abs_name);
+    Py_XDECREF(lazy_modules);
+    return res;
 }
 
 PyObject *
@@ -3302,8 +3309,10 @@ _PyEval_LazyImportFrom(PyThreadState *tstate, PyObject *v, PyObject *name)
             PyObject *mod_dict = PyModule_GetDict(mod);
             if (mod_dict != NULL) {
                 if (PyDict_GetItemRef(mod_dict, name, &ret) < 0) {
+                    Py_DECREF(mod);
                     return NULL;
                 } else if (ret != NULL) {
+                    Py_DECREF(mod);
                     return ret;
                 }
             }

Python/executor_cases.c.h

@@ -2461,17 +2461,18 @@
             if (PyLazyImport_CheckExact(v_o)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
                 PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, v_o);
-                stack_pointer = _PyFrame_GetStackPointer(frame);
-                if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
-                    JUMP_TO_LABEL(error);
-                }
-                _PyFrame_SetStackPointer(frame, stack_pointer);
                 Py_DECREF(v_o);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
-                v_o = l_v;
+                if (l_v == NULL) {
-                if (v_o == NULL) {
                     JUMP_TO_ERROR();
                 }
+                _PyFrame_SetStackPointer(frame, stack_pointer);
+                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                stack_pointer = _PyFrame_GetStackPointer(frame);
+                if (err < 0) {
+                    JUMP_TO_LABEL(error);
+                }
+                v_o = l_v;
             }
             v = PyStackRef_FromPyObjectSteal(v_o);
             stack_pointer[0] = v;
@@ -2495,18 +2496,21 @@
             if (PyLazyImport_CheckExact(res_o)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
                 PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, res_o);
-                stack_pointer = _PyFrame_GetStackPointer(frame);
-                if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
-                    JUMP_TO_LABEL(error);
-                }
-                res_o = l_v;
-                _PyFrame_SetStackPointer(frame, stack_pointer);
                 PyStackRef_CLOSE(res[0]);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
-                if (res_o == NULL) {
+                if (l_v == NULL) {
                     JUMP_TO_ERROR();
                 }
-                *res = PyStackRef_FromPyObjectSteal(res_o);
+                _PyFrame_SetStackPointer(frame, stack_pointer);
+                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                stack_pointer = _PyFrame_GetStackPointer(frame);
+                if (err < 0) {
+                    _PyFrame_SetStackPointer(frame, stack_pointer);
+                    Py_DECREF(l_v);
+                    stack_pointer = _PyFrame_GetStackPointer(frame);
+                    JUMP_TO_LABEL(error);
+                }
+                *res = PyStackRef_FromPyObjectSteal(l_v);
             }
             stack_pointer += 1;
             assert(WITHIN_STACK_BOUNDS());

Python/generated_cases.c.h

@@ -9193,18 +9193,21 @@
                 if (PyLazyImport_CheckExact(res_o)) {
                     _PyFrame_SetStackPointer(frame, stack_pointer);
                     PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, res_o);
-                    stack_pointer = _PyFrame_GetStackPointer(frame);
-                    if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
-                        JUMP_TO_LABEL(error);
-                    }
-                    res_o = l_v;
-                    _PyFrame_SetStackPointer(frame, stack_pointer);
                     PyStackRef_CLOSE(res[0]);
                     stack_pointer = _PyFrame_GetStackPointer(frame);
-                    if (res_o == NULL) {
+                    if (l_v == NULL) {
                         JUMP_TO_LABEL(error);
                     }
-                    *res = PyStackRef_FromPyObjectSteal(res_o);
+                    _PyFrame_SetStackPointer(frame, stack_pointer);
+                    int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                    stack_pointer = _PyFrame_GetStackPointer(frame);
+                    if (err < 0) {
+                        _PyFrame_SetStackPointer(frame, stack_pointer);
+                        Py_DECREF(l_v);
+                        stack_pointer = _PyFrame_GetStackPointer(frame);
+                        JUMP_TO_LABEL(error);
+                    }
+                    *res = PyStackRef_FromPyObjectSteal(l_v);
                 }
             }
             // _PUSH_NULL_CONDITIONAL
@@ -9409,17 +9412,18 @@
             if (PyLazyImport_CheckExact(v_o)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
                 PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, v_o);
+                Py_DECREF(v_o);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
-                if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
+                if (l_v == NULL) {
                     JUMP_TO_LABEL(error);
                 }
                 _PyFrame_SetStackPointer(frame, stack_pointer);
-                Py_DECREF(v_o);
+                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
-                v_o = l_v;
+                if (err < 0) {
-                if (v_o == NULL) {
                     JUMP_TO_LABEL(error);
                 }
+                v_o = l_v;
             }
             v = PyStackRef_FromPyObjectSteal(v_o);
             stack_pointer[0] = v;

Python/import.c

@@ -3706,6 +3706,7 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
 {
     PyObject *obj = NULL;
     PyObject *fromlist = NULL;
+    PyObject *import_func = NULL;
     assert(lazy_import != NULL);
     assert(PyLazyImport_CheckExact(lazy_import));
 
@@ -3723,7 +3724,7 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
     }
 
     int is_loading = PySet_Contains(importing, lazy_import);
-    if (is_loading < 0 ) {
+    if (is_loading < 0) {
         return NULL;
     } else if (is_loading == 1) {
         PyObject *name = _PyLazyImport_GetName(lazy_import);
@@ -3766,9 +3767,8 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
 
     PyObject *globals = PyEval_GetGlobals();
 
-    PyObject *import_func;
     if (PyMapping_GetOptionalItem(lz->lz_builtins, &_Py_ID(__import__), &import_func) < 0) {
-        return NULL;
+        goto error;
     }
     if (full) {
         obj = _PyEval_ImportNameWithImport(tstate,
@@ -3781,7 +3781,6 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
     } else {
         PyObject *name = PyUnicode_Substring(lz->lz_from, 0, dot);
         if (name == NULL) {
-            Py_DECREF(import_func);
             goto error;
         }
         obj = _PyEval_ImportNameWithImport(tstate,
@@ -3793,7 +3792,6 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
                                    _PyLong_GetZero());
         Py_DECREF(name);
     }
-    Py_DECREF(import_func);
     if (obj == NULL) {
         goto error;
     }
@@ -3891,6 +3889,7 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
     }
 
     Py_XDECREF(fromlist);
+    Py_XDECREF(import_func);
     return obj;
 }
 
@@ -4226,7 +4225,7 @@ _PyImport_LazyImportModuleLevelObject(PyThreadState *tstate,
     }
 
     PyInterpreterState *interp = tstate->interp;
-    assert(frame->f_globals == frame->f_locals); // should only be called in global scope
+    assert(_PyEval_GetFrame()->f_globals == _PyEval_GetFrame()->f_locals); // should only be called in global scope
 
     // Check if the filter disables the lazy import
     PyObject *filter = LAZY_IMPORTS_FILTER(interp);
@@ -5312,7 +5311,7 @@ _imp__set_lazy_attributes_impl(PyObject *module, PyObject *child_module,
                 goto error;
             }
 
-            if (PyDict_SetItem(child_dict, attr_name, lazy_module_attr) < 0) {
+            if (_PyModule_ReplaceLazyValue(child_dict, attr_name, lazy_module_attr) < 0) {
                 Py_DECREF(lazy_module_attr);
                 goto error;
             }

Python/specialize.c

@@ -8,6 +8,7 @@
 #include "pycore_dict.h"          // DICT_KEYS_UNICODE
 #include "pycore_function.h"      // _PyFunction_GetVersionForCurrentState()
 #include "pycore_interpframe.h"   // FRAME_SPECIALS_SIZE
+#include "pycore_lazyimportobject.h" // PyLazyImport_CheckExact
 #include "pycore_list.h"          // _PyListIterObject
 #include "pycore_long.h"          // _PyLong_IsNonNegativeCompact()
 #include "pycore_moduleobject.h"
@@ -541,6 +542,7 @@ _PyCode_Quicken(_Py_CODEUNIT *instructions, Py_ssize_t size, int enable_counters
 #define SPEC_FAIL_ATTR_BUILTIN_CLASS_METHOD 22
 #define SPEC_FAIL_ATTR_CLASS_METHOD_OBJ 23
 #define SPEC_FAIL_ATTR_OBJECT_SLOT 24
+#define SPEC_FAIL_ATTR_MODULE_LAZY_VALUE 25
 
 #define SPEC_FAIL_ATTR_INSTANCE_ATTRIBUTE 26
 #define SPEC_FAIL_ATTR_METACLASS_ATTRIBUTE 27
@@ -797,6 +799,14 @@ specialize_module_load_attr_lock_held(PyDictObject *dict, _Py_CODEUNIT *instr, P
         SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_OUT_OF_VERSIONS);
         return -1;
     }
+    PyObject *value = NULL;
+    if (PyDict_GetItemRef((PyObject *)dict, name, &value) < 0 ||
+        (value != NULL && PyLazyImport_CheckExact(value))) {
+        Py_XDECREF(value);
+        SPECIALIZATION_FAIL(SPEC_FAIL_ATTR_MODULE_LAZY_VALUE, SPEC_FAIL_OUT_OF_VERSIONS);
+        return -1;
+    }
+    Py_XDECREF(value);
     write_u32(cache->version, keys_version);
     cache->index = (uint16_t)index;
     specialize(instr, LOAD_ATTR_MODULE);