Make sure that mngr.__exit__() is always called in a with statement, even if there is an interrupt during the call to mngr.__enter__()

Author: markshannon

Lib/test/test_with.py

@@ -10,6 +10,7 @@
 import unittest
 from collections import deque
 from contextlib import _GeneratorContextManager, contextmanager, nullcontext
+from _testinternalcapi import SelfInterruptingContextManager
 
 
 def do_with(obj):
@@ -850,5 +851,21 @@ def exit_raises():
                                  expected)
 
 
+class InterruptDuringEnter(unittest.TestCase):
+
+    def test_exit_called_after_interrupt(self):
+        cm = SelfInterruptingContextManager()
+        self.assertFalse(cm.within())
+        try:
+            with cm:
+                self.assertTrue(cm.within())
+        except KeyboardInterrupt:
+            self.assertFalse(cm.within())
+            return
+        except:
+            self.fail("Wrong exception raised")
+        self.fail("No exception raised")
+
+
 if __name__ == '__main__':
     unittest.main()

Modules/_testinternalcapi.c

@@ -3196,6 +3196,66 @@ test_thread_state_ensure_from_view_interp_switch(PyObject *self, PyObject *unuse
     Py_RETURN_NONE;
 }
 
+/* Self interrupting context manager */
+
+typedef struct {
+    PyObject_HEAD
+    int within;
+} SelfInterruptingContextManagerObject;
+
+static PyObject *
+new_self_interrupting(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+    SelfInterruptingContextManagerObject *self =
+        (SelfInterruptingContextManagerObject *)type->tp_alloc(type, 0);
+    if (self != NULL) {
+        self->within = 0;
+    }
+    return (PyObject *)self;
+}
+
+static PyObject *
+self_interrupting_enter(PyObject *op, PyObject *Py_UNUSED(dummy))
+{
+    ((SelfInterruptingContextManagerObject *)op)->within = 1;
+    PyThreadState *tstate = PyThreadState_Get();
+    PyObject *ki = Py_NewRef(PyExc_KeyboardInterrupt);
+    PyObject *old_exc = _Py_atomic_exchange_ptr(&tstate->async_exc, ki);
+    _Py_set_eval_breaker_bit(tstate, _PY_ASYNC_EXCEPTION_BIT);
+    Py_XDECREF(old_exc);
+
+    return Py_NewRef(op);
+}
+
+static PyObject *
+self_interrupting_within(PyObject *op, PyObject *Py_UNUSED(dummy))
+{
+    return PyBool_FromLong(((SelfInterruptingContextManagerObject *)op)->within);
+}
+
+static PyObject *
+self_interrupting_exit(PyObject *op, PyObject *Py_UNUSED(args)) {
+    ((SelfInterruptingContextManagerObject *)op)->within = 0;
+    Py_RETURN_NONE;
+}
+
+static PyMethodDef self_interrupting_methods[] = {
+    {"__enter__", self_interrupting_enter, METH_NOARGS, NULL},
+    {"within", self_interrupting_within, METH_NOARGS, NULL},
+    {"__exit__", self_interrupting_exit, METH_VARARGS, NULL},
+    {NULL, NULL} /* sentinel */
+};
+
+static PyTypeObject SelfInterruptingContextManager_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "_testcapi.SelfInterruptingContextManager",
+    sizeof(SelfInterruptingContextManagerObject),
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE,
+    .tp_new = new_self_interrupting,
+    .tp_methods = self_interrupting_methods,
+};
+
+
 static PyMethodDef module_functions[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_eval_frame_stats", get_eval_frame_stats, METH_NOARGS, NULL},
@@ -3418,6 +3478,11 @@ module_exec(PyObject *module)
     }
 #endif
 
+    if (PyType_Ready(&SelfInterruptingContextManager_Type) < 0) {
+        return 1;
+    }
+    PyModule_AddObject(module, "SelfInterruptingContextManager", (PyObject *)&SelfInterruptingContextManager_Type);
+
     return 0;
 }
 

Modules/_testinternalcapi/test_cases.c.h

@@ -161,7 +161,7 @@ dummy_func(
         }
 
         replaced op(_CHECK_PERIODIC_AT_END, (--)) {
-            int err = check_periodics(tstate);
+            int err = check_periodics_at_end(tstate, frame);
             ERROR_IF(err != 0);
         }
 

Python/bytecodes.c

@@ -526,6 +526,22 @@ check_periodics(PyThreadState *tstate) {
     return 0;
 }
 
+static inline int
+check_periodics_at_end(PyThreadState *tstate, _PyInterpreterFrame *frame) {
+    _Py_CHECK_EMSCRIPTEN_SIGNALS_PERIODICALLY();
+    QSBR_QUIESCENT_STATE(tstate);
+    if (_Py_atomic_load_uintptr_relaxed(&tstate->eval_breaker) & _PY_EVAL_EVENTS_MASK) {
+        // Do not handle pending interrupts if the previous instruction was LOAD_SPECIAL
+        // This may also not handle interrupts if a cache looks like LOAD_SPECIAL,
+        // but this is benign as we won't skip periodic checks indefinitely.
+        if (frame->instr_ptr[-1].op.code == LOAD_SPECIAL) {
+            return 0;
+        }
+        return _Py_HandlePending(tstate);
+    }
+    return 0;
+}
+
 // Mark the generator as executing. Returns true if the state was changed,
 // false if it was already executing or finished.
 static inline bool