From 724b5c98bafd634b46b44b28724cf8ce20b8d041 Mon Sep 17 00:00:00 2001
From: Nocturn9x <hackhab@gmail.com>
Date: Thu, 11 May 2023 14:17:09 +0200
Subject: [PATCH] Fixed bugs with cancellations and exceptions

---
 aiosched/internals/syscalls.py |  2 ++
 aiosched/kernel.py             | 44 +++++++++++++++-------------------
 tests/deadlock.py              | 15 ++++++++++++
 tests/proxy.py                 |  6 ++++-
 4 files changed, 41 insertions(+), 26 deletions(-)
 create mode 100644 tests/deadlock.py

diff --git a/aiosched/internals/syscalls.py b/aiosched/internals/syscalls.py
index e695bde..188f2a6 100644
--- a/aiosched/internals/syscalls.py
+++ b/aiosched/internals/syscalls.py
@@ -166,6 +166,8 @@ async def cancel(task: Task, block: bool = False):
 
     if task.done():
         return
+    if task.scope and not task.scope.cancellable:
+        return
     await syscall("cancel", task)
     if block:
         await wait(task)
diff --git a/aiosched/kernel.py b/aiosched/kernel.py
index 5f23294..03e0478 100644
--- a/aiosched/kernel.py
+++ b/aiosched/kernel.py
@@ -126,10 +126,12 @@ class FIFOKernel:
         :return:
         """
 
-        self._sigint_handled = True
         if not currently_protected():
             self.run_ready.appendleft(self.entry_point)
+            self.entry_point.pending_exception = KeyboardInterrupt()
             self.handle_errors(self.run_task_step)
+        else:
+            self._sigint_handled = True
 
     def done(self) -> bool:
         """
@@ -352,28 +354,17 @@ class FIFOKernel:
                 # there are no more runnable tasks
                 return
             self.current_task = self.run_ready.popleft()
-        if self.current_task.pending_cancellation:
-            # We perform the deferred cancellation
-            # if it was previously scheduled
-            self.cancel(self.current_task)
-        if self.current_task.done():
-            return
         self._running = True
         _runner = self.current_task.run
         _data = [self.data.pop(self.current_task, None)]
+        if exc := self.current_task.pending_exception:
+            self.current_task.pending_exception = None
+            _runner = partial(self.current_task.throw, exc)
+            _data = []
         # Some debugging and internal chatter here
         self.current_task.steps += 1
         self.current_task.state = TaskState.RUN
         self.debugger.before_task_step(self.current_task)
-        if self._sigint_handled:
-            self._sigint_handled = False
-            self.current_task.throw(KeyboardInterrupt())
-            _runner = partial(self.current_task.throw, KeyboardInterrupt)
-            _data = []
-        elif exc := self.current_task.pending_exception:
-            self.current_task.pending_exception = None
-            _runner = partial(self.current_task.throw, exc)
-            _data = []
         # Run a single step with the calculation (i.e. until a yield
         # somewhere)
         method, args, kwargs = _runner(*_data)
@@ -456,7 +447,9 @@ class FIFOKernel:
                 else:
                     task = self.current_task
                 """
+                self._sigint_handled = False
                 task = self.entry_point
+                task.pending_exception = KeyboardInterrupt()
                 self.run_ready.appendleft(task)
                 self.handle_errors(self.run_task_step)
             elif not self.run_ready:
@@ -579,14 +572,14 @@ class FIFOKernel:
         it fails
         """
 
-        if not task.scope or task.scope.cancellable:
-            self.paused.discard(task)
-            self.io_release_task(task)
-            self.run_ready.append(task)
-            self.handle_errors(partial(task.throw, Cancelled(task)), task)
-            if task.state != TaskState.CANCELLED:
-                task.pending_cancellation = True
         self.reschedule_running()
+        self.paused.discard(task)
+        self.io_release_task(task)
+        self.run_ready.appendleft(task)
+        self.handle_errors(partial(task.throw, Cancelled()), task)
+        self.handle_errors(self.run_task_step)
+        if task.state != TaskState.CANCELLED:
+            task.pending_exception = Cancelled()
 
     def throw(self, task, error):
         """
@@ -682,7 +675,7 @@ class FIFOKernel:
         """
 
         task = Task(func.__name__ or repr(func), func(*args, **kwargs))
-        # We inject our magic secret variable into the coroutine's stack frame so
+        # We inject our magic secret variable into the coroutine's stack frame, so
         # we can look it up later
         task.coroutine.cr_frame.f_locals.setdefault(CTRLC_PROTECTION_ENABLED, False)
         task.scope = self.current_scope
@@ -736,7 +729,8 @@ class FIFOKernel:
                 return
             elif self.current_task.last_io[1] == resource:
                 # If the event to listen for has changed we just modify it
-                self.selector.modify(resource, evt_type, {evt_type: self.current_task})
+                key = self.selector.get_key(resource)
+                self.selector.modify(resource, evt_type | key.events, key.data.update({evt_type: self.current_task}))
                 self.current_task.last_io = (evt_type, resource)
                 self.debugger.on_io_schedule(resource, evt_type)
         elif not self.current_task.last_io or self.current_task.last_io[1] != resource:
diff --git a/tests/deadlock.py b/tests/deadlock.py
new file mode 100644
index 0000000..dfa9f81
--- /dev/null
+++ b/tests/deadlock.py
@@ -0,0 +1,15 @@
+import aiosched
+
+
+async def deadlock():
+    await aiosched.Event().wait()
+
+
+async def parent():
+    async with aiosched.skip_after(5):
+        await deadlock()
+    print("Done")
+
+
+aiosched.run(parent)
+print("Exited")
diff --git a/tests/proxy.py b/tests/proxy.py
index f01cddf..cdb8bcd 100644
--- a/tests/proxy.py
+++ b/tests/proxy.py
@@ -32,4 +32,8 @@ async def main():
     print("Two-way proxy shutting down")
 
 
-aiosched.run(main)
+try:
+    aiosched.run(main)
+except KeyboardInterrupt:
+    print(f"Ctrl+C caught")
+