Skip to content

Commit 31bcec6

Browse files
sfacklersfackler
committed
Add vectored read and write support
This functionality has lived for a while in the tokio ecosystem, where it can improve performance by minimizing copies.
1 parent 4772dc8 commit 31bcec6

File tree

29 files changed

+1033
-92
lines changed

29 files changed

+1033
-92
lines changed

‎src/libstd/io/buffered.rs

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use io::prelude::*;
55
use cmp;
66
use error;
77
use fmt;
8-
use io::{self, Initializer, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom};
8+
use io::{self, Initializer, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom, IoVec, IoVecMut};
99
use memchr;
1010

1111
/// The `BufReader` struct adds buffering to any reader.
@@ -235,6 +235,19 @@ impl<R: Read> Read for BufReader<R> {
235235
Ok(nread)
236236
}
237237

238+
fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
239+
let total_len = bufs.iter().map(|b| b.as_slice().len()).sum::<usize>();
240+
if self.pos == self.cap && total_len >= self.buf.len() {
241+
return self.inner.read_vectored(bufs);
242+
}
243+
let nread = {
244+
let mut rem = self.fill_buf()?;
245+
rem.read_vectored(bufs)?
246+
};
247+
self.consume(nread);
248+
Ok(nread)
249+
}
250+
238251
// we can't skip unconditionally because of the large buffer case in read.
239252
unsafe fn initializer(&self) -> Initializer {
240253
self.inner.initializer()
@@ -577,9 +590,25 @@ impl<W: Write> Write for BufWriter<W> {
577590
self.panicked = false;
578591
r
579592
} else {
580-
Write::write(&mut self.buf, buf)
593+
self.buf.write(buf)
594+
}
595+
}
596+
597+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
598+
let total_len = bufs.iter().map(|b| b.as_slice().len()).sum::<usize>();
599+
if self.buf.len() + total_len > self.buf.capacity() {
600+
self.flush_buf()?;
601+
}
602+
if total_len >= self.buf.capacity() {
603+
self.panicked = true;
604+
let r = self.inner.as_mut().unwrap().write_vectored(bufs);
605+
self.panicked = false;
606+
r
607+
} else {
608+
self.buf.write_vectored(bufs)
581609
}
582610
}
611+
583612
fn flush(&mut self) -> io::Result<()> {
584613
self.flush_buf().and_then(|()| self.get_mut().flush())
585614
}

‎src/libstd/io/cursor.rs

Lines changed: 212 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use io::prelude::*;
22

33
use core::convert::TryInto;
44
use cmp;
5-
use io::{self, Initializer, SeekFrom, Error, ErrorKind};
5+
use io::{self, Initializer, SeekFrom, Error, ErrorKind, IoVec, IoVecMut};
66

77
/// A `Cursor` wraps an in-memory buffer and provides it with a
88
/// [`Seek`] implementation.
@@ -221,6 +221,19 @@ impl<T> Read for Cursor<T> where T: AsRef<[u8]> {
221221
Ok(n)
222222
}
223223

224+
fn read_vectored(&mut self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
225+
let mut nread = 0;
226+
for buf in bufs {
227+
let buf = buf.as_mut_slice();
228+
let n = self.read(buf)?;
229+
nread += n;
230+
if n < buf.len() {
231+
break;
232+
}
233+
}
234+
Ok(nread)
235+
}
236+
224237
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
225238
let n = buf.len();
226239
Read::read_exact(&mut self.fill_buf()?, buf)?;
@@ -251,6 +264,24 @@ fn slice_write(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<us
251264
Ok(amt)
252265
}
253266

267+
fn slice_write_vectored(
268+
pos_mut: &mut u64,
269+
slice: &mut [u8],
270+
bufs: &[IoVec<'_>],
271+
) -> io::Result<usize>
272+
{
273+
let mut nwritten = 0;
274+
for buf in bufs {
275+
let buf = buf.as_slice();
276+
let n = slice_write(pos_mut, slice, buf)?;
277+
nwritten += n;
278+
if n < buf.len() {
279+
break;
280+
}
281+
}
282+
Ok(nwritten)
283+
}
284+
254285
// Resizing write implementation
255286
fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usize> {
256287
let pos: usize = (*pos_mut).try_into().map_err(|_| {
@@ -278,12 +309,31 @@ fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usi
278309
Ok(buf.len())
279310
}
280311

312+
fn vec_write_vectored(
313+
pos_mut: &mut u64,
314+
vec: &mut Vec<u8>,
315+
bufs: &[IoVec<'_>],
316+
) -> io::Result<usize>
317+
{
318+
let mut nwritten = 0;
319+
for buf in bufs {
320+
nwritten += vec_write(pos_mut, vec, buf.as_slice())?;
321+
}
322+
Ok(nwritten)
323+
}
324+
281325
#[stable(feature = "rust1", since = "1.0.0")]
282326
impl<'a> Write for Cursor<&'a mut [u8]> {
283327
#[inline]
284328
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
285329
slice_write(&mut self.pos, self.inner, buf)
286330
}
331+
332+
#[inline]
333+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
334+
slice_write_vectored(&mut self.pos, self.inner, bufs)
335+
}
336+
287337
fn flush(&mut self) -> io::Result<()> { Ok(()) }
288338
}
289339

@@ -292,6 +342,11 @@ impl<'a> Write for Cursor<&'a mut Vec<u8>> {
292342
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
293343
vec_write(&mut self.pos, self.inner, buf)
294344
}
345+
346+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
347+
vec_write_vectored(&mut self.pos, self.inner, bufs)
348+
}
349+
295350
fn flush(&mut self) -> io::Result<()> { Ok(()) }
296351
}
297352

