Tokio read_exact example

Cargo.toml

[package]
name = "example"
version = "0.1.0"
edition = "2021"

[dependencies]
bytes = "1.5"
tokio = { version = "1", features=[ "rt", "macros", "io-util" ] }

lib.rs

use std::io::{self, Read, Result};
use std::pin::Pin;
use std::task::{Context, Poll};

use bytes::{Buf, Bytes};
use tokio::io::{AsyncRead, ReadBuf};

struct TechnicallyCorrectReader {
    data: Bytes,
    frustration_counter: usize,
}

impl TechnicallyCorrectReader {
    fn new(data: Bytes) -> TechnicallyCorrectReader {
        TechnicallyCorrectReader {
            data: data.into(),
            frustration_counter: 0,
        }
    }
}

impl Read for TechnicallyCorrectReader {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        use std::cmp::min;
        if self.frustration_counter < 3 {
            self.frustration_counter += 1
        } else {
            self.frustration_counter = 0;
            return Err(io::Error::new(io::ErrorKind::Interrupted, "spurious error"));
        }

        let n = min(min(self.data.len(), self.frustration_counter), buf.len());
        buf[..n].copy_from_slice(&self.data[..n]);
        self.data.advance(n);
        Ok(n)
    }
}

struct AsyncTechnicallyCorrectReader {
    data: Bytes,
    frustration_counter: usize,
}

impl AsyncTechnicallyCorrectReader {
    fn new(data: Bytes) -> AsyncTechnicallyCorrectReader {
        AsyncTechnicallyCorrectReader {
            data: data.into(),
            frustration_counter: 0,
        }
    }
}

impl AsyncRead for AsyncTechnicallyCorrectReader {
    fn poll_read(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        use std::cmp::min;
        if self.frustration_counter < 3 {
            self.frustration_counter += 1
        } else {
            self.frustration_counter = 0;
            return Poll::Ready(Err(io::Error::new(
                io::ErrorKind::Interrupted,
                "spurious error",
            )));
        }

        let n = min(
            min(self.data.len(), self.frustration_counter),
            buf.remaining(),
        );
        buf.put_slice(&self.data[..n]);
        self.data.advance(n);
        Poll::Ready(Ok(()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use tokio::io::AsyncReadExt;

    #[test]
    fn test_behavior() {
        let data = b"this is a  very  boring test string";
        let mut buf = [0; 5];
        let mut reader = TechnicallyCorrectReader::new(data.to_vec().into());
        assert!(data.len() % 5 == 0);
        reader.read_exact(&mut buf).unwrap();
        assert_eq!(buf, *b"this ");
        reader.read_exact(&mut buf).unwrap();
        assert_eq!(buf, *b"is a ");
        reader.read_exact(&mut buf).unwrap();
        assert_eq!(buf, *b" very");
        reader.read_exact(&mut buf).unwrap();
        assert_eq!(buf, *b"  bor");
        reader.read_exact(&mut buf).unwrap();
        assert_eq!(buf, *b"ing t");
        reader.read_exact(&mut buf).unwrap();
        assert_eq!(buf, *b"est s");
        reader.read_exact(&mut buf).unwrap();
        assert_eq!(buf, *b"tring");
    }

    #[tokio::test]
    async fn test_behavior_async() {
        let data = b"this is a  very  boring test string";
        let mut buf = [0; 5];
        let mut reader = AsyncTechnicallyCorrectReader::new(data.to_vec().into());
        assert!(data.len() % 5 == 0);
        reader.read_exact(&mut buf).await.unwrap();
        assert_eq!(buf, *b"this ");
        reader.read_exact(&mut buf).await.unwrap();
        assert_eq!(buf, *b"is a ");
        reader.read_exact(&mut buf).await.unwrap();
        assert_eq!(buf, *b" very");
        reader.read_exact(&mut buf).await.unwrap();
        assert_eq!(buf, *b"  bor");
        reader.read_exact(&mut buf).await.unwrap();
        assert_eq!(buf, *b"ing t");
        reader.read_exact(&mut buf).await.unwrap();
        assert_eq!(buf, *b"est s");
        reader.read_exact(&mut buf).await.unwrap();
        assert_eq!(buf, *b"tring");
    }
}