Skip to content

Commit

Permalink
Implement read_buf for a few more types
Browse files Browse the repository at this point in the history
Implement read_buf for TcpStream, Stdin, StdinLock, ChildStdout,
ChildStderr (and internally for AnonPipe, Handle, Socket), so
that it skips buffer initialization.

The other provided methods like read_to_string and read_to_end are
implemented in terms of read_buf and so benefit from the optimization
as well.

This commit also implements read_vectored and is_read_vectored where
applicable.
  • Loading branch information
tmiasko committed Feb 22, 2023
1 parent b869e84 commit 8ffb6af
Show file tree
Hide file tree
Showing 15 changed files with 205 additions and 25 deletions.
20 changes: 19 additions & 1 deletion library/std/src/fs/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@ use crate::io::prelude::*;

use crate::env;
use crate::fs::{self, File, OpenOptions};
use crate::io::{ErrorKind, SeekFrom};
use crate::io::{BorrowedBuf, ErrorKind, SeekFrom};
use crate::mem::MaybeUninit;
use crate::path::Path;
use crate::str;
use crate::sync::Arc;
Expand Down Expand Up @@ -401,6 +402,23 @@ fn file_test_io_seek_read_write() {
check!(fs::remove_file(&filename));
}

#[test]
fn file_test_read_buf() {
let tmpdir = tmpdir();
let filename = &tmpdir.join("test");
check!(fs::write(filename, &[1, 2, 3, 4]));

let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
let mut buf = BorrowedBuf::from(buf.as_mut_slice());
let mut file = check!(File::open(filename));
check!(file.read_buf(buf.unfilled()));
assert_eq!(buf.filled(), &[1, 2, 3, 4]);
// File::read_buf should omit buffer initialization.
assert_eq!(buf.init_len(), 4);

check!(fs::remove_file(filename));
}

#[test]
fn file_test_stat_is_correct_on_is_file() {
let tmpdir = tmpdir();
Expand Down
13 changes: 12 additions & 1 deletion library/std/src/io/stdio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use crate::io::prelude::*;
use crate::cell::{Cell, RefCell};
use crate::fmt;
use crate::fs::File;
use crate::io::{self, BufReader, IoSlice, IoSliceMut, LineWriter, Lines};
use crate::io::{self, BorrowedCursor, BufReader, IoSlice, IoSliceMut, LineWriter, Lines};
use crate::sync::atomic::{AtomicBool, Ordering};
use crate::sync::{Arc, Mutex, MutexGuard, OnceLock, ReentrantMutex, ReentrantMutexGuard};
use crate::sys::stdio;
Expand Down Expand Up @@ -97,6 +97,10 @@ impl Read for StdinRaw {
handle_ebadf(self.0.read(buf), 0)
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
handle_ebadf(self.0.read_buf(buf), ())
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
handle_ebadf(self.0.read_vectored(bufs), 0)
}
Expand Down Expand Up @@ -418,6 +422,9 @@ impl Read for Stdin {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.lock().read(buf)
}
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.lock().read_buf(buf)
}
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.lock().read_vectored(bufs)
}
Expand Down Expand Up @@ -450,6 +457,10 @@ impl Read for StdinLock<'_> {
self.inner.read(buf)
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.inner.read_buf(buf)
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.inner.read_vectored(bufs)
}
Expand Down
10 changes: 9 additions & 1 deletion library/std/src/net/tcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ mod tests;
use crate::io::prelude::*;

