Use newer event loop implementation (based on asi2700-master branch)

Author: dbrignoli

uasyncio.core/uasyncio/core.py

@@ -1,120 +1,148 @@
-try:
-    import utime as time
-except ImportError:
-    import time
-import uheapq as heapq
-import logging
+import utime as time
+import utimeq
+import ucollections
+
+
+type_gen = type((lambda: (yield))())
+
+DEBUG = 0
+log = None
+
+def set_debug(val):
+    global DEBUG, log
+    DEBUG = val
+    if val:
+        import logging
+        log = logging.getLogger("uasyncio.core")
+        log.setLevel(val)
+
 
 class CancelledError(Exception):
     pass
 
-log = logging.getLogger("asyncio")
 
-type_gen = type((lambda: (yield))())
+class TimeoutError(CancelledError):
+    pass
+
 
 class EventLoop:
 
-    def __init__(self):
-        self.q = []
-        self.cnt = 0
+    def __init__(self, runq_len=16, waitq_len=16):
+        self.runq = ucollections.deque((), runq_len, True)
+        self.waitq = utimeq.utimeq(waitq_len)
+        # Current task being run. Task is a top-level coroutine scheduled
+        # in the event loop (sub-coroutines executed transparently by
+        # yield from/await, event loop "doesn't see" them).
+        self.cur_task = None
 
     def time(self):
-        return time.time()
+        return time.ticks_ms()
 
     def create_task(self, coro):
         # CPython 3.4.2
-        self.call_at(0, coro)
+        self.call_later_ms(0, coro)
         # CPython asyncio incompatibility: we don't return Task object
 
     def call_soon(self, callback, *args):
-        self.call_at(self.time(), callback, *args)
+        if __debug__ and DEBUG:
+            log.debug("Scheduling in runq: %s", (callback, args))
+        self.runq.append(callback)
+        if not isinstance(callback, type_gen):
+            self.runq.append(args)
 
     def call_later(self, delay, callback, *args):
-        self.call_at(self.time() + delay, callback, *args)
+        self.call_at_(time.ticks_add(self.time(), int(delay * 1000)), callback, args)
 
-    def call_at(self, time, callback, *args, exc=None):
-        # Including self.cnt is a workaround per heapq docs
-        if __debug__:
-            log.debug("Scheduling %s", (time, self.cnt, callback, args, exc))
-        heapq.heappush(self.q, (time, self.cnt, callback, args, exc, False))
-#        print(self.q)
-        self.cnt += 1
+    def call_later_ms(self, delay, callback, *args):
+        if not delay:
+            return self.call_soon(callback, *args)
+        self.call_at_(time.ticks_add(self.time(), delay), callback, args)
+
+    def call_at_(self, time, callback, args=()):
+        if __debug__ and DEBUG:
+            log.debug("Scheduling in waitq: %s", (time, callback, args))
+        self.waitq.push(time, callback, args)
 
     def wait(self, delay):
         # Default wait implementation, to be overriden in subclasses
         # with IO scheduling
-        if __debug__:
+        if __debug__ and DEBUG:
             log.debug("Sleeping for: %s", delay)
-        time.sleep(delay)
-
-    def cancel(self, callback, exc = CancelledError):
-        _id = id(callback)
-        for idx, item in enumerate(self.q):
-            t, cnt, cb, args, _exc, _discard = item
-            if id(cb) != _id:
-                continue
-            if __debug__:
-                log.debug("Setting discard flag on: %s at index %d", (t, cnt, cb, args, _exc), idx)
-            self.q[idx] = t, cnt, cb, args, _exc, True
-            self.call_at(0, cb, *args, exc=exc)
-        self.remove_polled_cb(callback)
+        time.sleep_ms(delay)
+
+    def cancel(self, coro, exc = CancelledError):
+        if isinstance(coro, type_gen):
+            try:
+                prev = coro.pend_throw(exc)
+                if __debug__ and DEBUG:
+                    log.debug("Cancelling %s asynchronously", coro)
+                if prev is False:
+                    _event_loop.remove_polled_cb(coro)
+                    _event_loop.call_soon(coro)
+            except TypeError:
+                # .pend_throw() works only on started coroutines
+                # Kill the coro right here, right now
+                # No need to worry about IO because the coro cannot be registered
+                # because it didn't get a chance to run yet
+                try:
+                    if __debug__ and DEBUG:
+                        log.debug("Cancelling %s synchronously", coro)
+                    coro.throw(exc)
+                except exc:
+                    pass
+        else:
+            raise UnimplementedError('Cancelling a callback is not supported')
 
     def run_forever(self):