@@ -300,6 +355,11 @@ impl Write for Cursor<Vec<u8>> {
300355
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
301356
vec_write(&mut self.pos, &mut self.inner, buf)
302357
}
358+
359+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
360+
vec_write_vectored(&mut self.pos, &mut self.inner, bufs)
361+
}
362+
303363
fn flush(&mut self) -> io::Result<()> { Ok(()) }
304364
}
305365

@@ -309,21 +369,30 @@ impl Write for Cursor<Box<[u8]>> {
309369
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
310370
slice_write(&mut self.pos, &mut self.inner, buf)
311371
}
372+
373+
#[inline]
374+
fn write_vectored(&mut self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
375+
slice_write_vectored(&mut self.pos, &mut self.inner, bufs)
376+
}
377+
312378
fn flush(&mut self) -> io::Result<()> { Ok(()) }
313379
}
314380

315381
#[cfg(test)]
316382
mod tests {
317383
use io::prelude::*;
318-
use io::{Cursor, SeekFrom};
384+
use io::{Cursor, SeekFrom, IoVec, IoVecMut};
319385

320386
#[test]
321387
fn test_vec_writer() {
322388
let mut writer = Vec::new();
323389
assert_eq!(writer.write(&[0]).unwrap(), 1);
324390
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
325391
assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
326-
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
392+
assert_eq!(writer.write_vectored(
393+
&[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
394+
).unwrap(), 3);
395+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
327396
assert_eq!(writer, b);
328397
}
329398

@@ -333,7 +402,10 @@ mod tests {
333402
assert_eq!(writer.write(&[0]).unwrap(), 1);
334403
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
335404
assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
336-
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
405+
assert_eq!(writer.write_vectored(
406+
&[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
407+
).unwrap(), 3);
408+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
337409
assert_eq!(&writer.get_ref()[..], b);
338410
}
339411

@@ -344,7 +416,10 @@ mod tests {
344416
assert_eq!(writer.write(&[0]).unwrap(), 1);
345417
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
346418
assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
347-
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
419+
assert_eq!(writer.write_vectored(
420+
&[IoVec::new(&[]), IoVec::new(&[8, 9]), IoVec::new(&[10])],
421+
).unwrap(), 3);
422+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
348423
assert_eq!(&writer.get_ref()[..], b);
349424
}
350425

@@ -366,6 +441,26 @@ mod tests {
366441
assert_eq!(&**writer.get_ref(), b);
367442
}
368443

444+
#[test]
445+
fn test_box_slice_writer_vectored() {
446+
let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
447+
assert_eq!(writer.position(), 0);
448+
assert_eq!(writer.write_vectored(&[IoVec::new(&[0])]).unwrap(), 1);
449+
assert_eq!(writer.position(), 1);
450+
assert_eq!(
451+
writer.write_vectored(&[IoVec::new(&[1, 2, 3]), IoVec::new(&[4, 5, 6, 7])]).unwrap(),
452+
7,
453+
);
454+
assert_eq!(writer.position(), 8);
455+
assert_eq!(writer.write_vectored(&[]).unwrap(), 0);
456+
assert_eq!(writer.position(), 8);
457+
458+
assert_eq!(writer.write_vectored(&[IoVec::new(&[8, 9])]).unwrap(), 1);
459+
assert_eq!(writer.write_vectored(&[IoVec::new(&[10])]).unwrap(), 0);
460+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
461+
assert_eq!(&**writer.get_ref(), b);
462+
}
463+
369464
#[test]
370465
fn test_buf_writer() {
371466
let mut buf = [0 as u8; 9];
@@ -387,6 +482,31 @@ mod tests {
387482
assert_eq!(buf, b);
388483
}
389484

485+
#[test]
486+
fn test_buf_writer_vectored() {
487+
let mut buf = [0 as u8; 9];
488+
{
489+
let mut writer = Cursor::new(&mut buf[..]);
490+
assert_eq!(writer.position(), 0);
491+
assert_eq!(writer.write_vectored(&[IoVec::new(&[0])]).unwrap(), 1);
492+
assert_eq!(writer.position(), 1);
493+
assert_eq!(
494+
writer.write_vectored(
495+
&[IoVec::new(&[1, 2, 3]), IoVec::new(&[4, 5, 6, 7])],
496+
).unwrap(),
497+
7,
498+
);
499+
assert_eq!(writer.position(), 8);
500+
assert_eq!(writer.write_vectored(&[]).unwrap(), 0);
501+
assert_eq!(writer.position(), 8);
502+
503+
assert_eq!(writer.write_vectored(&[IoVec::new(&[8, 9])]).unwrap(), 1);
504+
assert_eq!(writer.write_vectored(&[IoVec::new(&[10])]).unwrap(), 0);
505+
}
506+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
507+
assert_eq!(buf, b);
508+
}
509+
390510
#[test]
391511
fn test_buf_writer_seek() {
392512
let mut buf = [0 as u8; 8];
@@ -447,6 +567,35 @@ mod tests {
447567
assert_eq!(reader.read(&mut buf).unwrap(), 0);
448568
}
449569

570+
#[test]
571+
fn test_mem_reader_vectored() {
572+
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7]);
573+
let mut buf = [];
574+
assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
575+
assert_eq!(reader.position(), 0);
576+
let mut buf = [0];
577+
assert_eq!(
578+
reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
579+
1,
580+
);
581+
assert_eq!(reader.position(), 1);
582+
let b: &[_] = &[0];
583+
assert_eq!(buf, b);
584+
let mut buf1 = [0; 4];
585+
let mut buf2 = [0; 4];
586+
assert_eq!(
587+
reader.read_vectored(
588+
&mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
589+
).unwrap(),
590+
7,
591+
);
592+
let b1: &[_] = &[1, 2, 3, 4];
593+
let b2: &[_] = &[5, 6, 7];
594+
assert_eq!(buf1, b1);
595+
assert_eq!(&buf2[..3], b2);
596+
assert_eq!(reader.read(&mut buf).unwrap(), 0);
597+
}
598+
450599
#[test]
451600
fn test_boxed_slice_reader() {
452601
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice());
@@ -469,6 +618,35 @@ mod tests {
469618
assert_eq!(reader.read(&mut buf).unwrap(), 0);
470619
}
471620

621+
#[test]
622+
fn test_boxed_slice_reader_vectored() {
623+
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice());
624+
let mut buf = [];
625+
assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
626+
assert_eq!(reader.position(), 0);
627+
let mut buf = [0];
628+
assert_eq!(
629+
reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
630+
1,
631+
);
632+
assert_eq!(reader.position(), 1);
633+
let b: &[_] = &[0];
634+
assert_eq!(buf, b);
635+
let mut buf1 = [0; 4];
636+
let mut buf2 = [0; 4];
637+
assert_eq!(
638+
reader.read_vectored(
639+
&mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
640+
).unwrap(),
641+
7,
642+
);
643+
let b1: &[_] = &[1, 2, 3, 4];
644+
let b2: &[_] = &[5, 6, 7];
645+
assert_eq!(buf1, b1);
646+
assert_eq!(&buf2[..3], b2);
647+
assert_eq!(reader.read(&mut buf).unwrap(), 0);
648+
}
649+
472650
#[test]
473651
fn read_to_end() {
474652
let mut reader = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7]);
@@ -499,6 +677,35 @@ mod tests {
499677
assert_eq!(reader.read(&mut buf).unwrap(), 0);
500678
}
501679

680+
#[test]
681+
fn test_slice_reader_vectored() {
682+
let in_buf = vec![0, 1, 2, 3, 4, 5, 6, 7];
683+
let reader = &mut &in_buf[..];
684+
let mut buf = [];
685+
assert_eq!(reader.read_vectored(&mut [IoVecMut::new(&mut buf)]).unwrap(), 0);
686+
let mut buf = [0];
687+
assert_eq!(
688+
reader.read_vectored(&mut [IoVecMut::new(&mut []), IoVecMut::new(&mut buf)]).unwrap(),
689+
1,
690+
);
691+
assert_eq!(reader.len(), 7);
692+
let b: &[_] = &[0];
693+
assert_eq!(buf, b);
694+
let mut buf1 = [0; 4];
695+
let mut buf2 = [0; 4];
696+
assert_eq!(
697+
reader.read_vectored(
698+
&mut [IoVecMut::new(&mut buf1), IoVecMut::new(&mut buf2)],
699+
).unwrap(),
700+
7,
701+
);
702+
let b1: &[_] = &[1, 2, 3, 4];
703+
let b2: &[_] = &[5, 6, 7];
704+
assert_eq!(buf1, b1);
705+
assert_eq!(&buf2[..3], b2);
706+
assert_eq!(reader.read(&mut buf).unwrap(), 0);
707+
}
708+
502709
#[test]
503710
fn test_read_exact() {
504711
let in_buf = vec![0, 1, 2, 3, 4, 5, 6, 7];

0 commit comments

Comments
 (0)