Fix more bugs in Ctrl+C delivery and restrict_ki_to_checkpoints

structio/core/kernel.py

@@ -12,7 +12,7 @@ from structio.abc import (
 from structio.io import FdWrapper
 from structio.core.context import TaskPool, TaskScope
 from structio.core.task import Task, TaskState
-from structio.util.ki import CTRLC_PROTECTION_ENABLED, critical_section
+from structio.util.ki import CTRLC_PROTECTION_ENABLED, critical_section, currently_protected
 from structio.exceptions import (
     StructIOException,
     Cancelled,
@@ -167,9 +167,12 @@ class DefaultKernel(BaseKernel):
     def signal_notify(self, sig: int, frame: FrameType):
         match sig:
             case signal.SIGINT:
-                self._sigint_handled = True
-                # Poke the event loop with a stick ;)
-                self.policy.schedule(self.entry_point)
+                if self.restrict_ki_to_checkpoints:
+                    self._sigint_handled = True
+                    # Poke the event loop with a stick ;)
+                    self.policy.schedule(self._pick_ki_task())
+                else:
+                    self.raise_ki()
             case _:
                 pass
 
@@ -196,17 +199,15 @@ class DefaultKernel(BaseKernel):
         if self.current_task.pending_cancellation:
             self.cancel_task(self.current_task)
             return
-        elif self._sigint_handled and not critical_section(
-            self.current_task.coroutine.cr_frame
-        ):
+        elif self._sigint_handled:
             self.raise_ki(self.current_task)
-            return
         self.event("before_task_step", self.current_task)
         self.current_task.state = TaskState.RUNNING
         self.current_task.paused_when = 0
         self.current_pool = self.current_task.pool
         self.current_scope = self.current_task.scope
         data = self.handle(runner, self.current_task)
+        self._sigint_handled = False
         if data is not None:
             method, args, kwargs = data
             self.suspend()
@@ -237,9 +238,7 @@ class DefaultKernel(BaseKernel):
         self.policy.schedule(task)
 
     def check_cancelled(self, schedule: bool = True):
-        if self._sigint_handled:
-            self.raise_ki()
-        elif self.current_task.pending_cancellation:
+        if self.current_task.pending_cancellation:
             self.current_task: Task
             self.cancel_task(self.current_task)
         elif schedule:
@@ -320,18 +319,13 @@ class DefaultKernel(BaseKernel):
         """
 
         self._sigint_handled = False
-        task = task or self._pick_ki_task()
-        self.throw(task, KeyboardInterrupt())
-        if task.done():
-            self.close()
+        self.throw(task or self._pick_ki_task(), KeyboardInterrupt())
 
     def _tick(self):
         """
         Runs a single event loop tick
         """
 
-        if self._sigint_handled and not self.restrict_ki_to_checkpoints:
-            self.raise_ki()
         self.wakeup()
         self.check_scopes()
         self.step()

structio/core/managers/signals/sigint.py

@@ -1,11 +1,12 @@
 from structio.abc import SignalManager
-from structio.util.ki import currently_protected
 from structio.signals import set_signal_handler
-from structio.core.run import current_loop, current_task
+from structio.core.run import current_loop
 from types import FrameType
 import warnings
 import signal
 
+from structio.util.ki import currently_protected
+
 
 class SigIntManager(SignalManager):
     """

tests/processes.py

@@ -55,7 +55,7 @@ async def main_limiter():
     for i in range(0, pool.max_workers * 4 + 1):   # We iterate one more time than the limit: this is on purpose
         # This will stop every 2 iterations
         await pool.submit(shlex.split(f"""python -c 'print({i}); __import__("time").sleep(1)'"""))
-    print(f"Submitted {i + 1} processes")
+    print(f"[main] Submitted {i + 1} processes")
 
 
 def foo():
@@ -74,8 +74,9 @@ async def main_python():
 
 
 if __name__ == "__main__":
-    #structio.run(main_simple, "owo")
+    structio.run(main_simple, "owo")
     structio.run(main_limiter)
-    #structio.run(main_python)
+    # TODO: Needs fixing
+    # structio.run(main_python)