gh-132975: Improve Remote PDB interrupt handling

Lib/pdb.py

@@ -92,6 +92,7 @@
 import itertools
 import traceback
 import linecache
+import selectors
 import _colorize
 
 from contextlib import closing
@@ -2903,9 +2904,12 @@ def default(self, line):
 
 
 class _PdbClient:
-    def __init__(self, pid, sockfile, interrupt_script):
+    def __init__(self, pid, server_socket, interrupt_script):
         self.pid = pid
-        self.sockfile = sockfile
+        self.read_buf = b""
+        self.signal_read = None
+        self.signal_write = None
+        self.server_socket = server_socket
         self.interrupt_script = interrupt_script
         self.pdb_instance = Pdb()
         self.pdb_commands = set()
@@ -2947,8 +2951,7 @@ def _send(self, **kwargs):
         self._ensure_valid_message(kwargs)
         json_payload = json.dumps(kwargs)
         try:
-            self.sockfile.write(json_payload.encode() + b"\n")
-            self.sockfile.flush()
+            self.server_socket.sendall(json_payload.encode() + b"\n")
         except OSError:
             # This means that the client has abruptly disconnected, but we'll
             # handle that the next time we try to read from the client instead
@@ -2957,6 +2960,41 @@ def _send(self, **kwargs):
             # return an empty string because the socket may be half-closed.
             self.write_failed = True
 
+    def _readline(self):
+        # Wait for either a SIGINT or a line or EOF from the PDB server.
+        selector = selectors.DefaultSelector()
+        selector.register(self.signal_read, selectors.EVENT_READ)
+        selector.register(self.server_socket, selectors.EVENT_READ)
+
+        old_wakeup_fd = signal.set_wakeup_fd(
+            self.signal_write.fileno(),
+            warn_on_full_buffer=False,
+        )
+
+        got_sigint = False
+        def sigint_handler(*args, **kwargs):
+            nonlocal got_sigint
+            got_sigint = True
+
+        old_handler = signal.signal(signal.SIGINT, sigint_handler)
+        try:
+            while b"\n" not in self.read_buf:
+                for key, _ in selector.select():
+                    if key.fileobj == self.signal_read:
+                        self.signal_read.recv(1024)
+                        if got_sigint:
+                            raise KeyboardInterrupt
+                    elif key.fileobj == self.server_socket:
+                        self.read_buf += self.server_socket.recv(16 * 1024)
+                        if not self.read_buf:
+                            return b""
+
+            ret, sep, self.read_buf = self.read_buf.partition(b"\n")
+            return ret + sep
+        finally:
+            signal.set_wakeup_fd(old_wakeup_fd)
+            signal.signal(signal.SIGINT, old_handler)
+
     def read_command(self, prompt):
         reply = input(prompt)
 
@@ -3009,12 +3047,65 @@ def readline_completion(self, completer):
         finally:
             readline.set_completer(old_completer)
 
+    if not hasattr(signal, "pthread_sigmask"):
+        # On Windows, we must drop signals arriving while we're ignoring them.
+        @contextmanager
+        def _block_sigint(self):
+            old_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
+            try:
+                yield
+            finally:
+                signal.signal(signal.SIGINT, old_handler)
+
+        @contextmanager
+        def _handle_sigint(self, handler):
+            old_handler = signal.signal(signal.SIGINT, handler)
+            try:
+                yield
+            finally:
+                signal.signal(signal.SIGINT, old_handler)
+    else:
+        # On Unix, we can save them to be processed later.
+        @contextmanager
+        def _block_sigint(self):
+            signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT})
+            try:
+                yield
+            finally:
+                signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGINT})
+
+        @contextmanager
+        def _handle_sigint(self, handler):
+            old_handler = signal.signal(signal.SIGINT, handler)
+            try:
+                signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGINT})
+                yield
+            finally:
+                signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT})
+                signal.signal(signal.SIGINT, old_handler)
+
+    @contextmanager
+    def _signal_socket_pair(self):
+        self.signal_read, self.signal_write = socket.socketpair()
+        try:
+            with (closing(self.signal_read), closing(self.signal_write)):
+                self.signal_read.setblocking(False)
+                self.signal_write.setblocking(False)
+                yield
+        finally:
+            self.signal_read = self.signal_write = None
+
     def cmdloop(self):
-        with self.readline_completion(self.complete):
+        with (
+            self._block_sigint(),
+            self._signal_socket_pair(),
+            self.readline_completion(self.complete),
+        ):
             while not self.write_failed:
                 try:
-                    if not (payload_bytes := self.sockfile.readline()):
-                        break
+                    with self._handle_sigint(signal.default_int_handler):
+                        if not (payload_bytes := self._readline()):
+                            break
                 except KeyboardInterrupt:
                     self.send_interrupt()
                     continue
@@ -3031,11 +3122,20 @@ def cmdloop(self):
                 self.process_payload(payload)
 
     def send_interrupt(self):
-        print(
-            "\n*** Program will stop at the next bytecode instruction."
-            " (Use 'cont' to resume)."
-        )
-        sys.remote_exec(self.pid, self.interrupt_script)
+        if hasattr(signal, "pthread_kill"):
+            # On Unix, send a SIGINT to the remote process, which interrupts IO
+            # and makes it raise a KeyboardInterrupt on the main thread when
+            # PyErr_CheckSignals is called or the eval loop regains control.
+            os.kill(self.pid, signal.SIGINT)
+        else:
+            # On Windows, inject a remote script that calls Pdb.set_trace()
+            # when the eval loop regains control. This cannot interrupt IO, and
+            # also cannot interrupt statements executed at a PDB prompt.
+            print(
+                "\n*** Program will stop at the next bytecode instruction."
+                " (Use 'cont' to resume)."
+            )
+            sys.remote_exec(self.pid, self.interrupt_script)
 
     def process_payload(self, payload):
         match payload:
@@ -3061,7 +3161,8 @@ def process_payload(self, payload):
     def prompt_for_reply(self, prompt):
         while True:
             try:
-                payload = {"reply": self.read_command(prompt)}
+                with self._handle_sigint(signal.default_int_handler):
+                    payload = {"reply": self.read_command(prompt)}
             except EOFError:
                 payload = {"signal": "EOF"}
             except KeyboardInterrupt:
@@ -3101,7 +3202,8 @@ def complete(self, text, state):
             if self.write_failed:
                 return None
 
-            payload = self.sockfile.readline()
+            with self._handle_sigint(signal.default_int_handler):
+                payload = self._readline()
             if not payload:
                 return None
 
@@ -3168,9 +3270,6 @@ def attach(pid, commands=()):
             client_sock, _ = server.accept()
 
             with closing(client_sock):
-                sockfile = client_sock.makefile("rwb")
-
-            with closing(sockfile):
                 with tempfile.NamedTemporaryFile("w", delete_on_close=False) as interrupt_script:
                     interrupt_script.write(
                         'import pdb, sys\n'
@@ -3179,7 +3278,7 @@ def attach(pid, commands=()):
                     )
                     interrupt_script.close()
 
-                    _PdbClient(pid, sockfile, interrupt_script.name).cmdloop()
+                    _PdbClient(pid, client_sock, interrupt_script.name).cmdloop()
 
 
 # Post-Mortem interface