+        cur_task = [0, 0, 0]
         while True:
-            if self.q:
-                tnow = self.time()
-                if __debug__:
-                    log.debug('*'*20+' sched step start at %s, num tasks in queue %d', tnow, len(self.q))
-                t, cnt, cb, args, exc, discard = heapq.heappop(self.q)
-                delay = t - tnow
-                if __debug__:
-                    log.debug("Next coroutine to run in %s: %s", delay, (t, cnt, cb, args, exc))
-                if discard:
-                    if __debug__:
-                        log.debug("Discarding: %s", (t, cnt, cb, args, exc, discard))
-                    continue
-#                __main__.mem_info()
-                if delay > 0 and not exc:
-                    self.call_at(t, cb, *args)
-                    self.wait(delay)
-                    continue
-            else:
-                self.wait(-1)
-                # Assuming IO completion scheduled some tasks
-                continue
-            # cancelled callbacks aren't called and nor rescheduled
-            if callable(cb):
-                if not exc:
+            # Expire entries in waitq and move them to runq
+            tnow = self.time()
+            while self.waitq:
+                t = self.waitq.peektime()
+                delay = time.ticks_diff(t, tnow)
+                if delay > 0:
+                    break
+                self.waitq.pop(cur_task)
+                if __debug__ and DEBUG:
+                    log.debug("Moving from waitq to runq: %s", cur_task[1])
+                self.call_soon(cur_task[1], *cur_task[2])
+
+            # Process runq
+            l = len(self.runq)
+            if __debug__ and DEBUG:
+                log.debug("Entries in runq: %d", l)
+            while l:
+                cb = self.runq.popleft()
+                l -= 1
+                args = ()
+                if not isinstance(cb, type_gen):
+                    args = self.runq.popleft()
+                    l -= 1
+                    if __debug__ and DEBUG:
+                        log.info("Next callback to run: %s", (cb, args))
                     cb(*args)
-            else:
+                    continue
+
+                if __debug__ and DEBUG:
+                    log.info("Next coroutine to run: %s", (cb, args))
+                self.cur_task = cb
                 delay = 0
                 try:
-                    if __debug__:
-                        log.debug("Coroutine %s send args: %s, %s", cb, args, exc)
-                    if exc:
-                        try:
-                            ret = cb.throw(exc)
-                        except exc:
-                            # ret == None reschedules a canceled task, next round it should raise StopIteration
-                            ret = None
-                    elif args == ():
+                    if args is ():
                         ret = next(cb)
                     else:
                         ret = cb.send(*args)
-                    if __debug__:
-                        log.debug("Coroutine %s yield result: %s", cb, ret)
+                    if __debug__ and DEBUG:
+                        log.info("Coroutine %s yield result: %s", cb, ret)
                     if isinstance(ret, SysCall1):
                         arg = ret.arg
-                        if isinstance(ret, Sleep):
+                        if isinstance(ret, SleepMs):
                             delay = arg
                         elif isinstance(ret, IORead):
-#                            self.add_reader(ret.obj.fileno(), lambda self, c, f: self.call_soon(c, f), self, cb, ret.obj)
-#                            self.add_reader(ret.obj.fileno(), lambda c, f: self.call_soon(c, f), cb, ret.obj)
-#                            self.add_reader(arg.fileno(), lambda cb: self.call_soon(cb), cb)
                             self.add_reader(arg.fileno(), cb)
                             continue
                         elif isinstance(ret, IOWrite):
-#                            self.add_writer(arg.fileno(), lambda cb: self.call_soon(cb), cb)
                             self.add_writer(arg.fileno(), cb)
                             continue
                         elif isinstance(ret, IOReadDone):
@@ -123,18 +151,56 @@ def run_forever(self):
                             self.remove_writer(arg.fileno(), cb)
                         elif isinstance(ret, StopLoop):
                             return arg
+                        else:
+                            assert False, "Unknown syscall yielded: %r (of type %r)" % (ret, type(ret))
                     elif isinstance(ret, type_gen):
                         self.call_soon(ret)
+                    elif isinstance(ret, int):
+                        # Delay
+                        delay = ret
                     elif ret is None:
                         # Just reschedule
                         pass
