micropython · dpgeorge · May 18, 2020 · May 18, 2020 · May 18, 2020 · May 18, 2020
diff --git a/docs/library/micropython.rst b/docs/library/micropython.rst
@@ -113,6 +113,14 @@ Functions
    incoming stream of characters that is usually used for the REPL, in case
    that stream is used for other purposes.
 
+.. function:: scheduler_lock()
+
+   Lock the scheduler.  May be nested.
+
+.. function:: scheduler_unlock()
+
+   Unlock the scheduler.  May be nested.
+
 .. function:: schedule(func, arg)
 
    Schedule the function *func* to be executed "very soon".  The function

diff --git a/extmod/moduselect.c b/extmod/moduselect.c
@@ -44,6 +44,12 @@
 ///
 /// This module provides the select function.
 
+#define MICROPY_PY_USELECT_NOTIFIER (1)
+
+#if MICROPY_PY_USELECT_NOTIFIER
+STATIC mp_obj_t notifier_new(void);
+#endif
+
 typedef struct _poll_obj_t {
     mp_obj_t obj;
     mp_uint_t (*ioctl)(mp_obj_t obj, mp_uint_t request, uintptr_t arg, int *errcode);
@@ -199,6 +205,13 @@ STATIC mp_obj_t poll_register(size_t n_args, const mp_obj_t *args) {
     } else {
         flags = MP_STREAM_POLL_RD | MP_STREAM_POLL_WR;
     }
+    #if MICROPY_PY_USELECT_NOTIFIER
+    if (args[1] == mp_const_none) {
+        mp_obj_t s = notifier_new();
+        poll_map_add(&self->poll_map, &s, 1, flags, false);
+        return s;
+    }
+    #endif
     poll_map_add(&self->poll_map, &args[1], 1, flags, false);
     return mp_const_none;
 }
@@ -376,4 +389,68 @@ const mp_obj_module_t mp_module_uselect = {
     .globals = (mp_obj_dict_t *)&mp_module_select_globals,
 };
 
+#if MICROPY_PY_USELECT_NOTIFIER
+
+typedef struct _mp_obj_notifier_t {
+    mp_obj_base_t base;
+    volatile size_t avail;
+} mp_obj_notifier_t;
+
+STATIC mp_uint_t notifier_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *errcode) {
+    mp_obj_notifier_t *self = MP_OBJ_TO_PTR(self_in);
+    switch (request) {
+        case MP_STREAM_POLL: {
+            uintptr_t flags = arg;
+            mp_uint_t ret = MP_STREAM_POLL_WR & flags;
+            if (self->avail) {
+                ret |= MP_STREAM_POLL_RD & flags;
+            }
+            return ret;
+        }
+        default:
+            *errcode = MP_EINVAL;
+            return MP_STREAM_ERROR;
+    }
+}
+
+STATIC mp_obj_t notifier_set(mp_obj_t self_in) {
+    mp_obj_notifier_t *self = MP_OBJ_TO_PTR(self_in);
+    self->avail = 1;
+    return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(notifier_set_obj, notifier_set);
+
+STATIC mp_obj_t notifier_clear(mp_obj_t self_in) {
+    mp_obj_notifier_t *self = MP_OBJ_TO_PTR(self_in);
+    self->avail = 0;
+    return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(notifier_clear_obj, notifier_clear);
+
+STATIC const mp_rom_map_elem_t notifier_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR_set), MP_ROM_PTR(&notifier_set_obj) },
+    { MP_ROM_QSTR(MP_QSTR_clear), MP_ROM_PTR(&notifier_clear_obj) },
+};
+STATIC MP_DEFINE_CONST_DICT(notifier_locals_dict, notifier_locals_dict_table);
+
+STATIC const mp_stream_p_t notifier_stream_p = {
+    .ioctl = notifier_ioctl,
+};
+
+STATIC const mp_obj_type_t mp_type_notifier = {
+    { &mp_type_type },
+    .name = MP_QSTR_notifier,
+    .protocol = &notifier_stream_p,
+    .locals_dict = (mp_obj_dict_t *)&notifier_locals_dict,
+};
+
+STATIC mp_obj_t notifier_new(void) {
+    mp_obj_notifier_t *self = m_new_obj(mp_obj_notifier_t);
+    self->base.type = &mp_type_notifier;
+    self->avail = 0;
+    return MP_OBJ_FROM_PTR(self);
+}
+
+#endif // MICROPY_PY_USELECT_NOTIFIER
+
 #endif // MICROPY_PY_USELECT
diff --git a/extmod/uasyncio/core.py b/extmod/uasyncio/core.py
@@ -70,6 +70,7 @@ class IOQueue:
     def __init__(self):
         self.poller = select.poll()
         self.map = {}  # maps id(stream) to [task_waiting_read, task_waiting_write, stream]
+        self.notify = self.poller.register(None, select.POLLIN)
 
     def _enqueue(self, s, idx):
         if id(s) not in self.map:
@@ -111,6 +112,9 @@ def remove(self, task):
 
     def wait_io_event(self, dt):
         for s, ev in self.poller.ipoll(dt):
+            if s is self.notify:
+                s.clear()
+                continue
             sm = self.map[id(s)]
             # print('poll', s, sm, ev)
             if ev & ~select.POLLOUT and sm[0] is not None:
@@ -160,9 +164,6 @@ def run_until_complete(main_task=None):
             if t:
                 # A task waiting on _task_queue; "ph_key" is time to schedule task at
                 dt = max(0, ticks_diff(t.ph_key, ticks()))
