py/scheduler: Schedule SystemExit or SystemAbort.

laurensvalk · laurensvalk · commit c8072e4a8d01 · 2022-11-17T11:11:19.000+01:00
This implements mp_sched_system_exit_or_abort(), which can be used to schedule a SystemExit or SystemAbort on the main thread. This can be used by async sources to exit MicroPython if requested by a firmware process. Use abort=true if the MicroPython script must not be able to catch it, not even with a bare except.

The implementation is similar to mp_sched_keyboard_interrupt().

Signed-off-by: Laurens Valk &lt;laurens@pybricks.com&gt;
diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -659,7 +659,8 @@
 #endif
 
 // Whether to provide the SystemAbort exception, used to exit MicroPython
-// even if user code catches all exceptions.
+// even if user code catches all exceptions. This also provides
+// mp_sched_system_exit_or_abort() to schedule the mp_system_exception object.
 #ifndef MICROPY_ENABLE_SYSTEM_ABORT
 #define MICROPY_ENABLE_SYSTEM_ABORT (0)
 #endif
diff --git a/py/mpstate.h b/py/mpstate.h
@@ -165,6 +165,11 @@ typedef struct _mp_state_vm_t {
     mp_obj_exception_t mp_kbd_exception;
     #endif
 
+    #if MICROPY_ENABLE_SYSTEM_ABORT
+    // exception object of type SystemExit or SystemAbort
+    mp_obj_exception_t mp_system_exception;
+    #endif
+
     // dictionary with loaded modules (may be exposed as sys.modules)
     mp_obj_dict_t mp_loaded_modules_dict;
 
diff --git a/py/objexcept.c b/py/objexcept.c
@@ -377,8 +377,9 @@ MP_DEFINE_EXCEPTION(Exception, BaseException)
     */
 
 #if MICROPY_ENABLE_SYSTEM_ABORT
-// Exception that can't be raised or caught by user code, not even with a
-// bare except. Can be used to force MicroPython to exit from C code.
+// Exception that can't be raised or caught by user code, not even with a bare
+// except. Do not raise directly but use mp_sched_system_exit_or_abort(true) to
+// schedule it on the main thread.
 MP_DEFINE_EXCEPTION(SystemAbort, BaseException)
 #endif
 
diff --git a/py/runtime.c b/py/runtime.c
@@ -95,6 +95,14 @@ void mp_init(void) {
     MP_STATE_VM(mp_kbd_exception).args = (mp_obj_tuple_t *)&mp_const_empty_tuple_obj;
     #endif
 
+    #if MICROPY_ENABLE_SYSTEM_ABORT
+    // initialise the exception object for raising SystemExit or SystemAbort
+    MP_STATE_VM(mp_system_exception).traceback_alloc = 0;
+    MP_STATE_VM(mp_system_exception).traceback_len = 0;
+    MP_STATE_VM(mp_system_exception).traceback_data = NULL;
+    MP_STATE_VM(mp_system_exception).args = (mp_obj_tuple_t *)&mp_const_empty_tuple_obj;
+    #endif
+
     #if MICROPY_ENABLE_COMPILER
     // optimization disabled by default
     MP_STATE_VM(mp_optimise_value) = 0;
diff --git a/py/runtime.h b/py/runtime.h
@@ -75,6 +75,9 @@ void mp_deinit(void);
 
 void mp_sched_exception(mp_obj_t exc);
 void mp_sched_keyboard_interrupt(void);
+#if MICROPY_ENABLE_SYSTEM_ABORT
+void mp_sched_system_exit_or_abort(bool abort);
+#endif
 void mp_handle_pending(bool raise_exc);
 
 #if MICROPY_ENABLE_SCHEDULER
diff --git a/py/scheduler.c b/py/scheduler.c
@@ -51,6 +51,17 @@ void MICROPY_WRAP_MP_SCHED_KEYBOARD_INTERRUPT(mp_sched_keyboard_interrupt)(void)
 }
 #endif
 
+#if MICROPY_ENABLE_SYSTEM_ABORT
+// Schedules SystemExit or SystemAbort on the main thread. Can be used by
+// async sources to exit MicroPython. Use abort=true if the MicroPython script
+// must not be able to catch it, not even with a bare except.
+void MICROPY_WRAP_MP_SCHED_EXCEPTION(mp_sched_system_exit_or_abort)(bool abort) {
+    MP_STATE_VM(mp_system_exception).base.type = abort ?
+        &mp_type_SystemAbort : &mp_type_SystemExit;
+    mp_sched_exception(MP_OBJ_FROM_PTR(&MP_STATE_VM(mp_system_exception)));
+}
+#endif
+
 #if MICROPY_ENABLE_SCHEDULER
 
 #define IDX_MASK(i) ((i) & (MICROPY_SCHEDULER_DEPTH - 1))
