Skip to content

Commit 82dc127

Browse files
NoratriebNoratrieb
authored
Rollup merge of #108326 - tmiasko:read-buf, r=thomcc
Implement read_buf for a few more types 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.
2 parents caae551 + defa245 commit 82dc127

File tree

21 files changed

+250
-31
lines changed

21 files changed

+250
-31
lines changed

‎library/std/src/fs/tests.rs

+19-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@ use crate::io::prelude::*;
22

33
use crate::env;
44
use crate::fs::{self, File, OpenOptions};
5-
use crate::io::{ErrorKind, SeekFrom};
5+
use crate::io::{BorrowedBuf, ErrorKind, SeekFrom};
6+
use crate::mem::MaybeUninit;
67
use crate::path::Path;
78
use crate::str;
89
use crate::sync::Arc;
@@ -401,6 +402,23 @@ fn file_test_io_seek_read_write() {
401402
check!(fs::remove_file(&filename));
402403
}
403404

405+
#[test]
406+
fn file_test_read_buf() {
407+
let tmpdir = tmpdir();
408+
let filename = &tmpdir.join("test");
409+
check!(fs::write(filename, &[1, 2, 3, 4]));
410+
411+
let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
412+
let mut buf = BorrowedBuf::from(buf.as_mut_slice());
413+
let mut file = check!(File::open(filename));
414+
check!(file.read_buf(buf.unfilled()));
415+
assert_eq!(buf.filled(), &[1, 2, 3, 4]);
416+
// File::read_buf should omit buffer initialization.
417+
assert_eq!(buf.init_len(), 4);
418+
419+
check!(fs::remove_file(filename));
420+
}
421+
404422
#[test]
405423
fn file_test_stat_is_correct_on_is_file() {
406424
let tmpdir = tmpdir();

‎library/std/src/io/stdio.rs

+12-1
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use crate::io::prelude::*;
88
use crate::cell::{Cell, RefCell};
99
use crate::fmt;
1010
use crate::fs::File;
11-
use crate::io::{self, BufReader, IoSlice, IoSliceMut, LineWriter, Lines};
11+
use crate::io::{self, BorrowedCursor, BufReader, IoSlice, IoSliceMut, LineWriter, Lines};
1212
use crate::sync::atomic::{AtomicBool, Ordering};
1313
use crate::sync::{Arc, Mutex, MutexGuard, OnceLock, ReentrantMutex, ReentrantMutexGuard};
1414
use crate::sys::stdio;
@@ -97,6 +97,10 @@ impl Read for StdinRaw {
9797
handle_ebadf(self.0.read(buf), 0)
9898
}
9999

100+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
101+
handle_ebadf(self.0.read_buf(buf), ())
102+
}
103+
100104
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
101105
handle_ebadf(self.0.read_vectored(bufs), 0)
102106
}
@@ -418,6 +422,9 @@ impl Read for Stdin {
418422
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
419423
self.lock().read(buf)
420424
}
425+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
426+
self.lock().read_buf(buf)
427+
}
421428
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
422429
self.lock().read_vectored(bufs)
423430
}
@@ -450,6 +457,10 @@ impl Read for StdinLock<'_> {
450457
self.inner.read(buf)
451458
}
452459

460+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
461+
self.inner.read_buf(buf)
462+
}
463+
453464
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
454465
self.inner.read_vectored(bufs)
455466
}

‎library/std/src/net/tcp.rs

+9-1
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ mod tests;
66
use crate::io::prelude::*;
77

88
use crate::fmt;
9-
use crate::io::{self, IoSlice, IoSliceMut};
9+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
1010
use crate::iter::FusedIterator;
1111
use crate::net::{Shutdown, SocketAddr, ToSocketAddrs};
1212
use crate::sys_common::net as net_imp;
@@ -619,6 +619,10 @@ impl Read for TcpStream {
619619
self.0.read(buf)
620620
}
621621

622+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
623+
self.0.read_buf(buf)
624+
}
625+
622626
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
623627
self.0.read_vectored(bufs)
624628
}
@@ -653,6 +657,10 @@ impl Read for &TcpStream {
653657
self.0.read(buf)
654658
}
655659

660+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
661+
self.0.read_buf(buf)
662+
}
663+
656664
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
657665
self.0.read_vectored(bufs)
658666
}

‎library/std/src/net/tcp/tests.rs

+27-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use crate::fmt;
22
use crate::io::prelude::*;
3-
use crate::io::{ErrorKind, IoSlice, IoSliceMut};
3+
use crate::io::{BorrowedBuf, ErrorKind, IoSlice, IoSliceMut};
4+
use crate::mem::MaybeUninit;
45
use crate::net::test::{next_test_ip4, next_test_ip6};
56
use crate::net::*;
67
use crate::sync::mpsc::channel;
@@ -279,6 +280,31 @@ fn partial_read() {
279280
})
280281
}
281282

