From a952fe11516bd1f5fc1d756305f33936a012ac41 Mon Sep 17 00:00:00 2001
From: Andrew Leech <andrew.leech@planetinnovation.com.au>
Date: Wed, 23 Oct 2024 11:41:46 +1100
Subject: [PATCH] py/modsys: Extend atexit to hold multiple functions.

Signed-off-by: Andrew Leech <andrew@alelec.net>
---
 py/modsys.c                  | 37 ++++++++++++++++++++++++++++++++----
 py/mpconfig.h                |  2 +-
 py/runtime.h                 |  4 ++++
 tests/misc/sys_atexit.py     |  5 +++++
 tests/misc/sys_atexit.py.exp |  1 +
 5 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/py/modsys.c b/py/modsys.c
index e90ea2233a732..e2e0263ea23e0 100644
--- a/py/modsys.c
+++ b/py/modsys.c
@@ -208,13 +208,42 @@ static MP_DEFINE_CONST_FUN_OBJ_1(mp_sys_getsizeof_obj, mp_sys_getsizeof);
 #endif
 
 #if MICROPY_PY_SYS_ATEXIT
+typedef struct _m_atexit_node_t {
+    struct _m_atexit_node_t *prev;
+    struct _m_atexit_node_t *next;
+    mp_obj_t fn;
+} m_atexit_node_t;
+
 // atexit(callback): Callback is called when sys.exit is called.
 static mp_obj_t mp_sys_atexit(mp_obj_t obj) {
-    mp_obj_t old = MP_STATE_VM(sys_exitfunc);
-    MP_STATE_VM(sys_exitfunc) = obj;
-    return old;
+    m_atexit_node_t *node = m_malloc(sizeof(m_atexit_node_t));
+    if (MP_STATE_VM(sys_exitfunc) != mp_const_none) {
+        MP_STATE_VM(sys_exitfunc)->prev = node;
+    }
+    node->fn = obj;
+    node->prev = NULL;
+    node->next = MP_STATE_VM(sys_exitfunc);
+    MP_STATE_VM(sys_exitfunc) = node;
+    return obj;
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(mp_sys_atexit_obj, mp_sys_atexit);
+
+void mp_sys_atexit_execute(void) {
+    // walk through the linked list and execute each function
+    // Beware, the sys.settrace callback should be disabled before running sys.atexit.
+    while (MP_STATE_VM(sys_exitfunc) != mp_const_none) {
+        if (mp_obj_is_callable(MP_STATE_VM(sys_exitfunc)->fn)) {
+            nlr_buf_t nlr;
+            if (nlr_push(&nlr) == 0) {
+                mp_call_function_0(MP_STATE_VM(sys_exitfunc)->fn);
+            } else {
+                // Uncaught exception: print it out.
+                mp_obj_print_exception(&mp_plat_print, (mp_obj_t)nlr.ret_val);
+            }
+        }
+        MP_STATE_VM(sys_exitfunc) = MP_STATE_VM(sys_exitfunc)->next;
+    }
+}
 #endif
 
 #if MICROPY_PY_SYS_SETTRACE
@@ -364,7 +393,7 @@ MP_REGISTER_ROOT_POINTER(mp_obj_base_t * cur_exception);
 
 #if MICROPY_PY_SYS_ATEXIT
 // exposed through sys.atexit function
-MP_REGISTER_ROOT_POINTER(mp_obj_t sys_exitfunc);
+MP_REGISTER_ROOT_POINTER(struct _m_atexit_node_t *sys_exitfunc);
 #endif
 
 #if MICROPY_PY_SYS_ATTR_DELEGATION
diff --git a/py/mpconfig.h b/py/mpconfig.h
index 34eafa9e5debb..342f0966284b6 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -1488,7 +1488,7 @@ typedef double mp_float_t;
 
 // Whether to provide "sys.atexit" function (MicroPython extension)
 #ifndef MICROPY_PY_SYS_ATEXIT
-#define MICROPY_PY_SYS_ATEXIT (0)
+#define MICROPY_PY_SYS_ATEXIT (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
 #endif
 
 // Whether to provide the "sys.path" attribute (which forces module delegation
diff --git a/py/runtime.h b/py/runtime.h
index e8e5a758f8beb..37f49d4fb35ea 100644
--- a/py/runtime.h
+++ b/py/runtime.h
@@ -95,6 +95,10 @@ extern const byte mp_binary_op_method_name[];
 void mp_init(void);
 void mp_deinit(void);
 
+#if MICROPY_PY_SYS_ATEXIT
+void mp_sys_atexit_execute(void);
+#endif
+
 void mp_sched_exception(mp_obj_t exc);
 void mp_sched_keyboard_interrupt(void);
 #if MICROPY_ENABLE_VM_ABORT
diff --git a/tests/misc/sys_atexit.py b/tests/misc/sys_atexit.py
index e9c5693f975e3..86f3552a4245c 100644
--- a/tests/misc/sys_atexit.py
+++ b/tests/misc/sys_atexit.py
@@ -15,6 +15,11 @@ def do_at_exit():
     print("done at exit:", some_var)
 
 
+@sys.atexit
+def do_at_exit_2():
+    print("done at exit last")
+
+
 sys.atexit(do_at_exit)
 
 some_var = "ok"
diff --git a/tests/misc/sys_atexit.py.exp b/tests/misc/sys_atexit.py.exp
index 3cbdae9a5a0aa..5965ca541549d 100644
--- a/tests/misc/sys_atexit.py.exp
+++ b/tests/misc/sys_atexit.py.exp
@@ -1,2 +1,3 @@
 done before exit
 done at exit: ok
+done at exit last