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
@@ -2677,7 +2678,7 @@ def _send(self, **kwargs):
         try:
             self._sockfile.write(json_payload.encode() + b"\n")
             self._sockfile.flush()
-        except OSError:
+        except (OSError, ValueError):
             # This means that the client has abruptly disconnected, but we'll
             # handle that the next time we try to read from the client instead
             # of trying to handle it from everywhere _send() may be called.
@@ -2903,9 +2904,14 @@ 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.sigint_received = False
+        self.raise_on_sigint = False
+        self.server_socket = server_socket
         self.interrupt_script = interrupt_script
         self.pdb_instance = Pdb()
         self.pdb_commands = set()
@@ -2947,8 +2953,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,8 +2962,40 @@ def _send(self, **kwargs):
             # return an empty string because the socket may be half-closed.
             self.write_failed = True
 
+    def _readline(self):
+        if self.sigint_received:
+            # There's a pending unhandled SIGINT. Handle it now.
+            self.sigint_received = False
+            raise KeyboardInterrupt
+
+        # 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)
+
+        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 self.sigint_received:
+                        # If not, we're reading wakeup events for sigints that
+                        # we've previously handled, and can ignore them.
+                        self.sigint_received = False
+                        raise KeyboardInterrupt
+                elif key.fileobj == self.server_socket:
+                    data = self.server_socket.recv(16 * 1024)
+                    self.read_buf += data
+                    if not data and b"\n" not in self.read_buf:
+                        # EOF without a full final line. Drop the partial line.
+                        self.read_buf = b""
+                        return b""
+
+        ret, sep, self.read_buf = self.read_buf.partition(b"\n")
+        return ret + sep
+
     def read_command(self, prompt):
-        reply = input(prompt)
+        with self._sigint_raises_keyboard_interrupt():
+            reply = input(prompt)
 
         if self.state == "dumb":
             # No logic applied whatsoever, just pass the raw reply back.
@@ -2984,7 +3021,8 @@ def read_command(self, prompt):
         # Otherwise, valid first line of a multi-line statement
         continue_prompt = "...".ljust(len(prompt))
         while codeop.compile_command(reply, "<stdin>", "single") is None:
-            reply += "\n" + input(continue_prompt)
+            with self._sigint_raises_keyboard_interrupt():
+                reply += "\n" + input(continue_prompt)
 
         return prefix + reply
 
@@ -3009,11 +3047,70 @@ def readline_completion(self, completer):
         finally:
             readline.set_completer(old_completer)
 
+    @contextmanager
+    def _sigint_handler(self):
+        # Signal handling strategy:
+        # - When we call input() we want a SIGINT to raise KeyboardInterrupt
+        # - Otherwise we want to write to the wakeup FD and set a flag.
+        #   We'll break out of select() when the wakeup FD is written to,
+        #   and we'll check the flag whenever we're about to accept input.
+        def handler(signum, frame):
+            self.sigint_received = True
+            if self.raise_on_sigint:
+                # One-shot; don't raise again until the flag is set again.
+                self.raise_on_sigint = False
+                self.sigint_received = False
+                raise KeyboardInterrupt
+
+        sentinel = object()
+        old_handler = sentinel
+        old_wakeup_fd = sentinel
+
+        self.signal_read, self.signal_write = socket.socketpair()
+        with (closing(self.signal_read), closing(self.signal_write)):
+            self.signal_read.setblocking(False)
+            self.signal_write.setblocking(False)
+
+            try:
+                old_handler = signal.signal(signal.SIGINT, handler)
+
+                try:
+                    old_wakeup_fd = signal.set_wakeup_fd(
+                        self.signal_write.fileno(),
+                        warn_on_full_buffer=False,
+                    )
+                    yield
+                finally:
+                    # Restore the old wakeup fd if we installed a new one
+                    if old_wakeup_fd is not sentinel:
+                        signal.set_wakeup_fd(old_wakeup_fd)
+                    self.signal_read = self.signal_write = None
+            finally:
+                if old_handler is not sentinel:
+                    # Restore the old handler if we installed a new one
+                    signal.signal(signal.SIGINT, old_handler)
+
+    @contextmanager
+    def _sigint_raises_keyboard_interrupt(self):
+        if self.sigint_received:
+            # There's a pending unhandled SIGINT. Handle it now.
+            self.sigint_received = False
+            raise KeyboardInterrupt
+
+        try:
+            self.raise_on_sigint = True
+            yield
+        finally:
+            self.raise_on_sigint = False
+
     def cmdloop(self):
-        with self.readline_completion(self.complete):
+        with (
+            self._sigint_handler(),
+            self.readline_completion(self.complete),
+        ):
             while not self.write_failed:
                 try:
-                    if not (payload_bytes := self.sockfile.readline()):
+                    if not (payload_bytes := self._readline()):
                         break
                 except KeyboardInterrupt:
                     self.send_interrupt()
@@ -3031,11 +3128,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:
@@ -3101,7 +3207,7 @@ def complete(self, text, state):
             if self.write_failed:
                 return None
 
-            payload = self.sockfile.readline()
+            payload = self._readline()
             if not payload:
                 return None
 
@@ -3168,9 +3274,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 +3282,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