283+
#[test]
284+
fn read_buf() {
285+
each_ip(&mut |addr| {
286+
let srv = t!(TcpListener::bind(&addr));
287+
let t = thread::spawn(move || {
288+
let mut s = t!(TcpStream::connect(&addr));
289+
s.write_all(&[1, 2, 3, 4]).unwrap();
290+
});
291+
292+
let mut s = t!(srv.accept()).0;
293+
let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
294+
let mut buf = BorrowedBuf::from(buf.as_mut_slice());
295+
t!(s.read_buf(buf.unfilled()));
296+
assert_eq!(buf.filled(), &[1, 2, 3, 4]);
297+
298+
// FIXME: sgx uses default_read_buf that initializes the buffer.
299+
if cfg!(not(target_env = "sgx")) {
300+
// TcpStream::read_buf should omit buffer initialization.
301+
assert_eq!(buf.init_len(), 4);
302+
}
303+
304+
t.join().ok().expect("thread panicked");
305+
})
306+
}
307+
282308
#[test]
283309
fn read_vectored() {
284310
each_ip(&mut |addr| {

‎library/std/src/process.rs

+9-1
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ use crate::convert::Infallible;
110110
use crate::ffi::OsStr;
111111
use crate::fmt;
112112
use crate::fs;
113-
use crate::io::{self, IoSlice, IoSliceMut};
113+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
114114
use crate::num::NonZeroI32;
115115
use crate::path::Path;
116116
use crate::str;
@@ -354,6 +354,10 @@ impl Read for ChildStdout {
354354
self.inner.read(buf)
355355
}
356356

357+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
358+
self.inner.read_buf(buf)
359+
}
360+
357361
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
358362
self.inner.read_vectored(bufs)
359363
}
@@ -419,6 +423,10 @@ impl Read for ChildStderr {
419423
self.inner.read(buf)
420424
}
421425

426+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
427+
self.inner.read_buf(buf)
428+
}
429+
422430
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
423431
self.inner.read_vectored(bufs)
424432
}

‎library/std/src/process/tests.rs

+33-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
use crate::io::prelude::*;
22

33
use super::{Command, Output, Stdio};
4-
use crate::io::ErrorKind;
4+
use crate::io::{BorrowedBuf, ErrorKind};
5+
use crate::mem::MaybeUninit;
56
use crate::str;
67

78
fn known_command() -> Command {
@@ -119,6 +120,37 @@ fn stdin_works() {
119120
assert_eq!(out, "foobar\n");
120121
}
121122

123+
#[test]
124+
#[cfg_attr(any(target_os = "vxworks"), ignore)]
125+
fn child_stdout_read_buf() {
126+
let mut cmd = if cfg!(target_os = "windows") {
127+
let mut cmd = Command::new("cmd");
128+
cmd.arg("/C").arg("echo abc");
129+
cmd
130+
} else {
131+
let mut cmd = shell_cmd();
132+
cmd.arg("-c").arg("echo abc");
133+
cmd
134+
};
135+
cmd.stdin(Stdio::null());
136+
cmd.stdout(Stdio::piped());
137+
let child = cmd.spawn().unwrap();
138+
139+
let mut stdout = child.stdout.unwrap();
140+
let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
141+
let mut buf = BorrowedBuf::from(buf.as_mut_slice());
142+
stdout.read_buf(buf.unfilled()).unwrap();
143+
144+
// ChildStdout::read_buf should omit buffer initialization.
145+
if cfg!(target_os = "windows") {
146+
assert_eq!(buf.filled(), b"abc\r\n");
147+
assert_eq!(buf.init_len(), 5);
148+
} else {
149+
assert_eq!(buf.filled(), b"abc\n");
150+
assert_eq!(buf.init_len(), 4);
151+
};
152+
}
153+
122154
#[test]
123155
#[cfg_attr(any(target_os = "vxworks"), ignore)]
124156
fn test_process_status() {

‎library/std/src/sys/sgx/fd.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use fortanix_sgx_abi::Fd;
22

33
use super::abi::usercalls;
4-
use crate::io::{self, IoSlice, IoSliceMut};
4+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
55
use crate::mem;
66
use crate::sys::{AsInner, FromInner, IntoInner};
77

@@ -30,6 +30,10 @@ impl FileDesc {
3030
usercalls::read(self.fd, &mut [IoSliceMut::new(buf)])
3131
}
3232

33+
pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
34+
crate::io::default_read_buf(|b| self.read(b), buf)
35+
}
36+
3337
pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
3438
usercalls::read(self.fd, bufs)
3539
}