+                    elif ret is False:
+                        # Don't reschedule
+                        continue
                     else:
                         assert False, "Unsupported coroutine yield value: %r (of type %r)" % (ret, type(ret))
                 except StopIteration as e:
-                    if __debug__:
+                    if __debug__ and DEBUG:
                         log.debug("Coroutine finished: %s", cb)
+                    self.remove_polled_cb(cb)
+                    continue
+                except CancelledError as e:
+                    if __debug__ and DEBUG:
+                        log.debug("Coroutine cancelled: %s", cb)
+                    self.remove_polled_cb(cb)
                     continue
-                self.call_later(delay, cb, *args)
+                except Exception as e:
+                    if log:
+                        log.error("Coroutine %s exception: %s", cb, e)
+                # Currently all syscalls don't return anything, so we don't
+                # need to feed anything to the next invocation of coroutine.
+                # If that changes, need to pass that value below.
+                if delay:
+                    self.call_later_ms(delay, cb)
+                else:
+                    self.call_soon(cb)
+
+            # Wait until next waitq task or I/O availability
+            delay = 0
+            if not self.runq:
+                if self.waitq:
+                    tnow = self.time()
+                    t = self.waitq.peektime()
+                    delay = time.ticks_diff(t, tnow)
+                    if delay < 0:
+                        delay = 0
+                else:
+                    if __debug__ and DEBUG:
+                        log.info("No more tasks to execute, waiting forever")
+                    delay = -1
+            self.wait(delay)
 
     def run_until_complete(self, coro):
         def _run_and_stop():
@@ -143,6 +209,9 @@ def _run_and_stop():
         self.call_soon(_run_and_stop())
         self.run_forever()
 
+    def stop(self):
+        self.call_soon((lambda: (yield StopLoop(0)))())
+
     def close(self):
         pass
 
@@ -161,9 +230,6 @@ class SysCall1(SysCall):
     def __init__(self, arg):
         self.arg = arg
 
-class Sleep(SysCall1):
-    pass
-
 class StopLoop(SysCall1):
     pass
 
@@ -182,14 +248,77 @@ class IOWriteDone(SysCall1):
 
 _event_loop = None
 _event_loop_class = EventLoop
-def get_event_loop():
+def get_event_loop(runq_len=16, waitq_len=16):
     global _event_loop
     if _event_loop is None:
-        _event_loop = _event_loop_class()
+        _event_loop = _event_loop_class(runq_len, waitq_len)
     return _event_loop
 
 def sleep(secs):
-    yield Sleep(secs)
+    yield int(secs * 1000)
+
+# Implementation of sleep_ms awaitable with zero heap memory usage
+class SleepMs(SysCall1):
+
+    def __init__(self):
+        self.v = None
+        self.arg = None
+
+    def __call__(self, arg):
+        self.v = arg
+        #print("__call__")
+        return self
+
+    def __iter__(self):
+        #print("__iter__")
+        return self
+
+    def __next__(self):
+        if self.v is not None:
+            #print("__next__ syscall enter")
+            self.arg = self.v
+            self.v = None
+            return self
+        #print("__next__ syscall exit")
+        _stop_iter.__traceback__ = None
+        raise _stop_iter
+
+_stop_iter = StopIteration()
+sleep_ms = SleepMs()
+
+
+def cancel(coro, exc=CancelledError):
+    _event_loop.cancel(coro, exc=exc)
+
+
+class TimeoutObj:
+    def __init__(self, coro):
+        self.coro = coro
+
+
+def wait_for_ms(coro, timeout):
+
+    def waiter(coro, timeout_obj):
+        res = yield from coro
+        if __debug__ and DEBUG:
+            log.debug("waiter: cancelling %s", timeout_obj)
+        timeout_obj.coro = None
+        return res
+
+    def timeout_func(timeout_obj):
+        if timeout_obj.coro:
+            if __debug__ and DEBUG:
+                log.debug("timeout_func: cancelling %s", timeout_obj.coro)
+            cancel(timeout_obj.coro, exc=TimeoutError())
+
+    timeout_obj = TimeoutObj(_event_loop.cur_task)
+    _event_loop.call_later_ms(timeout, timeout_func, timeout_obj)
+    return (yield from waiter(coro, timeout_obj))
+
+
+def wait_for(coro, timeout):
+    return wait_for_ms(coro, int(timeout * 1000))
+
 
 def coroutine(f):
     return f