gh-132975: Improve Remote PDB interrupt handling (#133223)

Author: godlygeek

Lib/pdb.py

@@ -77,6 +77,7 @@
 import json
 import token
 import types
+import atexit
 import codeop
 import pprint
 import signal
@@ -92,11 +93,12 @@
 import itertools
 import traceback
 import linecache
+import selectors
+import threading
 import _colorize
 import _pyrepl.utils
 
-from contextlib import closing
-from contextlib import contextmanager
+from contextlib import ExitStack, closing, contextmanager
 from rlcompleter import Completer
 from types import CodeType
 from warnings import deprecated
@@ -2670,12 +2672,21 @@ async def set_trace_async(*, header=None, commands=None):
 # Remote PDB
 
 class _PdbServer(Pdb):
-    def __init__(self, sockfile, owns_sockfile=True, **kwargs):
+    def __init__(
+        self,
+        sockfile,
+        signal_server=None,
+        owns_sockfile=True,
+        **kwargs,
+    ):
         self._owns_sockfile = owns_sockfile
         self._interact_state = None
         self._sockfile = sockfile
         self._command_name_cache = []
         self._write_failed = False
+        if signal_server:
+            # Only started by the top level _PdbServer, not recursive ones.
+            self._start_signal_listener(signal_server)
         super().__init__(colorize=False, **kwargs)
 
     @staticmethod
@@ -2731,15 +2742,49 @@ def _ensure_valid_message(self, msg):
                     f"PDB message doesn't follow the schema! {msg}"
                 )
 
+    @classmethod
+    def _start_signal_listener(cls, address):
+        def listener(sock):
+            with closing(sock):
+                # Check if the interpreter is finalizing every quarter of a second.
+                # Clean up and exit if so.
+                sock.settimeout(0.25)
+                sock.shutdown(socket.SHUT_WR)
+                while not shut_down.is_set():
+                    try:
+                        data = sock.recv(1024)
+                    except socket.timeout:
+                        continue
+                    if data == b"":
+                        return  # EOF
+                    signal.raise_signal(signal.SIGINT)
+
+        def stop_thread():
+            shut_down.set()
+            thread.join()
+
+        # Use a daemon thread so that we don't detach until after all non-daemon
+        # threads are done. Use an atexit handler to stop gracefully at that point,
+        # so that our thread is stopped before the interpreter is torn down.
+        shut_down = threading.Event()
+        thread = threading.Thread(
+            target=listener,
+            args=[socket.create_connection(address, timeout=5)],
+            daemon=True,
+        )
+        atexit.register(stop_thread)
+        thread.start()
+
     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()
-        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
+        except (OSError, ValueError):
+            # We get an OSError if the network connection has dropped, and a
+            # ValueError if detach() if the sockfile has been closed. We'll
+            # handle this the next time we try to read from the client instead
             # of trying to handle it from everywhere _send() may be called.
             # Track this with a flag rather than assuming readline() will ever
             # return an empty string because the socket may be half-closed.
@@ -2967,10 +3012,15 @@ def default(self, line):
 
 
 class _PdbClient:
-    def __init__(self, pid, sockfile, interrupt_script):
+    def __init__(self, pid, server_socket, interrupt_sock):
         self.pid = pid
-        self.sockfile = sockfile
-        self.interrupt_script = interrupt_script
+        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_sock = interrupt_sock
         self.pdb_instance = Pdb()
         self.pdb_commands = set()
         self.completion_matches = []
@@ -3012,8 +3062,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
@@ -3022,10 +3071,44 @@ def _send(self, **kwargs):
             # return an empty string because the socket may be half-closed.
             self.write_failed = True
 
-    def read_command(self, prompt):
-        self.multiline_block = False
-        reply = input(prompt)
+    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_input(self, prompt, multiline_block):
+        self.multiline_block = multiline_block
+        with self._sigint_raises_keyboard_interrupt():
+            return input(prompt)
 
+    def read_command(self, prompt):
+        reply = self.read_input(prompt, multiline_block=False)
         if self.state == "dumb":
             # No logic applied whatsoever, just pass the raw reply back.
             return reply
@@ -3048,10 +3131,9 @@ def read_command(self, prompt):
             return prefix + reply
 
         # Otherwise, valid first line of a multi-line statement
-        self.multiline_block = True
-        continue_prompt = "...".ljust(len(prompt))
+        more_prompt = "...".ljust(len(prompt))
         while codeop.compile_command(reply, "<stdin>", "single") is None:
-            reply += "\n" + input(continue_prompt)
+            reply += "\n" + self.read_input(more_prompt, multiline_block=True)
 
         return prefix + reply
 