‎library/std/src/sys/sgx/net.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::error;
22
use crate::fmt;
3-
use crate::io::{self, IoSlice, IoSliceMut};
3+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
44
use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr, ToSocketAddrs};
55
use crate::sync::Arc;
66
use crate::sys::fd::FileDesc;
@@ -144,6 +144,10 @@ impl TcpStream {
144144
self.inner.inner.read(buf)
145145
}
146146

147+
pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
148+
self.inner.inner.read_buf(buf)
149+
}
150+
147151
pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
148152
self.inner.inner.read_vectored(bufs)
149153
}

‎library/std/src/sys/unix/fd.rs

+9
Original file line numberDiff line numberDiff line change
@@ -469,6 +469,15 @@ impl<'a> Read for &'a FileDesc {
469469
fn read_buf(&mut self, cursor: BorrowedCursor<'_>) -> io::Result<()> {
470470
(**self).read_buf(cursor)
471471
}
472+
473+
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
474+
(**self).read_vectored(bufs)
475+
}
476+
477+
#[inline]
478+
fn is_read_vectored(&self) -> bool {
479+
(**self).is_read_vectored()
480+
}
472481
}
473482

474483
impl AsInner<OwnedFd> for FileDesc {

‎library/std/src/sys/unix/net.rs

+22-6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::cmp;
22
use crate::ffi::CStr;
3-
use crate::io::{self, IoSlice, IoSliceMut};
3+
use crate::io::{self, BorrowedBuf, BorrowedCursor, IoSlice, IoSliceMut};
44
use crate::mem;
55
use crate::net::{Shutdown, SocketAddr};
66
use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
@@ -242,19 +242,35 @@ impl Socket {
242242
self.0.duplicate().map(Socket)
243243
}
244244

245-
fn recv_with_flags(&self, buf: &mut [u8], flags: c_int) -> io::Result<usize> {
245+
fn recv_with_flags(&self, mut buf: BorrowedCursor<'_>, flags: c_int) -> io::Result<()> {
246246
let ret = cvt(unsafe {
247-
libc::recv(self.as_raw_fd(), buf.as_mut_ptr() as *mut c_void, buf.len(), flags)
247+
libc::recv(
248+
self.as_raw_fd(),
249+
buf.as_mut().as_mut_ptr() as *mut c_void,
250+
buf.capacity(),
251+
flags,
252+
)
248253
})?;
249-
Ok(ret as usize)
254+
unsafe {
255+
buf.advance(ret as usize);
256+
}
257+
Ok(())
250258
}
251259

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

256266
pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
257-
self.recv_with_flags(buf, MSG_PEEK)
267+
let mut buf = BorrowedBuf::from(buf);
268+
self.recv_with_flags(buf.unfilled(), MSG_PEEK)?;
269+
Ok(buf.len())
270+
}
271+
272+
pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
273+
self.recv_with_flags(buf, 0)
258274
}
259275

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

‎library/std/src/sys/unix/pipe.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::io::{self, IoSlice, IoSliceMut};
1+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
22
use crate::mem;
33
use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
44
use crate::sys::fd::FileDesc;
@@ -49,6 +49,10 @@ impl AnonPipe {
4949
self.0.read(buf)
5050
}
5151

52+
pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
53+
self.0.read_buf(buf)
54+
}
55+
5256
pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
5357
self.0.read_vectored(bufs)
5458
}

‎library/std/src/sys/unix/stdio.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::io::{self, IoSlice, IoSliceMut};
1+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
22
use crate::mem::ManuallyDrop;
33
use crate::os::unix::io::FromRawFd;
44
use crate::sys::fd::FileDesc;
@@ -18,6 +18,10 @@ impl io::Read for Stdin {
1818
unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read(buf) }
1919
}
2020

21+
fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
22+
unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_buf(buf) }
23+
}
24+
2125
fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
2226
unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_vectored(bufs) }
2327
}

‎library/std/src/sys/unsupported/net.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use crate::fmt;
2-
use crate::io::{self, IoSlice, IoSliceMut};
2+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
33
use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr};
44
use crate::sys::unsupported;
55
use crate::time::Duration;
@@ -39,6 +39,10 @@ impl TcpStream {
3939
self.0
4040
}
4141

42+
pub fn read_buf(&self, _buf: BorrowedCursor<'_>) -> io::Result<()> {
43+
self.0
44+
}
45+
4246
pub fn read_vectored(&self, _: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
4347
self.0
4448
}

‎library/std/src/sys/unsupported/pipe.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::io::{self, IoSlice, IoSliceMut};
1+
use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
22

33
pub struct AnonPipe(!);
44

@@ -7,6 +7,10 @@ impl AnonPipe {
77
self.0
88
}
99

10+
pub fn read_buf(&self, _buf: BorrowedCursor<'_>) -> io::Result<()> {
11+
self.0
12+
}
13+
1014
pub fn read_vectored(&self, _bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
1115
self.0
1216
}

0 commit comments

Comments
 (0)