use crate::fmt;
use crate::io::{self, IoSlice, IoSliceMut};
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
use crate::iter::FusedIterator;
use crate::net::{Shutdown, SocketAddr, ToSocketAddrs};
use crate::sys_common::net as net_imp;
Expand Down Expand Up @@ -619,6 +619,10 @@ impl Read for TcpStream {
self.0.read(buf)
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.0.read_buf(buf)
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.0.read_vectored(bufs)
}
Expand Down Expand Up @@ -653,6 +657,10 @@ impl Read for &TcpStream {
self.0.read(buf)
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.0.read_buf(buf)
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.0.read_vectored(bufs)
}
Expand Down
24 changes: 23 additions & 1 deletion library/std/src/net/tcp/tests.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use crate::fmt;
use crate::io::prelude::*;
use crate::io::{ErrorKind, IoSlice, IoSliceMut};
use crate::io::{BorrowedBuf, ErrorKind, IoSlice, IoSliceMut};
use crate::mem::MaybeUninit;
use crate::net::test::{next_test_ip4, next_test_ip6};
use crate::net::*;
use crate::sync::mpsc::channel;
Expand Down Expand Up @@ -279,6 +280,27 @@ fn partial_read() {
})
}

#[test]
fn read_buf() {
each_ip(&mut |addr| {
let srv = t!(TcpListener::bind(&addr));
let t = thread::spawn(move || {
let mut s = t!(TcpStream::connect(&addr));
s.write_all(&[1, 2, 3, 4]).unwrap();
});

let mut s = t!(srv.accept()).0;
let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
let mut buf = BorrowedBuf::from(buf.as_mut_slice());
t!(s.read_buf(buf.unfilled()));
assert_eq!(buf.filled(), &[1, 2, 3, 4]);
// TcpStream::read_buf should omit buffer initialization.
assert_eq!(buf.init_len(), 4);

t.join().ok().expect("thread panicked");
})
}