-            elif not _io_queue.map:
-                # No tasks can be woken so finished running
-                return
             # print('(poll {})'.format(dt), len(_io_queue.map))
             _io_queue.wait_io_event(dt)
 

diff --git a/extmod/uasyncio/event.py b/extmod/uasyncio/event.py
@@ -1,9 +1,11 @@
 # MicroPython uasyncio module
 # MIT license; Copyright (c) 2019-2020 Damien P. George
 
+from micropython import scheduler_lock, scheduler_unlock
 from . import core
 
 # Event class for primitive events that can be waited on, set, and cleared
+# The following methods are safe to call from a scheduled callback: is_set, set, clear
 class Event:
     def __init__(self):
         self.state = False  # False=unset; True=set
@@ -13,19 +15,32 @@ def is_set(self):
         return self.state
 
     def set(self):
+        if self.state:
+            return
         # Event becomes set, schedule any tasks waiting on it
+        scheduler_lock()
+        signal = False
         while self.waiting.peek():
             core._task_queue.push_head(self.waiting.pop_head())
+            signal = True
+        if signal:
+            # signal poll to finish
+            core._io_queue.notify.set()
         self.state = True
+        scheduler_unlock()
 
     def clear(self):
         self.state = False
 
     async def wait(self):
+        scheduler_lock()
         if not self.state:
             # Event not set, put the calling task on the event's waiting queue
             self.waiting.push_head(core.cur_task)
             # Set calling task's data to the event's queue so it can be removed if needed
                core.cur_task.data = self.waiting
+            scheduler_unlock()
             yield
+        else:
+            scheduler_unlock()
         return True
diff --git a/extmod/uasyncio/task.py b/extmod/uasyncio/task.py
@@ -4,6 +4,7 @@
 # This file contains the core TaskQueue based on a pairing heap, and the core Task class.
 # They can optionally be replaced by C implementations.
 
+import micropython
 from . import core
 
 
@@ -152,6 +153,7 @@ def cancel(self):
         # Can't cancel self (not supported yet).
         if self is core.cur_task:
             raise RuntimeError("can't cancel self")
+        micropython.scheduler_lock()
         # If Task waits on another task then forward the cancel to the one it's waiting on.
         while isinstance(self.data, Task):
             self = self.data
@@ -165,4 +167,5 @@ def cancel(self):
             core._task_queue.remove(self)
             core._task_queue.push_head(self)
         self.data = core.CancelledError
+        micropython.scheduler_unlock()
         return True
diff --git a/ports/stm32/mpconfigport.h b/ports/stm32/mpconfigport.h
@@ -376,31 +376,46 @@ static inline mp_uint_t disable_irq(void) {
 #define MICROPY_END_ATOMIC_SECTION(state)  enable_irq(state)
 
 #if MICROPY_PY_THREAD
-#define MICROPY_EVENT_POLL_HOOK \
+#define MICROPY_EVENT_WAIT_ATOMIC \
     do { \
-        extern void mp_handle_pending(bool); \
-        mp_handle_pending(true); \
-        if (pyb_thread_enabled) { \
-            MP_THREAD_GIL_EXIT(); \
-            pyb_thread_yield(); \
-            MP_THREAD_GIL_ENTER(); \
+        mp_uint_t atomic_state = MICROPY_BEGIN_ATOMIC_SECTION(); \
+        if (!mp_sched_any_pending()) { \
+            if (pyb_thread_enabled) { \
+                MICROPY_END_ATOMIC_SECTION(atomic_state); \
+                MP_THREAD_GIL_EXIT(); \
+                pyb_thread_yield(); \
+                MP_THREAD_GIL_ENTER(); \
+            } else { \
+                __WFI(); \
+                MICROPY_END_ATOMIC_SECTION(atomic_state); \
+            } \
         } else { \
-            __WFI(); \
+            MICROPY_END_ATOMIC_SECTION(atomic_state); \
         } \
-    } while (0);
+    } while (0)
 
 #define MICROPY_THREAD_YIELD() pyb_thread_yield()
 #else
+
+#define MICROPY_EVENT_WAIT_ATOMIC \
+    do { \
+        mp_uint_t atomic_state = MICROPY_BEGIN_ATOMIC_SECTION(); \
+        if (!mp_sched_any_pending()) { \
+            __WFI(); \
+        } \
+        MICROPY_END_ATOMIC_SECTION(atomic_state); \
+    } while (0)
+
+#define MICROPY_THREAD_YIELD()
+#endif
+
 #define MICROPY_EVENT_POLL_HOOK \
     do { \
+        MICROPY_EVENT_WAIT_ATOMIC; \
         extern void mp_handle_pending(bool); \
         mp_handle_pending(true); \
-        __WFI(); \
     } while (0);
 
-#define MICROPY_THREAD_YIELD()
-#endif
-
 // The LwIP interface must run at a raised IRQ priority
 #define MICROPY_PY_LWIP_ENTER   uint32_t atomic_state = raise_irq_pri(IRQ_PRI_PENDSV);
 #define MICROPY_PY_LWIP_REENTER atomic_state = raise_irq_pri(IRQ_PRI_PENDSV);

diff --git a/ports/stm32/usbd_hid_interface.c b/ports/stm32/usbd_hid_interface.c
@@ -26,7 +26,7 @@
 
 #include "usbd_hid_interface.h"
 
-#include "py/mpstate.h"
+#include "py/runtime.h"
 #include "py/mperrno.h"
 #include "py/mphal.h"
 #include "usb.h"