py/modmicropython: Add heap_locked function to test state of heap.

pi-anl · dpgeorge · commit 86bfabec1145 · 2020-03-11T16:54:16.000+11:00
This commit adds micropython.heap_locked() which returns the current
lock-depth of the heap, and can be used by Python code to check if the heap
is locked or not.  This new function is configured via
MICROPY_PY_MICROPYTHON_HEAP_LOCKED and is disabled by default.

This commit also changes the return value of micropython.heap_unlock() so
it returns the current lock-depth as well.
diff --git a/docs/library/micropython.rst b/docs/library/micropython.rst
@@ -82,17 +82,26 @@ Functions
 
 .. function:: heap_lock()
 .. function:: heap_unlock()
+.. function:: heap_locked()
 
    Lock or unlock the heap.  When locked no memory allocation can occur and a
    `MemoryError` will be raised if any heap allocation is attempted.
+   `heap_locked()` returns a true value if the heap is currently locked.
 
    These functions can be nested, ie `heap_lock()` can be called multiple times
    in a row and the lock-depth will increase, and then `heap_unlock()` must be
    called the same number of times to make the heap available again.
 
+   Both `heap_unlock()` and `heap_locked()` return the current lock depth
+   (after unlocking for the former) as a non-negative integer, with 0 meaning
+   the heap is not locked.
+
    If the REPL becomes active with the heap locked then it will be forcefully
    unlocked.
 
+   Note: `heap_locked()` is not enabled on most ports by default,
+   requires `MICROPY_PY_MICROPYTHON_HEAP_LOCKED`.
+
 .. function:: kbd_intr(chr)
 
    Set the character that will raise a `KeyboardInterrupt` exception.  By
diff --git a/ports/unix/variants/coverage/mpconfigvariant.h b/ports/unix/variants/coverage/mpconfigvariant.h
@@ -59,6 +59,7 @@
 #define MICROPY_PY_COLLECTIONS_NAMEDTUPLE__ASDICT (1)
 #define MICROPY_PY_UCRYPTOLIB          (1)
 #define MICROPY_PY_UCRYPTOLIB_CTR      (1)
+#define MICROPY_PY_MICROPYTHON_HEAP_LOCKED (1)
 
 // TODO these should be generic, not bound to fatfs
 #define mp_type_fileio mp_type_vfs_posix_fileio
diff --git a/py/modmicropython.c b/py/modmicropython.c
@@ -130,9 +130,16 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_0(mp_micropython_heap_lock_obj, mp_micropython_he
 
 STATIC mp_obj_t mp_micropython_heap_unlock(void) {
     gc_unlock();
-    return mp_const_none;
+    return MP_OBJ_NEW_SMALL_INT(MP_STATE_MEM(gc_lock_depth));
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(mp_micropython_heap_unlock_obj, mp_micropython_heap_unlock);
+
+#if MICROPY_PY_MICROPYTHON_HEAP_LOCKED
+STATIC mp_obj_t mp_micropython_heap_locked(void) {
+    return MP_OBJ_NEW_SMALL_INT(MP_STATE_MEM(gc_lock_depth));
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(mp_micropython_heap_locked_obj, mp_micropython_heap_locked);
+#endif
 #endif
 
 #if MICROPY_ENABLE_EMERGENCY_EXCEPTION_BUF && (MICROPY_EMERGENCY_EXCEPTION_BUF_SIZE == 0)
@@ -184,6 +191,9 @@ STATIC const mp_rom_map_elem_t mp_module_micropython_globals_table[] = {
     #if MICROPY_ENABLE_GC
     { MP_ROM_QSTR(MP_QSTR_heap_lock), MP_ROM_PTR(&mp_micropython_heap_lock_obj) },
     { MP_ROM_QSTR(MP_QSTR_heap_unlock), MP_ROM_PTR(&mp_micropython_heap_unlock_obj) },
+    #if MICROPY_PY_MICROPYTHON_HEAP_LOCKED
+    { MP_ROM_QSTR(MP_QSTR_heap_locked), MP_ROM_PTR(&mp_micropython_heap_locked_obj) },
+    #endif
     #endif
     #if MICROPY_KBD_EXCEPTION
     { MP_ROM_QSTR(MP_QSTR_kbd_intr), MP_ROM_PTR(&mp_micropython_kbd_intr_obj) },
diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -1069,6 +1069,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_MICROPYTHON_STACK_USE (MICROPY_PY_MICROPYTHON_MEM_INFO)
 #endif
 
+// Whether to provide the "micropython.heap_locked" function
+#ifndef MICROPY_PY_MICROPYTHON_HEAP_LOCKED
+#define MICROPY_PY_MICROPYTHON_HEAP_LOCKED (0)
+#endif
+
 // Whether to provide "array" module. Note that large chunk of the
 // underlying code is shared with "bytearray" builtin type, so to
 // get real savings, it should be disabled too.
diff --git a/tests/micropython/heap_lock.py b/tests/micropython/heap_lock.py
@@ -5,6 +5,7 @@
 l = []
 l2 = list(range(100))
 
+micropython.heap_lock()
 micropython.heap_lock()
 
 # general allocation on the heap
@@ -19,6 +20,14 @@
 except MemoryError:
     print('MemoryError')
 
+print(micropython.heap_unlock())
+
+# Should still fail
+try:
+    print([])
+except MemoryError:
+    print('MemoryError')
+
 micropython.heap_unlock()
 
 # check that allocation works after an unlock
diff --git a/tests/micropython/heap_lock.py.exp b/tests/micropython/heap_lock.py.exp
@@ -1,3 +1,5 @@
 MemoryError
 MemoryError
+1
+MemoryError
 []
diff --git a/tests/micropython/heap_locked.py b/tests/micropython/heap_locked.py
@@ -0,0 +1,12 @@
+# test micropython.heap_locked()
+
+import micropython
+
+if not hasattr(micropython, "heap_locked"):
+    print("SKIP")
+    raise SystemExit
+
+micropython.heap_lock()
+print(micropython.heap_locked())
+micropython.heap_unlock()
+print(micropython.heap_locked())
diff --git a/tests/micropython/heap_locked.py.exp b/tests/micropython/heap_locked.py.exp
@@ -0,0 +1,2 @@
+1
+0
