gh-132975: Improve Remote PDB interrupt handling

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