#[test]
fn read_vectored() {
each_ip(&mut |addr| {
Expand Down
10 changes: 9 additions & 1 deletion library/std/src/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ use crate::convert::Infallible;
use crate::ffi::OsStr;
use crate::fmt;
use crate::fs;
use crate::io::{self, IoSlice, IoSliceMut};
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
use crate::num::NonZeroI32;
use crate::path::Path;
use crate::str;
Expand Down Expand Up @@ -354,6 +354,10 @@ impl Read for ChildStdout {
self.inner.read(buf)
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.inner.read_buf(buf)
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.inner.read_vectored(bufs)
}
Expand Down Expand Up @@ -419,6 +423,10 @@ impl Read for ChildStderr {
self.inner.read(buf)
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.inner.read_buf(buf)
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.inner.read_vectored(bufs)
}
Expand Down
31 changes: 30 additions & 1 deletion library/std/src/process/tests.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
use crate::io::prelude::*;

use super::{Command, Output, Stdio};
use crate::io::ErrorKind;
use crate::io::{BorrowedBuf, ErrorKind};
use crate::mem::MaybeUninit;
use crate::str;

fn known_command() -> Command {
Expand Down Expand Up @@ -119,6 +120,34 @@ fn stdin_works() {
assert_eq!(out, "foobar\n");
}

#[test]
#[cfg_attr(any(target_os = "vxworks"), ignore)]
fn child_stdout_read_buf() {
let mut cmd = shell_cmd();
if cfg!(target_os = "windows") {
cmd.arg("/C").arg("echo abc");
} else {
cmd.arg("-c").arg("echo abc");
};
cmd.stdin(Stdio::null());
cmd.stdout(Stdio::piped());
let child = cmd.spawn().unwrap();

let mut stdout = child.stdout.unwrap();
let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
let mut buf = BorrowedBuf::from(buf.as_mut_slice());
stdout.read_buf(buf.unfilled()).unwrap();

// ChildStdout::read_buf should omit buffer initialization.
if cfg!(target_os = "windows") {
assert_eq!(buf.filled(), b"abc\r\n");
assert_eq!(buf.init_len(), 5);
} else {
assert_eq!(buf.filled(), b"abc\n");
assert_eq!(buf.init_len(), 4);
};
}

#[test]
#[cfg_attr(any(target_os = "vxworks"), ignore)]
fn test_process_status() {
Expand Down
9 changes: 9 additions & 0 deletions library/std/src/sys/unix/fd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,15 @@ impl<'a> Read for &'a FileDesc {
fn read_buf(&mut self, cursor: BorrowedCursor<'_>) -> io::Result<()> {
(**self).read_buf(cursor)
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
(**self).read_vectored(bufs)
}

#[inline]
fn is_read_vectored(&self) -> bool {
(**self).is_read_vectored()
}
}

impl AsInner<OwnedFd> for FileDesc {
Expand Down
28 changes: 22 additions & 6 deletions library/std/src/sys/unix/net.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use crate::cmp;
use crate::ffi::CStr;
use crate::io::{self, IoSlice, IoSliceMut};
use crate::io::{self, BorrowedBuf, BorrowedCursor, IoSlice, IoSliceMut};
use crate::mem;
use crate::net::{Shutdown, SocketAddr};
use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
Expand Down Expand Up @@ -240,19 +240,35 @@ impl Socket {
self.0.duplicate().map(Socket)
}

fn recv_with_flags(&self, buf: &mut [u8], flags: c_int) -> io::Result<usize> {
fn recv_with_flags(&self, mut buf: BorrowedCursor<'_>, flags: c_int) -> io::Result<()> {
let ret = cvt(unsafe {
libc::recv(self.as_raw_fd(), buf.as_mut_ptr() as *mut c_void, buf.len(), flags)
libc::recv(
self.as_raw_fd(),
buf.as_mut().as_mut_ptr() as *mut c_void,
buf.capacity(),
flags,
)
})?;
Ok(ret as usize)
unsafe {
buf.advance(ret as usize);
}
Ok(())
}

pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
self.recv_with_flags(buf, 0)
let mut buf = BorrowedBuf::from(buf);
self.recv_with_flags(buf.unfilled(), 0)?;
Ok(buf.len())
}

pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
self.recv_with_flags(buf, MSG_PEEK)
let mut buf = BorrowedBuf::from(buf);
self.recv_with_flags(buf.unfilled(), MSG_PEEK)?;
Ok(buf.len())
}

pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.recv_with_flags(buf, 0)
}

pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
Expand Down
6 changes: 5 additions & 1 deletion library/std/src/sys/unix/pipe.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::io::{self, IoSlice, IoSliceMut};
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
use crate::mem;
use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
use crate::sys::fd::FileDesc;
Expand Down Expand Up @@ -49,6 +49,10 @@ impl AnonPipe {
self.0.read(buf)
}

pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
self.0.read_buf(buf)
}

pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.0.read_vectored(bufs)
}
Expand Down
6 changes: 5 additions & 1 deletion library/std/src/sys/unix/stdio.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::io::{self, IoSlice, IoSliceMut};
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
use crate::mem::ManuallyDrop;
use crate::os::unix::io::FromRawFd;
use crate::sys::fd::FileDesc;
Expand All @@ -18,6 +18,10 @@ impl io::Read for Stdin {
unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read(buf) }
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_buf(buf) }
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_vectored(bufs) }
}
Expand Down
6 changes: 5 additions & 1 deletion library/std/src/sys/unsupported/pipe.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::io::{self, IoSlice, IoSliceMut};
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};

pub struct AnonPipe(!);

Expand All @@ -7,6 +7,10 @@ impl AnonPipe {
self.0
}

pub fn read_buf(&self, _buf: BorrowedCursor<'_>) -> io::Result<()> {
self.0
}

pub fn read_vectored(&self, _bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
self.0
}
Expand Down
9 changes: 9 additions & 0 deletions library/std/src/sys/windows/handle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -327,7 +327,16 @@ impl<'a> Read for &'a Handle {
(**self).read(buf)
}

fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
(**self).read_buf(buf)
}

fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
(**self).read_vectored(bufs)
}

#[inline]
fn is_read_vectored(&self) -> bool {
(**self).is_read_vectored()
}
}
Loading

0 comments on commit 8ffb6af

Please sign in to comment.