Don't hold socket locks across blocking I/O (#8746)

accept/recv/send kept the socket RwLock during the syscall. After
fork, the child copied that held lock, so _ResourceSharer._afterfork
hung in listener.close().

Unskip fork connection pickling, pool task ref release, and spawn
thread-pool terminate. Add extra_tests for fork+accept close and
fork after DupFd.

Assisted-by: Grok:4.6

Lib/test/_test_multiprocessing.py

@@ -3166,7 +3166,6 @@ class _TestPool(BaseTestCase):
         # check that we indeed waited for all jobs
         self.assertGreater(time.monotonic() - t_start, 0.9)
 
-    @unittest.skip("TODO: RUSTPYTHON; reference counting differences")
     def test_release_task_refs(self):
         # Issue #29861: task arguments and results should not be kept
         # alive after we are done with them.

Lib/test/test_multiprocessing_fork/test_processes.py

@@ -3,9 +3,5 @@ from test._test_multiprocessing import install_tests_in_module_dict
 
 install_tests_in_module_dict(globals(), 'fork', only_type="processes")
 
-class WithProcessesTestPicklingConnections(WithProcessesTestPicklingConnections):  # TODO: RUSTPYTHON
-    @unittest.skip('TODO: RUSTPYTHON; hangs')
-    def test_pickling(self): super().test_pickling()
-
 if __name__ == '__main__':
     unittest.main()

Lib/test/test_multiprocessing_spawn/test_threads.py

@@ -3,10 +3,5 @@ from test._test_multiprocessing import install_tests_in_module_dict
 
 install_tests_in_module_dict(globals(), 'spawn', only_type="threads")
 
-import os, sys  # TODO: RUSTPYTHON
-class WithThreadsTestPool(WithThreadsTestPool):  # TODO: RUSTPYTHON
-    @unittest.skip("TODO: RUSTPYTHON; flaky environment pollution when running rustpython -m test --fail-env-changed due to unknown reason")
-    def test_terminate(self): super().test_terminate()  # TODO: RUSTPYTHON
-
 if __name__ == '__main__':
     unittest.main()

crates/stdlib/src/socket.rs

@@ -45,6 +45,7 @@ mod _socket {
     }
 
     use core::{
+        mem::ManuallyDrop,
         net::{Ipv4Addr, Ipv6Addr, SocketAddr},
         time::Duration,
     };
@@ -981,17 +982,17 @@ mod _socket {
 
     impl Read for &PySocket {
         fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-            (&mut &*self.sock()?).read(buf)
+            (&mut &*self.sock_snapshot()?).read(buf)
         }
     }
 
     impl Write for &PySocket {
         fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-            (&mut &*self.sock()?).write(buf)
+            (&mut &*self.sock_snapshot()?).write(buf)
         }
 
         fn flush(&mut self) -> std::io::Result<()> {
-            (&mut &*self.sock()?).flush()
+            (&mut &*self.sock_snapshot()?).flush()
         }
     }
 
@@ -1006,6 +1007,21 @@ mod _socket {
                 .ok_or_else(|| io::Error::from_raw_os_error(CLOSED_ERR))
         }
 
