diff --git a/aiosched/context.py b/aiosched/context.py index d0b9b97..a13fd89 100644 --- a/aiosched/context.py +++ b/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 """ diff --git a/aiosched/kernel.py b/aiosched/kernel.py index 9a00b5d..5807362 100644 --- a/aiosched/kernel.py +++ b/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): diff --git a/aiosched/task.py b/aiosched/task.py index 01d842e..b669686 100644 --- a/aiosched/task.py +++ b/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): """ diff --git a/tests/proxy.py b/tests/proxy.py index 69edc1c..f01cddf 100644 --- a/tests/proxy.py +++ b/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) diff --git a/tests/timeout.py b/tests/timeout.py deleted file mode 100644 index e69de29..0000000