[3.11] gh-96522: Fix deadlock in pty.spawn (GH-96639) (#104655)

(cherry picked from commit 9c5aa8967b)

Co-authored-by: Youfu Zhang <1315097+zhangyoufu@users.noreply.github.com>

Lib/pty.py

@@ -121,12 +121,6 @@ def fork():
     # Parent and child process.
     return pid, master_fd
 
-def _writen(fd, data):
-    """Write all the data to a descriptor."""
-    while data:
-        n = os.write(fd, data)
-        data = data[n:]
-
 def _read(fd):
     """Default read function."""
     return os.read(fd, 1024)
@@ -136,9 +130,42 @@ def _copy(master_fd, master_read=_read, stdin_read=_read):
     Copies
             pty master -> standard output   (master_read)
             standard input -> pty master    (stdin_read)"""
-    fds = [master_fd, STDIN_FILENO]
-    while fds:
-        rfds, _wfds, _xfds = select(fds, [], [])
+    if os.get_blocking(master_fd):
+        # If we write more than tty/ndisc is willing to buffer, we may block
+        # indefinitely. So we set master_fd to non-blocking temporarily during
+        # the copy operation.
+        os.set_blocking(master_fd, False)
+        try:
+            _copy(master_fd, master_read=master_read, stdin_read=stdin_read)
+        finally:
+            # restore blocking mode for backwards compatibility
+            os.set_blocking(master_fd, True)
+        return
+    high_waterlevel = 4096
+    stdin_avail = master_fd != STDIN_FILENO
+    stdout_avail = master_fd != STDOUT_FILENO
+    i_buf = b''
+    o_buf = b''
+    while 1:
+        rfds = []
+        wfds = []
+        if stdin_avail and len(i_buf) < high_waterlevel:
+            rfds.append(STDIN_FILENO)
+        if stdout_avail and len(o_buf) < high_waterlevel:
+            rfds.append(master_fd)
+        if stdout_avail and len(o_buf) > 0:
+            wfds.append(STDOUT_FILENO)
+        if len(i_buf) > 0:
+            wfds.append(master_fd)
+
+        rfds, wfds, _xfds = select(rfds, wfds, [])
+
+        if STDOUT_FILENO in wfds:
+            try:
+                n = os.write(STDOUT_FILENO, o_buf)
+                o_buf = o_buf[n:]
+            except OSError:
+                stdout_avail = False
 
         if master_fd in rfds:
             # Some OSes signal EOF by returning an empty byte string,
@@ -150,15 +177,18 @@ def _copy(master_fd, master_read=_read, stdin_read=_read):
             if not data:  # Reached EOF.
                 return    # Assume the child process has exited and is
                           # unreachable, so we clean up.
-            else:
-                os.write(STDOUT_FILENO, data)
+            o_buf += data
 
-        if STDIN_FILENO in rfds:
+        if master_fd in wfds:
+            n = os.write(master_fd, i_buf)
+            i_buf = i_buf[n:]
+
+        if stdin_avail and STDIN_FILENO in rfds:
             data = stdin_read(STDIN_FILENO)
             if not data:
-                fds.remove(STDIN_FILENO)
+                stdin_avail = False
             else:
-                _writen(master_fd, data)
+                i_buf += data
 
 def spawn(argv, master_read=_read, stdin_read=_read):
     """Create a spawned process."""