Reuse unlocked buffers in os.write (#9107)

* Reuse unlocked buffers in os.write

Assisted-by: Codex:gpt-6

* Snapshot buffers without immutable storage

Assisted-by: Codex:gpt-6

crates/common/src/borrow.rs

@@ -39,10 +39,9 @@ impl_from!('a, T, BorrowedValue<'a, T>,
 impl<'a, T: ?Sized> BorrowedValue<'a, T> {
     /// Whether reaching the value holds a lock that other threads wait on.
     ///
-    /// An immutable object hands out a plain reference and answers `false`;
-    /// one whose storage can change hands out a guard. A caller about to wait
-    /// for something unrelated -- a peer, a file, a signal -- can use this to
-    /// decide whether it may keep the borrow for the duration.
+    /// A plain reference answers `false`, including references to externally
+    /// managed memory. This does not guarantee that the storage is immutable
+    /// or safe to borrow while another thread runs.
     #[must_use]
     pub const fn is_locked(&self) -> bool {
         !matches!(self, Self::Ref(_))

crates/vm/src/function/buffer.rs

@@ -1,7 +1,7 @@
 use crate::{
     AsObject, PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject,
     VirtualMachine,
-    builtins::{PyStr, PyStrRef},
+    builtins::{PyBytes, PyStr, PyStrRef},
     common::borrow::{BorrowedValue, BorrowedValueMut},
     protocol::{BufferFlags, PyBuffer},
 };
@@ -55,13 +55,12 @@ impl ArgBytesLike {
     /// `borrow_buf` may answer with a lock that every other thread writing to
     /// the same object waits on, and a thread waiting on a lock never reaches
     /// a safepoint, so keeping one across a wait for a peer, a pipe or a
-    /// signal stops the world from being stopped at all. Bytes reached that
-    /// way are copied out first. Bytes that lock nothing -- an immutable
-    /// object's -- are borrowed where they lie, which is all CPython holds in
-    /// either case.
+    /// signal stops the world from being stopped at all. Only storage owned
+    /// by `bytes` can be borrowed across that wait. Other exporters are copied:
+    /// an unlocked or readonly buffer may still alias writable raw memory.
     pub fn borrow_buf_unlocked(&self, vm: &VirtualMachine) -> PyResult<UnlockedBuf<'_>> {
         let borrowed = self.borrow_buf();
-        if !borrowed.is_locked() {
+        if !borrowed.is_locked() && self.source_object().downcast_ref::<PyBytes>().is_some() {
             return Ok(UnlockedBuf::Borrowed(borrowed));
         }
         let mut copy = Vec::new();

crates/vm/src/stdlib/_ctypes/function.rs

@@ -2020,3 +2020,41 @@ impl Py<PyCThunk> {
         self.callable.clone()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{TryFromObject, class::PyClassImpl, function::ArgBytesLike};
+
+    #[test]
+    fn unlocked_raw_buffers_are_copied_even_when_readonly() {
+        crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
+            let _class = RawMemoryBuffer::make_static_type();
+            for readonly in [false, true] {
+                let mut storage = *b"before";
+                let view = memoryview_at_impl(
+                    storage.as_mut_ptr() as usize,
+                    storage.len() as isize,
+                    readonly,
+                    vm,
+                )
+                .unwrap();
+                let buffer = ArgBytesLike::try_from_object(vm, view).unwrap();
+                let snapshot = buffer.borrow_buf_unlocked(vm).unwrap();
+                assert_ne!(snapshot.as_ptr(), storage.as_ptr());
+                storage.copy_from_slice(b"after!");
+                assert_eq!(&*snapshot, b"before");
+            }
+
+            let bytes = vm.ctx.new_bytes(b"bytes".to_vec());
+            let view = crate::builtins::PyMemoryView::from_object(bytes.as_object(), vm)
+                .unwrap()
+                .into_pyobject(vm);
+            for obj in [bytes.clone().into(), view] {
+                let buffer = ArgBytesLike::try_from_object(vm, obj).unwrap();
+                let borrowed = buffer.borrow_buf_unlocked(vm).unwrap();
+                assert_eq!(borrowed.as_ptr(), bytes.as_bytes().as_ptr());
+            }
+        });
+    }
+}

crates/vm/src/stdlib/os.rs

@@ -460,9 +460,9 @@ pub(super) mod _os {
 
     #[pyfunction]
     fn write(fd: crt_fd::Borrowed<'_>, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult<usize> {
-        let owned = data.with_ref(|b| b.to_vec());
+        let buffer = data.borrow_buf_unlocked(vm)?;
         loop {
-            match vm.allow_threads(|| crt_fd::write(fd, &owned)) {
+            match vm.allow_threads(|| crt_fd::write(fd, &buffer)) {
                 Ok(n) => return Ok(n),
                 Err(e) if e.raw_os_error() == Some(libc::EINTR) => {
                     vm.check_signals()?;