From f6ac0c2eb79157f7dab2b852ea4eb9fb73baf384 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Thu, 13 Oct 2022 13:12:17 -0400
Subject: [PATCH] Always restore at least `trio`'s sigint handler

We can get it during runtime startup and stash on a new
`Lock._trio_handler`. Always at least revert to this handler to
guarantee graceful kbi handling despite mucking about with our own
handler in debug mode.
---
 tractor/_debug.py | 42 +++++++++++++++++++++++++++++-------------
 tractor/_root.py  | 10 ++++++++--
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/tractor/_debug.py b/tractor/_debug.py
index 751c646..819dd9a 100644
--- a/tractor/_debug.py
+++ b/tractor/_debug.py
@@ -75,8 +75,10 @@ class Lock:
     # placeholder for function to set a ``trio.Event`` on debugger exit
     # pdb_release_hook: Optional[Callable] = None
 
+    _trio_handler: Callable | None = None
+
     # actor-wide variable pointing to current task name using debugger
-    local_task_in_debug: Optional[str] = None
+    local_task_in_debug: str | None = None
 
     # NOTE: set by the current task waiting on the root tty lock from
     # the CALLER side of the `lock_tty_for_child()` context entry-call
@@ -111,13 +113,10 @@ class Lock:
 
     @classmethod
     def unshield_sigint(cls):
-        if cls._orig_sigint_handler is not None:
-            # restore original sigint handler
-            signal.signal(
-                signal.SIGINT,
-                cls._orig_sigint_handler
-            )
-
+        # always restore (some) sigint handler, either
+        # the prior or at least ``trio``'s.
+        orig = cls._orig_sigint_handler or cls._trio_handler
+        signal.signal(signal.SIGINT, orig)
         cls._orig_sigint_handler = None
 
     @classmethod
@@ -544,7 +543,7 @@ def shield_sigint(
 
 ) -> None:
     '''
-    Specialized debugger compatible SIGINT handler.
+    Specialized, debugger-aware SIGINT handler.
 
     In childred we always ignore to avoid deadlocks since cancellation
     should always be managed by the parent supervising actor. The root
@@ -601,6 +600,8 @@ def shield_sigint(
         # which has already terminated to unlock.
         and any_connected
     ):
+        # we are root and some actor is in debug mode
+        # if uid_in_debug is not None:
         name = uid_in_debug[0]
         if name != 'root':
             log.pdb(
@@ -611,6 +612,22 @@ def shield_sigint(
             log.pdb(
                 "Ignoring SIGINT while in debug mode"
             )
+    elif (
+        is_root_process()
+    ):
+        log.pdb(
+            "Ignoring SIGINT since debug mode is enabled"
+        )
+
+        # revert back to ``trio`` handler asap!
+        Lock.unshield_sigint()
+        if (
+            Lock._root_local_task_cs_in_debug
+            and not Lock._root_local_task_cs_in_debug.cancel_called
+        ):
+            Lock._root_local_task_cs_in_debug.cancel()
+
+        # raise KeyboardInterrupt
 
     # child actor that has locked the debugger
     elif not is_root_process():
@@ -636,10 +653,9 @@ def shield_sigint(
         # https://github.com/goodboy/tractor/issues/320
         # elif debug_mode():
 
-    else:
-        log.pdb(
-            "Ignoring SIGINT since debug mode is enabled"
-        )
+    else:  # XXX: shouldn't ever get here?
+        print("WTFWTFWTF")
+        raise KeyboardInterrupt
 
     # NOTE: currently (at least on ``fancycompleter`` 0.9.2)
     # it lookks to be that the last command that was run (eg. ll)
diff --git a/tractor/_root.py b/tractor/_root.py
index 16c4bb8..e43121c 100644
--- a/tractor/_root.py
+++ b/tractor/_root.py
@@ -23,6 +23,7 @@ from functools import partial
 import importlib
 import logging
 import os
+import signal
 from typing import (
     Optional,
 )
@@ -76,14 +77,19 @@ async def open_root_actor(
     rpc_module_paths: Optional[list] = None,
 
 ) -> typing.Any:
-    """Async entry point for ``tractor``.
+    '''
+    Runtime init entry point for ``tractor``.
 
-    """
+    '''
     # Override the global debugger hook to make it play nice with
     # ``trio``, see:
     # https://github.com/python-trio/trio/issues/1155#issuecomment-742964018
     os.environ['PYTHONBREAKPOINT'] = 'tractor._debug._set_trace'
 
+    # attempt to retreive ``trio``'s sigint handler and stash it
+    # on our debugger lock state.
+    _debug.Lock._trio_handler = signal.getsignal(signal.SIGINT)
+
     # mark top most level process as root actor
     _state._runtime_vars['_is_root'] = True