+        /// Snapshot the inner socket without holding `self.sock` across a syscall.
+        ///
+        /// Blocking accept/recv/send must not keep the RwLock: `close()` needs
+        /// the write side, and a lock held by a thread that did not survive
+        /// `fork()` stays locked forever in the child.
+        pub(crate) fn sock_snapshot(&self) -> io::Result<ManuallyDrop<Socket>> {
+            let guard = self.sock()?;
+            let fd = sock_fileno(&guard);
+            drop(guard);
+            // SAFETY: PySocket remains the owner of `fd`. Drop is suppressed so
+            // this wrapper does not close it. A concurrent close() may
+            // invalidate `fd`; the syscall then fails with EBADF.
+            Ok(ManuallyDrop::new(unsafe { sock_from_raw_unchecked(fd) }))
+        }
+
         fn init_inner(
             &self,
             family: i32,
@@ -1117,7 +1133,7 @@ mod _socket {
 
             loop {
                 if deadline.is_some() || matches!(wait_kind, SockWaitKind::Connect) {
-                    let sock = self.sock()?;
+                    let sock = self.sock_snapshot()?;
                     sock_wait_deadline(&sock, wait_kind, deadline.as_ref(), vm)?;
                 }
 
@@ -1441,7 +1457,7 @@ mod _socket {
         ) -> Result<(), IoOrPyException> {
             let sock_addr = self.extract_address(address, caller, vm)?;
 
-            let sock = self.sock()?;
+            let sock = self.sock_snapshot()?;
             let err = match vm.allow_threads(|| sock.connect(&sock_addr)) {
                 Ok(()) => return Ok(()),
                 Err(e) => e,
@@ -1461,7 +1477,7 @@ mod _socket {
 
             if wait_connect {
                 self.sock_op(vm, SockWaitKind::Connect, || {
-                    let sock = self.sock()?;
+                    let sock = self.sock_snapshot()?;
                     let err = sock.take_error()?;
                     match err {
                         Some(e) => Err(e),
@@ -1710,8 +1726,9 @@ mod _socket {
         ) -> Result<(RawSocket, PyObjectRef), IoOrPyException> {
             // Use accept_raw() instead of accept() to avoid socket2's set_common_flags()
             // which tries to set SO_NOSIGPIPE and fails with EINVAL on Unix domain sockets on macOS
-            let (sock, addr) =
-                self.sock_op(vm, SockWaitKind::Read, || self.sock()?.accept_raw())?;
+            let (sock, addr) = self.sock_op(vm, SockWaitKind::Read, || {
+                self.sock_snapshot()?.accept_raw()
+            })?;
             let fd = into_sock_fileno(sock);
             Ok((fd, get_addr_tuple(&addr, vm)))
         }
@@ -1728,9 +1745,9 @@ mod _socket {
             buffer
                 .try_reserve_exact(bufsize)
                 .map_err(|_| vm.no_memory_error())?;
-            let sock = self.sock()?;
             let n = self.sock_op(vm, SockWaitKind::Read, || {
-                sock.recv_with_flags(buffer.spare_capacity_mut(), flags)
+                self.sock_snapshot()?
+                    .recv_with_flags(buffer.spare_capacity_mut(), flags)
             })?;
             unsafe { buffer.set_len(n) };
             Ok(buffer)
@@ -1745,7 +1762,6 @@ mod _socket {
             vm: &VirtualMachine,
         ) -> Result<usize, IoOrPyException> {
             let flags = flags.unwrap_or(0);
-            let sock = self.sock()?;
 
             // Handle nbytes parameter
             let read_len = if let OptionalArg::Present(nbytes) = nbytes {
@@ -1759,7 +1775,8 @@ mod _socket {
 
             let mut scratch = alloc_recv_scratch(read_len, vm)?;
             let n = self.sock_op(vm, SockWaitKind::Read, || {
-                sock.recv_with_flags(&mut scratch.spare_capacity_mut()[..read_len], flags)
+                self.sock_snapshot()?
+                    .recv_with_flags(&mut scratch.spare_capacity_mut()[..read_len], flags)
             })?;
             unsafe { scratch.set_len(n) };
             buf.borrow_buf_mut()[..n].copy_from_slice(&scratch);
@@ -1782,7 +1799,7 @@ mod _socket {
                 .try_reserve_exact(bufsize)
                 .map_err(|_| vm.no_memory_error())?;
             let (n, addr) = self.sock_op(vm, SockWaitKind::Read, || {
-                self.sock()?
+                self.sock_snapshot()?
                     .recv_from_with_flags(buffer.spare_capacity_mut(), flags)
             })?;
             unsafe { buffer.set_len(n) };
@@ -1812,10 +1829,10 @@ mod _socket {
                 OptionalArg::Missing => buf.len(),
             };
             let flags = flags.unwrap_or(0);
-            let sock = self.sock()?;
             let mut scratch = alloc_recv_scratch(read_len, vm)?;
             let (n, addr) = self.sock_op(vm, SockWaitKind::Read, || {
-                sock.recv_from_with_flags(&mut scratch.spare_capacity_mut()[..read_len], flags)
+                self.sock_snapshot()?
+                    .recv_from_with_flags(&mut scratch.spare_capacity_mut()[..read_len], flags)
             })?;
             unsafe { scratch.set_len(n) };
             buf.borrow_buf_mut()[..n].copy_from_slice(&scratch);
@@ -1833,7 +1850,7 @@ mod _socket {
             let buf = bytes.borrow_buf_unlocked(vm)?;
             let buf = &*buf;
             self.sock_op(vm, SockWaitKind::Write, || {
-                self.sock()?.send_with_flags(buf, flags)
+                self.sock_snapshot()?.send_with_flags(buf, flags)
             })
         }
 
@@ -1858,7 +1875,7 @@ mod _socket {
                 let interval = deadline.as_ref().map(|d| d.time_until()).transpose()?;
                 self.sock_op_timeout_err(vm, SockWaitKind::Write, interval, || {
                     let subbuf = &buf[buf_offset..];
-                    buf_offset += self.sock()?.send_with_flags(subbuf, flags)?;
+                    buf_offset += self.sock_snapshot()?.send_with_flags(subbuf, flags)?;
                     Ok(())
                 })?;
                 vm.check_signals()?;
@@ -1890,7 +1907,7 @@ mod _socket {
             let buf = bytes.borrow_buf_unlocked(vm)?;
             let buf = &*buf;
             self.sock_op(vm, SockWaitKind::Write, || {
-                self.sock()?.send_to_with_flags(buf, &addr, flags)
+                self.sock_snapshot()?.send_to_with_flags(buf, &addr, flags)
             })
         }
 
@@ -1948,7 +1965,7 @@ mod _socket {
             }
 
             self.sock_op(vm, SockWaitKind::Write, || {
-                let sock = self.sock()?;
+                let sock = self.sock_snapshot()?;
                 sock.sendmsg(&msg, flags)
             })
             .map_err(|e| e.into_pyexception(vm))
@@ -1984,7 +2001,7 @@ mod _socket {
             let iv = iv.map(|iv| iv.borrow_buf().to_vec());
 
             self.sock_op(vm, SockWaitKind::Write, || {
-                let sock = self.sock()?;
+                let sock = self.sock_snapshot()?;
                 let fd = unsafe { BorrowedFd::borrow_raw(sock_fileno(&sock)) };
                 host_socket::sendmsg_afalg(fd, &buffers, op, iv.as_deref(), assoclen, flags)
             })
@@ -2017,7 +2034,7 @@ mod _socket {
 
             let msg = self
                 .sock_op(vm, SockWaitKind::Read, || {
-                    let sock = self.sock()?;
+                    let sock = self.sock_snapshot()?;
                     let fd = unsafe { std::os::fd::BorrowedFd::borrow_raw(sock_fileno(&sock)) };
                     host_socket::recvmsg(fd, bufsize, ancbufsize, flags)
                 })

crates/stdlib/src/ssl.rs

@@ -2527,7 +2527,7 @@ mod _ssl {
                 // Use select with the effective timeout
                 let py_socket: PyRef<PySocket> = self.io.socket_object(vm).try_into_value(vm)?;
                 let socket = py_socket
-                    .sock()
+                    .sock_snapshot()
                     .map_err(|e| vm.new_os_error(format!("Failed to get socket: {e}")))?;
 
                 sock_wait(&socket, wait_kind, timeout, vm)
@@ -2563,7 +2563,7 @@ mod _ssl {
 
                 let py_socket: PyRef<PySocket> = self.io.socket_object(vm).try_into_value(vm)?;
                 let socket = py_socket
-                    .sock()
+                    .sock_snapshot()
                     .map_err(|e| vm.new_os_error(format!("Failed to get socket: {e}")))?;
 
                 sock_wait(&socket, wait_kind, timeout, vm).map_err(|e| e.into_pyexception(vm))

extra_tests/snippets/stdlib_multiprocessing_fork_sharer.py

@@ -0,0 +1,31 @@
+"""Fork after the resource sharer thread is running must not hang the child.
+
+DupFd starts a daemon thread blocked in Listener.accept(). The child of a
+later fork used to stall in _ResourceSharer._afterfork while closing that
+listener.
+"""
+
+import multiprocessing
+import os
+
+
+def child():
+    print("child-ok", flush=True)
+
+
+if __name__ == "__main__":
+    multiprocessing.freeze_support()
+    if not hasattr(os, "fork") or "fork" not in multiprocessing.get_all_start_methods():
+        print("skipped (no fork)")
+        raise SystemExit(0)
+
+    from multiprocessing.resource_sharer import DupFd
+
+    DupFd(os.dup(0))
+    ctx = multiprocessing.get_context("fork")
+    proc = ctx.Process(target=child)
+    proc.daemon = True
+    proc.start()
+    proc.join(timeout=8)
+    assert proc.exitcode == 0, proc.exitcode
+    print("ok")

extra_tests/snippets/stdlib_socket_fork_accept.py

@@ -0,0 +1,55 @@
+"""Closing a listening socket after fork must not hang.
+
+A thread blocked in accept() must not keep the socket lock across the
+syscall. After fork the child copies that lock; close() in the child
+would wait forever for a thread that does not exist there.
+"""
+
+import os
+import socket
+import threading
+import time
+
+if not hasattr(os, "fork"):
+    print("skipped (no fork)")
+    raise SystemExit(0)
+
+listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+listener.bind(("127.0.0.1", 0))
+listener.listen(1)
+started = threading.Event()
+
+
+def accept_forever():
+    started.set()
+    try:
+        listener.accept()
+    except OSError:
+        pass
+
+
+t = threading.Thread(target=accept_forever, daemon=True)
+t.start()
+assert started.wait(5)
+
+# Give accept() time to enter the syscall.
+time.sleep(0.05)
+
+pid = os.fork()
+if pid == 0:
+    listener.close()
+    os._exit(0)
+
+_, status = os.waitpid(pid, 0)
+assert os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, status
+
+# Wake the parent accept thread. Closing the fd from another thread is
+# not a portable way to interrupt accept().
+addr = listener.getsockname()
+with socket.create_connection(addr, timeout=2):
+    pass
+t.join(timeout=5)
+assert not t.is_alive()
+listener.close()
+print("ok")