@@ -3076,11 +3158,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)
+            finally:
+                self.signal_read = self.signal_write = None
+                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()
@@ -3098,11 +3239,17 @@ 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 self.interrupt_sock is not None:
+            # Write to a socket that the PDB server listens on. This triggers
+            # the remote to raise a SIGINT for itself. We do this because
+            # Windows doesn't allow triggering SIGINT remotely.
+            # See https://stackoverflow.com/a/35792192 for many more details.
+            self.interrupt_sock.sendall(signal.SIGINT.to_bytes())
+        else:
+            # On Unix we can just send a SIGINT to the remote process.
+            # This is preferable to using the signal thread approach that we
+            # use on Windows because it can interrupt IO in the main thread.
+            os.kill(self.pid, signal.SIGINT)
 
     def process_payload(self, payload):
         match payload:
@@ -3172,7 +3319,7 @@ def complete(self, text, state):
             if self.write_failed:
                 return None
 
-            payload = self.sockfile.readline()
+            payload = self._readline()
             if not payload:
                 return None
 
@@ -3189,11 +3336,18 @@ def complete(self, text, state):
             return None
 
 
-def _connect(host, port, frame, commands, version):
+def _connect(*, host, port, frame, commands, version, signal_raising_thread):
     with closing(socket.create_connection((host, port))) as conn:
         sockfile = conn.makefile("rwb")
 
-    remote_pdb = _PdbServer(sockfile)
+    # The client requests this thread on Windows but not on Unix.
+    # Most tests don't request this thread, to keep them simpler.
+    if signal_raising_thread:
+        signal_server = (host, port)
+    else:
+        signal_server = None
+
+    remote_pdb = _PdbServer(sockfile, signal_server=signal_server)
     weakref.finalize(remote_pdb, sockfile.close)
 
     if Pdb._last_pdb_instance is not None:
@@ -3214,43 +3368,48 @@ def _connect(host, port, frame, commands, version):
 
 def attach(pid, commands=()):
     """Attach to a running process with the given PID."""
-    with closing(socket.create_server(("localhost", 0))) as server:
+    with ExitStack() as stack:
+        server = stack.enter_context(
+            closing(socket.create_server(("localhost", 0)))
+        )
         port = server.getsockname()[1]
 
-        with tempfile.NamedTemporaryFile("w", delete_on_close=False) as connect_script:
-            connect_script.write(
-                textwrap.dedent(
-                    f"""
-                    import pdb, sys
-                    pdb._connect(
-                        host="localhost",
-                        port={port},
-                        frame=sys._getframe(1),
-                        commands={json.dumps("\n".join(commands))},
-                        version={_PdbServer.protocol_version()},
-                    )
-                    """
+        connect_script = stack.enter_context(
+            tempfile.NamedTemporaryFile("w", delete_on_close=False)
+        )
+
+        use_signal_thread = sys.platform == "win32"
+
+        connect_script.write(
+            textwrap.dedent(
+                f"""
+                import pdb, sys
+                pdb._connect(
+                    host="localhost",
+                    port={port},
+                    frame=sys._getframe(1),
+                    commands={json.dumps("\n".join(commands))},
+                    version={_PdbServer.protocol_version()},
+                    signal_raising_thread={use_signal_thread!r},
                 )
+                """
             )
-            connect_script.close()
-            sys.remote_exec(pid, connect_script.name)
-
-            # TODO Add a timeout? Or don't bother since the user can ^C?
-            client_sock, _ = server.accept()
+        )
+        connect_script.close()
+        sys.remote_exec(pid, connect_script.name)
 
-            with closing(client_sock):
-                sockfile = client_sock.makefile("rwb")
+        # TODO Add a timeout? Or don't bother since the user can ^C?
+        client_sock, _ = server.accept()
+        stack.enter_context(closing(client_sock))
 
-            with closing(sockfile):
-                with tempfile.NamedTemporaryFile("w", delete_on_close=False) as interrupt_script:
-                    interrupt_script.write(
-                        'import pdb, sys\n'
-                        'if inst := pdb.Pdb._last_pdb_instance:\n'
-                        '    inst.set_trace(sys._getframe(1))\n'
-                    )
-                    interrupt_script.close()
+        if use_signal_thread:
+            interrupt_sock, _ = server.accept()
+            stack.enter_context(closing(interrupt_sock))
+            interrupt_sock.setblocking(False)
+        else:
+            interrupt_sock = None
 
-                    _PdbClient(pid, sockfile, interrupt_script.name).cmdloop()
+        _PdbClient(pid, client_sock, interrupt_sock).cmdloop()
 
 
 # Post-Mortem interface