io: impl AsyncSeek for BufReader/BufWriter/BufStream (#3491)

This commit is contained in:
Taiki Endo
2021-05-14 19:32:00 +09:00
committed by GitHub
parent e188e99ca3
commit 8c395dfe61
5 changed files with 798 additions and 5 deletions
+121 -3
View File
@@ -1,11 +1,11 @@
use crate::io::util::DEFAULT_BUF_SIZE;
use crate::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
use pin_project_lite::pin_project;
use std::io;
use std::io::{self, SeekFrom};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{cmp, fmt};
use std::{cmp, fmt, mem};
pin_project! {
/// The `BufReader` struct adds buffering to any reader.
@@ -30,6 +30,7 @@ pin_project! {
pub(super) buf: Box<[u8]>,
pub(super) pos: usize,
pub(super) cap: usize,
pub(super) seek_state: SeekState,
}
}
@@ -48,6 +49,7 @@ impl<R: AsyncRead> BufReader<R> {
buf: buffer.into_boxed_slice(),
pos: 0,
cap: 0,
seek_state: SeekState::Init,
}
}
@@ -141,6 +143,122 @@ impl<R: AsyncRead> AsyncBufRead for BufReader<R> {
}
}
#[derive(Debug, Clone, Copy)]
pub(super) enum SeekState {
/// start_seek has not been called.
Init,
/// start_seek has been called, but poll_complete has not yet been called.
Start(SeekFrom),
/// Waiting for completion of the first poll_complete in the `n.checked_sub(remainder).is_none()` branch.
PendingOverflowed(i64),
/// Waiting for completion of poll_complete.
Pending,
}
/// Seek to an offset, in bytes, in the underlying reader.
///
/// The position used for seeking with `SeekFrom::Current(_)` is the
/// position the underlying reader would be at if the `BufReader` had no
/// internal buffer.
///
/// Seeking always discards the internal buffer, even if the seek position
/// would otherwise fall within it. This guarantees that calling
/// `.into_inner()` immediately after a seek yields the underlying reader
/// at the same position.
///
/// See [`AsyncSeek`] for more details.
///
/// Note: In the edge case where you're seeking with `SeekFrom::Current(n)`
/// where `n` minus the internal buffer length overflows an `i64`, two
/// seeks will be performed instead of one. If the second seek returns
/// `Err`, the underlying reader will be left at the same position it would
/// have if you called `seek` with `SeekFrom::Current(0)`.
impl<R: AsyncRead + AsyncSeek> AsyncSeek for BufReader<R> {
fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> {
// We needs to call seek operation multiple times.
// And we should always call both start_seek and poll_complete,
// as start_seek alone cannot guarantee that the operation will be completed.
// poll_complete receives a Context and returns a Poll, so it cannot be called
// inside start_seek.
*self.project().seek_state = SeekState::Start(pos);
Ok(())
}
fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
let res = match mem::replace(self.as_mut().project().seek_state, SeekState::Init) {
SeekState::Init => {
// 1.x AsyncSeek recommends calling poll_complete before start_seek.
// We don't have to guarantee that the value returned by
// poll_complete called without start_seek is correct,
// so we'll return 0.
return Poll::Ready(Ok(0));
}
SeekState::Start(SeekFrom::Current(n)) => {
let remainder = (self.cap - self.pos) as i64;
// it should be safe to assume that remainder fits within an i64 as the alternative
// means we managed to allocate 8 exbibytes and that's absurd.
// But it's not out of the realm of possibility for some weird underlying reader to
// support seeking by i64::min_value() so we need to handle underflow when subtracting
// remainder.
if let Some(offset) = n.checked_sub(remainder) {
self.as_mut()
.get_pin_mut()
.start_seek(SeekFrom::Current(offset))?;
self.as_mut().get_pin_mut().poll_complete(cx)?
} else {
// seek backwards by our remainder, and then by the offset
self.as_mut()
.get_pin_mut()
.start_seek(SeekFrom::Current(-remainder))?;
if self.as_mut().get_pin_mut().poll_complete(cx)?.is_pending() {
*self.as_mut().project().seek_state = SeekState::PendingOverflowed(n);
return Poll::Pending;
}
// https://github.com/rust-lang/rust/pull/61157#issuecomment-495932676
self.as_mut().discard_buffer();
self.as_mut()
.get_pin_mut()
.start_seek(SeekFrom::Current(n))?;
self.as_mut().get_pin_mut().poll_complete(cx)?
}
}
SeekState::PendingOverflowed(n) => {
if self.as_mut().get_pin_mut().poll_complete(cx)?.is_pending() {
*self.as_mut().project().seek_state = SeekState::PendingOverflowed(n);
return Poll::Pending;
}
// https://github.com/rust-lang/rust/pull/61157#issuecomment-495932676
self.as_mut().discard_buffer();
self.as_mut()
.get_pin_mut()
.start_seek(SeekFrom::Current(n))?;
self.as_mut().get_pin_mut().poll_complete(cx)?
}
SeekState::Start(pos) => {
// Seeking with Start/End doesn't care about our buffer length.
self.as_mut().get_pin_mut().start_seek(pos)?;
self.as_mut().get_pin_mut().poll_complete(cx)?
}
SeekState::Pending => self.as_mut().get_pin_mut().poll_complete(cx)?,
};
match res {
Poll::Ready(res) => {
self.discard_buffer();
Poll::Ready(Ok(res))
}
Poll::Pending => {
*self.as_mut().project().seek_state = SeekState::Pending;
Poll::Pending
}
}
}
}
impl<R: AsyncRead + AsyncWrite> AsyncWrite for BufReader<R> {
fn poll_write(
self: Pin<&mut Self>,
+4
View File
@@ -94,9 +94,11 @@ impl<RW> From<BufWriter<BufReader<RW>>> for BufStream<RW> {
buf: rbuf,
pos,
cap,
seek_state: rseek_state,
},
buf: wbuf,
written,
seek_state: wseek_state,
} = b;
BufStream {
@@ -105,10 +107,12 @@ impl<RW> From<BufWriter<BufReader<RW>>> for BufStream<RW> {
inner,
buf: wbuf,
written,
seek_state: wseek_state,
},
buf: rbuf,
pos,
cap,
seek_state: rseek_state,
},
}
}
+60 -2
View File
@@ -1,9 +1,9 @@
use crate::io::util::DEFAULT_BUF_SIZE;
use crate::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
use pin_project_lite::pin_project;
use std::fmt;
use std::io::{self, Write};
use std::io::{self, SeekFrom, Write};
use std::pin::Pin;
use std::task::{Context, Poll};
@@ -34,6 +34,7 @@ pin_project! {
pub(super) inner: W,
pub(super) buf: Vec<u8>,
pub(super) written: usize,
pub(super) seek_state: SeekState,
}
}
@@ -50,6 +51,7 @@ impl<W: AsyncWrite> BufWriter<W> {
inner,
buf: Vec::with_capacity(cap),
written: 0,
seek_state: SeekState::Init,
}
}
@@ -142,6 +144,62 @@ impl<W: AsyncWrite> AsyncWrite for BufWriter<W> {
}
}
#[derive(Debug, Clone, Copy)]
pub(super) enum SeekState {
/// start_seek has not been called.
Init,
/// start_seek has been called, but poll_complete has not yet been called.
Start(SeekFrom),
/// Waiting for completion of poll_complete.
Pending,
}
/// Seek to the offset, in bytes, in the underlying writer.
///
/// Seeking always writes out the internal buffer before seeking.
impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> {
fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> {
// We need to flush the internal buffer before seeking.
// It receives a `Context` and returns a `Poll`, so it cannot be called
// inside `start_seek`.
*self.project().seek_state = SeekState::Start(pos);
Ok(())
}
fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
let pos = match self.seek_state {
SeekState::Init => {
return self.project().inner.poll_complete(cx);
}
SeekState::Start(pos) => Some(pos),
SeekState::Pending => None,
};
// Flush the internal buffer before seeking.
ready!(self.as_mut().flush_buf(cx))?;
let mut me = self.project();
if let Some(pos) = pos {
// Ensure previous seeks have finished before starting a new one
ready!(me.inner.as_mut().poll_complete(cx))?;
if let Err(e) = me.inner.as_mut().start_seek(pos) {
*me.seek_state = SeekState::Init;
return Poll::Ready(Err(e));
}
}
match me.inner.poll_complete(cx) {
Poll::Ready(res) => {
*me.seek_state = SeekState::Init;
Poll::Ready(res)
}
Poll::Pending => {
*me.seek_state = SeekState::Pending;
Poll::Pending
}
}
}
}
impl<W: AsyncWrite + AsyncRead> AsyncRead for BufWriter<W> {
fn poll_read(
self: Pin<&mut Self>,
+362
View File
@@ -0,0 +1,362 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]
// https://github.com/rust-lang/futures-rs/blob/1803948ff091b4eabf7f3bf39e16bbbdefca5cc8/futures/tests/io_buf_reader.rs
use futures::task::{noop_waker_ref, Context, Poll};
use std::cmp;
use std::io::{self, Cursor};
use std::pin::Pin;
use tokio::io::{
AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt, BufReader,
ReadBuf, SeekFrom,
};
macro_rules! run_fill_buf {
($reader:expr) => {{
let mut cx = Context::from_waker(noop_waker_ref());
loop {
if let Poll::Ready(x) = Pin::new(&mut $reader).poll_fill_buf(&mut cx) {
break x;
}
}
}};
}
struct MaybePending<'a> {
inner: &'a [u8],
ready_read: bool,
ready_fill_buf: bool,
}
impl<'a> MaybePending<'a> {
fn new(inner: &'a [u8]) -> Self {
Self {
inner,
ready_read: false,
ready_fill_buf: false,
}
}
}
impl AsyncRead for MaybePending<'_> {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
if self.ready_read {
self.ready_read = false;
Pin::new(&mut self.inner).poll_read(cx, buf)
} else {
self.ready_read = true;
cx.waker().wake_by_ref();
Poll::Pending
}
}
}
impl AsyncBufRead for MaybePending<'_> {
fn poll_fill_buf(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
if self.ready_fill_buf {
self.ready_fill_buf = false;
if self.inner.is_empty() {
return Poll::Ready(Ok(&[]));
}
let len = cmp::min(2, self.inner.len());
Poll::Ready(Ok(&self.inner[0..len]))
} else {
self.ready_fill_buf = true;
Poll::Pending
}
}
fn consume(mut self: Pin<&mut Self>, amt: usize) {
self.inner = &self.inner[amt..];
}
}
#[tokio::test]
async fn test_buffered_reader() {
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
let mut reader = BufReader::with_capacity(2, inner);
let mut buf = [0, 0, 0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 3);
assert_eq!(buf, [5, 6, 7]);
assert_eq!(reader.buffer(), []);
let mut buf = [0, 0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 2);
assert_eq!(buf, [0, 1]);
assert_eq!(reader.buffer(), []);
let mut buf = [0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 1);
assert_eq!(buf, [2]);
assert_eq!(reader.buffer(), [3]);
let mut buf = [0, 0, 0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 1);
assert_eq!(buf, [3, 0, 0]);
assert_eq!(reader.buffer(), []);
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 1);
assert_eq!(buf, [4, 0, 0]);
assert_eq!(reader.buffer(), []);
assert_eq!(reader.read(&mut buf).await.unwrap(), 0);
}
#[tokio::test]
async fn test_buffered_reader_seek() {
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
let mut reader = BufReader::with_capacity(2, Cursor::new(inner));
assert_eq!(reader.seek(SeekFrom::Start(3)).await.unwrap(), 3);
assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]);
assert!(reader
.seek(SeekFrom::Current(i64::min_value()))
.await
.is_err());
assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]);
assert_eq!(reader.seek(SeekFrom::Current(1)).await.unwrap(), 4);
assert_eq!(run_fill_buf!(reader).unwrap(), &[1, 2][..]);
Pin::new(&mut reader).consume(1);
assert_eq!(reader.seek(SeekFrom::Current(-2)).await.unwrap(), 3);
}
#[tokio::test]
async fn test_buffered_reader_seek_underflow() {
// gimmick reader that yields its position modulo 256 for each byte
struct PositionReader {
pos: u64,
}
impl AsyncRead for PositionReader {
fn poll_read(
mut self: Pin<&mut Self>,
_: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let b = buf.initialize_unfilled();
let len = b.len();
for x in b {
*x = self.pos as u8;
self.pos = self.pos.wrapping_add(1);
}
buf.advance(len);
Poll::Ready(Ok(()))
}
}
impl AsyncSeek for PositionReader {
fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> {
match pos {
SeekFrom::Start(n) => {
self.pos = n;
}
SeekFrom::Current(n) => {
self.pos = self.pos.wrapping_add(n as u64);
}
SeekFrom::End(n) => {
self.pos = u64::max_value().wrapping_add(n as u64);
}
}
Ok(())
}
fn poll_complete(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<u64>> {
Poll::Ready(Ok(self.pos))
}
}
let mut reader = BufReader::with_capacity(5, PositionReader { pos: 0 });
assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1, 2, 3, 4][..]);
assert_eq!(
reader.seek(SeekFrom::End(-5)).await.unwrap(),
u64::max_value() - 5
);
assert_eq!(run_fill_buf!(reader).unwrap().len(), 5);
// the following seek will require two underlying seeks
let expected = 9_223_372_036_854_775_802;
assert_eq!(
reader
.seek(SeekFrom::Current(i64::min_value()))
.await
.unwrap(),
expected
);
assert_eq!(run_fill_buf!(reader).unwrap().len(), 5);
// seeking to 0 should empty the buffer.
assert_eq!(reader.seek(SeekFrom::Current(0)).await.unwrap(), expected);
assert_eq!(reader.get_ref().pos, expected);
}
#[tokio::test]
async fn test_short_reads() {
/// A dummy reader intended at testing short-reads propagation.
struct ShortReader {
lengths: Vec<usize>,
}
impl AsyncRead for ShortReader {
fn poll_read(
mut self: Pin<&mut Self>,
_: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
if !self.lengths.is_empty() {
buf.advance(self.lengths.remove(0));
}
Poll::Ready(Ok(()))
}
}
let inner = ShortReader {
lengths: vec![0, 1, 2, 0, 1, 0],
};
let mut reader = BufReader::new(inner);
let mut buf = [0, 0];
assert_eq!(reader.read(&mut buf).await.unwrap(), 0);
assert_eq!(reader.read(&mut buf).await.unwrap(), 1);
assert_eq!(reader.read(&mut buf).await.unwrap(), 2);
assert_eq!(reader.read(&mut buf).await.unwrap(), 0);
assert_eq!(reader.read(&mut buf).await.unwrap(), 1);
assert_eq!(reader.read(&mut buf).await.unwrap(), 0);
assert_eq!(reader.read(&mut buf).await.unwrap(), 0);
}
#[tokio::test]
async fn maybe_pending() {
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
let mut reader = BufReader::with_capacity(2, MaybePending::new(inner));
let mut buf = [0, 0, 0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 3);
assert_eq!(buf, [5, 6, 7]);
assert_eq!(reader.buffer(), []);
let mut buf = [0, 0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 2);
assert_eq!(buf, [0, 1]);
assert_eq!(reader.buffer(), []);
let mut buf = [0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 1);
assert_eq!(buf, [2]);
assert_eq!(reader.buffer(), [3]);
let mut buf = [0, 0, 0];
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 1);
assert_eq!(buf, [3, 0, 0]);
assert_eq!(reader.buffer(), []);
let nread = reader.read(&mut buf).await.unwrap();
assert_eq!(nread, 1);
assert_eq!(buf, [4, 0, 0]);
assert_eq!(reader.buffer(), []);
assert_eq!(reader.read(&mut buf).await.unwrap(), 0);
}
#[tokio::test]
async fn maybe_pending_buf_read() {
let inner = MaybePending::new(&[0, 1, 2, 3, 1, 0]);
let mut reader = BufReader::with_capacity(2, inner);
let mut v = Vec::new();
reader.read_until(3, &mut v).await.unwrap();
assert_eq!(v, [0, 1, 2, 3]);
v.clear();
reader.read_until(1, &mut v).await.unwrap();
assert_eq!(v, [1]);
v.clear();
reader.read_until(8, &mut v).await.unwrap();
assert_eq!(v, [0]);
v.clear();
reader.read_until(9, &mut v).await.unwrap();
assert_eq!(v, []);
}
// https://github.com/rust-lang/futures-rs/pull/1573#discussion_r281162309
#[tokio::test]
async fn maybe_pending_seek() {
struct MaybePendingSeek<'a> {
inner: Cursor<&'a [u8]>,
ready: bool,
seek_res: Option<io::Result<()>>,
}
impl<'a> MaybePendingSeek<'a> {
fn new(inner: &'a [u8]) -> Self {
Self {
inner: Cursor::new(inner),
ready: true,
seek_res: None,
}
}
}
impl AsyncRead for MaybePendingSeek<'_> {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.inner).poll_read(cx, buf)
}
}
impl AsyncBufRead for MaybePendingSeek<'_> {
fn poll_fill_buf(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<io::Result<&[u8]>> {
let this: *mut Self = &mut *self as *mut _;
Pin::new(&mut unsafe { &mut *this }.inner).poll_fill_buf(cx)
}
fn consume(mut self: Pin<&mut Self>, amt: usize) {
Pin::new(&mut self.inner).consume(amt)
}
}
impl AsyncSeek for MaybePendingSeek<'_> {
fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> {
self.seek_res = Some(Pin::new(&mut self.inner).start_seek(pos));
Ok(())
}
fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
if self.ready {
self.ready = false;
self.seek_res.take().unwrap_or(Ok(()))?;
Pin::new(&mut self.inner).poll_complete(cx)
} else {
self.ready = true;
cx.waker().wake_by_ref();
Poll::Pending
}
}
}
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
let mut reader = BufReader::with_capacity(2, MaybePendingSeek::new(inner));
assert_eq!(reader.seek(SeekFrom::Current(3)).await.unwrap(), 3);
assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]);
assert!(reader
.seek(SeekFrom::Current(i64::min_value()))
.await
.is_err());
assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]);
assert_eq!(reader.seek(SeekFrom::Current(1)).await.unwrap(), 4);
assert_eq!(run_fill_buf!(reader).unwrap(), &[1, 2][..]);
Pin::new(&mut reader).consume(1);
assert_eq!(reader.seek(SeekFrom::Current(-2)).await.unwrap(), 3);
}
+251
View File
@@ -0,0 +1,251 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]
// https://github.com/rust-lang/futures-rs/blob/1803948ff091b4eabf7f3bf39e16bbbdefca5cc8/futures/tests/io_buf_writer.rs
use futures::task::{Context, Poll};
use std::io::{self, Cursor};
use std::pin::Pin;
use tokio::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, BufWriter, SeekFrom};
struct MaybePending {
inner: Vec<u8>,
ready: bool,
}
impl MaybePending {
fn new(inner: Vec<u8>) -> Self {
Self {
inner,
ready: false,
}
}
}
impl AsyncWrite for MaybePending {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
if self.ready {
self.ready = false;
Pin::new(&mut self.inner).poll_write(cx, buf)
} else {
self.ready = true;
cx.waker().wake_by_ref();
Poll::Pending
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.inner).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.inner).poll_shutdown(cx)
}
}
#[tokio::test]
async fn buf_writer() {
let mut writer = BufWriter::with_capacity(2, Vec::new());
writer.write(&[0, 1]).await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(*writer.get_ref(), [0, 1]);
writer.write(&[2]).await.unwrap();
assert_eq!(writer.buffer(), [2]);
assert_eq!(*writer.get_ref(), [0, 1]);
writer.write(&[3]).await.unwrap();
assert_eq!(writer.buffer(), [2, 3]);
assert_eq!(*writer.get_ref(), [0, 1]);
writer.flush().await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(*writer.get_ref(), [0, 1, 2, 3]);
writer.write(&[4]).await.unwrap();
writer.write(&[5]).await.unwrap();
assert_eq!(writer.buffer(), [4, 5]);
assert_eq!(*writer.get_ref(), [0, 1, 2, 3]);
writer.write(&[6]).await.unwrap();
assert_eq!(writer.buffer(), [6]);
assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5]);
writer.write(&[7, 8]).await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8]);
writer.write(&[9, 10, 11]).await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
writer.flush().await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
}
#[tokio::test]
async fn buf_writer_inner_flushes() {
let mut w = BufWriter::with_capacity(3, Vec::new());
w.write(&[0, 1]).await.unwrap();
assert_eq!(*w.get_ref(), []);
w.flush().await.unwrap();
let w = w.into_inner();
assert_eq!(w, [0, 1]);
}
#[tokio::test]
async fn buf_writer_seek() {
let mut w = BufWriter::with_capacity(3, Cursor::new(Vec::new()));
w.write_all(&[0, 1, 2, 3, 4, 5]).await.unwrap();
w.write_all(&[6, 7]).await.unwrap();
assert_eq!(w.seek(SeekFrom::Current(0)).await.unwrap(), 8);
assert_eq!(&w.get_ref().get_ref()[..], &[0, 1, 2, 3, 4, 5, 6, 7][..]);
assert_eq!(w.seek(SeekFrom::Start(2)).await.unwrap(), 2);
w.write_all(&[8, 9]).await.unwrap();
w.flush().await.unwrap();
assert_eq!(&w.into_inner().into_inner()[..], &[0, 1, 8, 9, 4, 5, 6, 7]);
}
#[tokio::test]
async fn maybe_pending_buf_writer() {
let mut writer = BufWriter::with_capacity(2, MaybePending::new(Vec::new()));
writer.write(&[0, 1]).await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(&writer.get_ref().inner, &[0, 1]);
writer.write(&[2]).await.unwrap();
assert_eq!(writer.buffer(), [2]);
assert_eq!(&writer.get_ref().inner, &[0, 1]);
writer.write(&[3]).await.unwrap();
assert_eq!(writer.buffer(), [2, 3]);
assert_eq!(&writer.get_ref().inner, &[0, 1]);
writer.flush().await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(&writer.get_ref().inner, &[0, 1, 2, 3]);
writer.write(&[4]).await.unwrap();
writer.write(&[5]).await.unwrap();
assert_eq!(writer.buffer(), [4, 5]);
assert_eq!(&writer.get_ref().inner, &[0, 1, 2, 3]);
writer.write(&[6]).await.unwrap();
assert_eq!(writer.buffer(), [6]);
assert_eq!(writer.get_ref().inner, &[0, 1, 2, 3, 4, 5]);
writer.write(&[7, 8]).await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(writer.get_ref().inner, &[0, 1, 2, 3, 4, 5, 6, 7, 8]);
writer.write(&[9, 10, 11]).await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(
writer.get_ref().inner,
&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
);
writer.flush().await.unwrap();
assert_eq!(writer.buffer(), []);
assert_eq!(
&writer.get_ref().inner,
&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
);
}
#[tokio::test]
async fn maybe_pending_buf_writer_inner_flushes() {
let mut w = BufWriter::with_capacity(3, MaybePending::new(Vec::new()));
w.write(&[0, 1]).await.unwrap();
assert_eq!(&w.get_ref().inner, &[]);
w.flush().await.unwrap();
let w = w.into_inner().inner;
assert_eq!(w, [0, 1]);
}
#[tokio::test]
async fn maybe_pending_buf_writer_seek() {
struct MaybePendingSeek {
inner: Cursor<Vec<u8>>,
ready_write: bool,
ready_seek: bool,
seek_res: Option<io::Result<()>>,
}
impl MaybePendingSeek {
fn new(inner: Vec<u8>) -> Self {
Self {
inner: Cursor::new(inner),
ready_write: false,
ready_seek: false,
seek_res: None,
}
}
}
impl AsyncWrite for MaybePendingSeek {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
if self.ready_write {
self.ready_write = false;
Pin::new(&mut self.inner).poll_write(cx, buf)
} else {
self.ready_write = true;
cx.waker().wake_by_ref();
Poll::Pending
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.inner).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.inner).poll_shutdown(cx)
}
}
impl AsyncSeek for MaybePendingSeek {
fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> {
self.seek_res = Some(Pin::new(&mut self.inner).start_seek(pos));
Ok(())
}
fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
if self.ready_seek {
self.ready_seek = false;
self.seek_res.take().unwrap_or(Ok(()))?;
Pin::new(&mut self.inner).poll_complete(cx)
} else {
self.ready_seek = true;
cx.waker().wake_by_ref();
Poll::Pending
}
}
}
let mut w = BufWriter::with_capacity(3, MaybePendingSeek::new(Vec::new()));
w.write_all(&[0, 1, 2, 3, 4, 5]).await.unwrap();
w.write_all(&[6, 7]).await.unwrap();
assert_eq!(w.seek(SeekFrom::Current(0)).await.unwrap(), 8);
assert_eq!(
&w.get_ref().inner.get_ref()[..],
&[0, 1, 2, 3, 4, 5, 6, 7][..]
);
assert_eq!(w.seek(SeekFrom::Start(2)).await.unwrap(), 2);
w.write_all(&[8, 9]).await.unwrap();
w.flush().await.unwrap();
assert_eq!(
&w.into_inner().inner.into_inner()[..],
&[0, 1, 8, 9, 4, 5, 6, 7]
);
}