Fix some issues with the throw syscall, two-way proxy now working

aiosched/context.py

@@ -84,7 +84,7 @@ class TaskPool:
     """
     An asynchronous context manager that automatically waits
     for all tasks spawned within it and cancels itself when
-    an exception occurs. Contexts can be nested and will
+    an exception occurs. Pools can be nested and will
     cancel inner ones if an exception is raised inside them
     """
 

aiosched/kernel.py

@@ -126,10 +126,10 @@ class FIFOKernel:
         :return:
         """
 
-        if currently_protected():
-            self._sigint_handled = True
-        else:
-            raise KeyboardInterrupt
+        self._sigint_handled = True
+        if not currently_protected():
+            self.run_ready.appendleft(self.entry_point)
+            self.handle_errors(self.run_task_step)
 
     def done(self) -> bool:
         """
@@ -352,13 +352,15 @@ class FIFOKernel:
                 # there are no more runnable tasks
                 return
             self.current_task = self.run_ready.popleft()
-        self._running = True
-        _runner = self.current_task.run
-        _data = [self.data.pop(self.current_task, None)]
         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)]
         # Some debugging and internal chatter here
         self.current_task.steps += 1
         self.current_task.state = TaskState.RUN
@@ -368,6 +370,10 @@ class FIFOKernel:
             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)
@@ -587,10 +593,9 @@ class FIFOKernel:
         Throws the given exception into the given task
         """
 
-        self.paused.discard(task)
-        self.io_release_task(task)
-        self.handle_errors(partial(task.throw, error), task)
+        task.pending_exception = error
         self.run_ready.appendleft(task)
+        self.handle_errors(partial(task.throw, error), task)
         self.reschedule_running()
 
     def handle_errors(self, func: Callable, task: Task | None = None):
@@ -615,6 +620,7 @@ class FIFOKernel:
             # most of this code below is just useful for internal/debugging purposes
             task.state = TaskState.FINISHED
             task.result = ret.value
+            self.io_release_task(self.current_task)
             self.wait(task)
         except Cancelled:
             # When a task needs to be cancelled, aiosched tries to do it gracefully
@@ -667,8 +673,6 @@ class FIFOKernel:
         if task is not self.current_task:
             task.joiners.add(self.current_task)
         if task.done():
-            self.paused.discard(task)
-            self.io_release_task(task)
             self.run_ready.extend(task.joiners)
 
     def spawn(self, func: Callable[..., Coroutine[Any, Any, Any]], *args, **kwargs):

aiosched/task.py

@@ -85,6 +85,8 @@ class Task:
     propagate: bool = True
     # The task's scope
     scope: "TaskScope" = field(default=None, repr=False)
+    # Is there any pending exception?
+    pending_exception: BaseException | None = None
 
     def run(self, what: Any | None = None):
         """

tests/proxy.py

@@ -1,7 +1,6 @@
 import aiosched
 import socket
 
-# TODO: This is borked
 
 # Pls notice me njsmith senpai :>
 
@@ -22,6 +21,7 @@ async def proxy_two_way(a: aiosched.socket.AsyncSocket, b: aiosched.socket.Async
 
 
 async def main():
+    print("Starting two-way proxy server")
     async with aiosched.skip_after(10):
         a = aiosched.socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         b = aiosched.socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -29,6 +29,7 @@ async def main():
         await b.connect(("localhost", 54321))
         async with a, b:
             await proxy_two_way(a, b)
+    print("Two-way proxy shutting down")
 
 
 aiosched.run(main)