tokio-io-utility-0.7.6/.cargo_vcs_info.json0000644000000001360000000000100143010ustar { "git": { "sha1": "903ac3811bb8b4a505c167146d31041958c2f896" }, "path_in_vcs": "" }tokio-io-utility-0.7.6/.github/dependabot.yml000064400000000000000000000004360072674642500173140ustar 00000000000000version: 2 updates: - package-ecosystem: "github-actions" # Workflow files stored in the # default location of `.github/workflows` directory: "/" schedule: interval: "daily" - package-ecosystem: "cargo" directory: "/" schedule: interval: "daily" tokio-io-utility-0.7.6/.github/workflows/rust.yml000064400000000000000000000027770072674642500202530ustar 00000000000000name: Rust env: CARGO_TERM_COLOR: always on: push: paths-ignore: - 'README.md' - 'LICENSE' - '.gitignore' pull_request: paths-ignore: - 'README.md' - 'LICENSE' - '.gitignore' jobs: check: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 - uses: actions/cache@v3 with: path: | ~/.cargo/bin/ ~/.cargo/registry/index/ ~/.cargo/registry/cache/ ~/.cargo/git/db/ target/ key: ${{ github.event.repository.name }}-${{ runner.os }}-cargo-check-v3 - name: Install latest stable cargo uses: actions-rs/toolchain@v1 with: toolchain: stable override: true components: rustfmt, clippy - name: Install cargo nightly run: rustup install nightly - name: Run check run: | cargo fmt --all -- --check cargo clippy --all ./build_doc.sh build: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 - uses: actions/cache@v3 with: path: | ~/.cargo/bin/ ~/.cargo/registry/index/ ~/.cargo/registry/cache/ ~/.cargo/git/db/ target/ key: ${{ github.event.repository.name }}-${{ runner.os }}-cargo-test-v6 - name: Install latest nightly uses: actions-rs/toolchain@v1 with: toolchain: nightly override: true components: rust-src, miri - name: Run tests run: ./run_test.sh tokio-io-utility-0.7.6/.gitignore000064400000000000000000000006310072674642500151110ustar 00000000000000# Generated by Cargo # will have compiled files and executables /target/ # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html Cargo.lock # These are backup files generated by rustfmt **/*.rs.bk # Added by cargo # # already existing elements were commented out /target #Cargo.lock tokio-io-utility-0.7.6/Cargo.toml0000644000000023250000000000100123010ustar # THIS FILE IS AUTOMATICALLY GENERATED BY CARGO # # When uploading crates to the registry Cargo will automatically # "normalize" Cargo.toml files for maximal compatibility # with all versions of Cargo and also rewrite `path` dependencies # to registry (e.g., crates.io) dependencies. # # If you are reading this file be aware that the original Cargo.toml # will likely look very different (and much more reasonable). # See Cargo.toml.orig for the original contents. [package] edition = "2021" rust-version = "1.60" name = "tokio-io-utility" version = "0.7.6" authors = ["Jiahao XU "] description = "Some helper functions for tokio::io." keywords = [ "tokio", "async", ] categories = ["asynchronous"] license = "MIT" repository = "https://github.com/openssh-rust/tokio-io-utility.git" resolver = "2" [package.metadata.docs.rs] all-features = true rustdoc-args = [ "--cfg", "docsrs", ] [dependencies.bytes] version = "1.2.1" optional = true [dependencies.tokio] version = "1.0" features = ["io-util"] [dev-dependencies.rayon] version = "1.5.1" [dev-dependencies.tokio] version = "1" features = ["rt"] [dev-dependencies.tokio-pipe] version = "0.2.5" [features] read-exact-to-bytes = ["bytes"] tokio-io-utility-0.7.6/Cargo.toml.orig000064400000000000000000000017700072674642500160150ustar 00000000000000[package] name = "tokio-io-utility" version = "0.7.6" edition = "2021" rust-version = "1.60" authors = ["Jiahao XU "] license = "MIT" description = "Some helper functions for tokio::io." repository = "https://github.com/openssh-rust/tokio-io-utility.git" keywords = ["tokio", "async"] categories = ["asynchronous"] # docs.rs-specific configuration, shamelessly copied from # https://stackoverflow.com/a/61417700/8375400. # # To test locally, use ` ./build_doc.sh` [package.metadata.docs.rs] # document all features all-features = true # defines the configuration attribute `docsrs` rustdoc-args = ["--cfg", "docsrs"] [features] read-exact-to-bytes = ["bytes"] [dependencies] tokio = { version = "1.0", features = ["io-util"] } bytes = { version = "1.2.1", optional = true } [dev-dependencies] tokio-pipe = "0.2.5" # Do not enable feature macro here! # It will cause the compilation of tests with address sanitizer to fail. tokio = { version = "1", features = ["rt"] } rayon = "1.5.1" tokio-io-utility-0.7.6/LICENSE000064400000000000000000000020520072674642500141250ustar 00000000000000MIT License Copyright (c) 2021 Jiahao XU Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. tokio-io-utility-0.7.6/README.md000064400000000000000000000014350072674642500144030ustar 00000000000000# tokio-io-utility [![Rust](https://github.com/NobodyXu/tokio-io-utility/actions/workflows/rust.yml/badge.svg)](https://github.com/NobodyXu/tokio-io-utility/actions/workflows/rust.yml) [![crate.io downloads](https://img.shields.io/crates/d/tokio-io-utility)](https://crates.io/crates/tokio-io-utility) [![crate.io version](https://img.shields.io/crates/v/tokio-io-utility)](https://crates.io/crates/tokio-io-utility) [![docs](https://docs.rs/tokio-io-utility/badge.svg)](https://docs.rs/tokio-io-utility) Provide some helper functions for - reading into `Vec`, `Bytes` or other container - writing `IoSlice`s - Initializing `[MaybeUninit>]` - `ReusableIoSlices` to reuse a `Vec` of `IoSlice<'_>` without worrying about lifetime. ## How to run tests ``` ./run_test.sh ``` tokio-io-utility-0.7.6/build_doc.sh000075500000000000000000000001250072674642500154020ustar 00000000000000#!/bin/sh export RUSTDOCFLAGS="--cfg docsrs" exec cargo +nightly doc --all-features tokio-io-utility-0.7.6/run_test.sh000075500000000000000000000030660072674642500153300ustar 00000000000000#!/bin/bash -ex arch=$(uname -m) if [ "$arch" = "arm64" ]; then arch=aarch64 fi os=unknown-linux-gnu if uname -v | grep -s Darwin >/dev/null; then os=apple-darwin fi target="$arch-$os" rep=$(seq 1 10) for _ in $rep; do cargo test --all-features -- --nocapture done export RUSTFLAGS='-Zsanitizer=address' cargo +nightly test --all-features async_read_utility -- --nocapture cargo +nightly test --all-features async_write_utility -- --nocapture cargo +nightly test --all-features io_slice_ext -- --nocapture cargo +nightly test --all-features reusable_io_slices -- --nocapture cargo +nightly test --all-features init_maybeuninit_io_slice -- --nocapture for _ in $rep; do cargo +nightly test --all-features queue -- --nocapture done # thread sanitizer reports false positive in crossbeam #export RUSTFLAGS='-Zsanitizer=thread' # #for _ in $rep; do # cargo +nightly test \ # -Z build-std \ # --target "$target" \ # --all-features queue::tests::test_par -- --nocapture #done unset RUSTFLAGS export MIRIFLAGS="-Zmiri-disable-isolation" cargo +nightly miri test \ -Z build-std \ --target "$target" \ --all-features io_slice_ext -- --nocapture cargo +nightly miri test \ -Z build-std \ --target "$target" \ --all-features reusable_io_slices -- --nocapture cargo +nightly miri test \ -Z build-std \ --target "$target" \ --all-features init_maybeuninit_io_slice -- --nocapture exec cargo +nightly miri test \ -Z build-std \ --target "$target" \ --all-features queue::tests::test_seq -- --nocapture tokio-io-utility-0.7.6/src/async_read_utility/bytes_impl.rs000064400000000000000000000173030072674642500223240ustar 00000000000000use super::{read_to_container_rng, Container, ReadToContainerRngFuture}; use std::{ future::Future, io::Result, marker::Unpin, mem::MaybeUninit, pin::Pin, task::{Context, Poll}, }; use bytes::{BufMut, BytesMut}; use tokio::io::{AsyncRead, ReadBuf}; #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] impl Container for BytesMut { fn reserve(&mut self, n: usize) { BytesMut::reserve(self, n) } fn len(&self) -> usize { BytesMut::len(self) } fn capacity(&self) -> usize { BytesMut::capacity(self) } unsafe fn spare_mut(&mut self) -> &mut [MaybeUninit] { self.chunk_mut().as_uninit_slice_mut() } unsafe fn advance(&mut self, n: usize) { self.advance_mut(n) } } /// Returned future of [`read_exact_to_bytes`]. #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] pub struct ReadExactToBytesFuture<'a, R: ?Sized>(ReadToBytesRngFuture<'a, R>); impl Future for ReadExactToBytesFuture<'_, R> { type Output = Result<()>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::new(&mut self.0).poll(cx) } } /// * `nread` - bytes to read in /// /// Return [`std::io::ErrorKind::UnexpectedEof`] on Eof. /// /// NOTE that this function does not modify any existing data. /// /// # Cancel safety /// /// It is cancel safe and dropping the returned future will not stop the /// wakeup from happening. #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] pub fn read_exact_to_bytes<'a, R: AsyncRead + ?Sized + Unpin>( reader: &'a mut R, bytes: &'a mut BytesMut, nread: usize, ) -> ReadExactToBytesFuture<'a, R> { ReadExactToBytesFuture(read_to_bytes_rng(reader, bytes, nread..=nread)) } /// Returned future of [`read_to_bytes_rng`]. #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] pub struct ReadToBytesRngFuture<'a, Reader: ?Sized>(ReadToContainerRngFuture<'a, BytesMut, Reader>); impl Future for ReadToBytesRngFuture<'_, Reader> { type Output = Result<()>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::new(&mut self.0).poll(cx) } } /// * `rng` - The start of the range specify the minimum of bytes to read in, /// while the end of the range specify the maximum of bytes that /// can be read in. /// If the lower bound is not specified, it is default to 0. /// If the upper bound is not specified, it is default to the /// capacity of `bytes`. /// The lower bound must not be larger than the upper bound. /// /// Return [`std::io::ErrorKind::UnexpectedEof`] on Eof. /// /// NOTE that this function does not modify any existing data. /// /// # Cancel safety /// /// It is cancel safe and dropping the returned future will not stop the /// wakeup from happening. #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] pub fn read_to_bytes_rng<'a, R: AsyncRead + ?Sized + Unpin>( reader: &'a mut R, bytes: &'a mut BytesMut, rng: impl std::ops::RangeBounds, ) -> ReadToBytesRngFuture<'a, R> { ReadToBytesRngFuture(read_to_container_rng(reader, bytes, rng)) } /// Return future of [`read_to_bytes_until_end`] /// /// Return number of bytes read in on `Poll::Ready`. #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] pub struct ReadToBytesUntilEndFuture<'a, Reader: ?Sized> { reader: &'a mut Reader, bytes: &'a mut BytesMut, cnt: usize, } impl ReadToBytesUntilEndFuture<'_, Reader> { /// Read into the bytes and adjust the internal length of it. /// Return 0 on EOF. fn poll_read_to_end(&mut self, cx: &mut Context<'_>) -> Poll> { // This uses an adaptive system to extend the vector when it fills. We want to // avoid paying to allocate and zero a huge chunk of memory if the reader only // has 4 bytes while still making large reads if the reader does have a ton // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every // time is 4,500 times (!) slower than this if the reader has a very small // amount of data to return. self.bytes.reserve(32); // Safety: // // The slice returned is not read from and only uninitialized bytes // would be written to it. let mut read_buf = ReadBuf::uninit(unsafe { self.bytes.spare_mut() }); ready!(Pin::new(&mut *self.reader).poll_read(cx, &mut read_buf))?; let filled = read_buf.filled().len(); // safety: // // `read_buf.filled().len()` return number of bytes read in. unsafe { self.bytes.advance(filled) }; Poll::Ready(Ok(filled)) } } impl Future for ReadToBytesUntilEndFuture<'_, Reader> { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = Pin::into_inner(self); loop { let n = ready!(this.poll_read_to_end(cx))?; if n == 0 { break Poll::Ready(Ok(this.cnt)); } this.cnt += n; } } } /// NOTE that this function does not modify any existing data. /// /// # Cancel safety /// /// It is cancel safe and dropping the returned future will not stop the /// wakeup from happening. #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] pub fn read_to_bytes_until_end<'a, R: AsyncRead + ?Sized + Unpin>( reader: &'a mut R, bytes: &'a mut BytesMut, ) -> ReadToBytesUntilEndFuture<'a, R> { ReadToBytesUntilEndFuture { reader, bytes, cnt: 0, } } #[cfg(test)] mod tests { use super::*; use tokio::io::AsyncWriteExt; #[test] fn test_read_exact_to_bytes() { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let (mut r, mut w) = tokio_pipe::pipe().unwrap(); let w_task = tokio::spawn(async move { for n in 1..=255 { w.write_u8(n).await.unwrap(); } }); let r_task = tokio::spawn(async move { let mut buffer = bytes::BytesMut::new(); buffer.put_u8(0); read_exact_to_bytes(&mut r, &mut buffer, 255).await.unwrap(); for (i, each) in buffer.iter().enumerate() { assert_eq!(*each as usize, i); } }); r_task.await.unwrap(); w_task.await.unwrap(); }); } #[test] fn test_read_to_bytes_until_end() { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let (mut r, mut w) = tokio_pipe::pipe().unwrap(); let w_task = tokio::spawn(async move { for n in 1..=255 { w.write_u8(n).await.unwrap(); } }); let r_task = tokio::spawn(async move { let mut buffer = bytes::BytesMut::new(); buffer.put_u8(0); assert_eq!( read_to_bytes_until_end(&mut r, &mut buffer).await.unwrap(), 255 ); for (i, each) in buffer.iter().enumerate() { assert_eq!(*each as usize, i); } }); r_task.await.unwrap(); w_task.await.unwrap(); }); } } tokio-io-utility-0.7.6/src/async_read_utility/inner.rs000064400000000000000000000133310072674642500212650ustar 00000000000000use std::{ cmp, future::Future, io::{ErrorKind, Result}, marker::Unpin, mem::MaybeUninit, ops::Bound::*, pin::Pin, task::{Context, Poll}, }; use tokio::io::{AsyncRead, ReadBuf}; pub trait Container { /// Reserve at least `n` bytes that can be used in /// [`Container::spare_mut`]. fn reserve(&mut self, n: usize); /// Number of initialized bytes. fn len(&self) -> usize; /// If there is no initialized bytes in the container, return `true`. fn is_empty(&self) -> bool { self.len() == 0 } /// Return the capacity reserved. fn capacity(&self) -> usize; /// The returned uninit slice must not be empty. /// /// NOTE that the returned uninit slice might be smaller /// than bytes reserved in [`Container::reserve`] or /// ([`Container::capacity`] - [`Container::len`]). /// /// This is because that the container might be a ring buffer. /// If you consume all uninit slices, then the sum of their lengths /// must be equal to the spare capacity ([`Container::capacity`] - /// [`Container::len`]). /// /// # Safety /// /// The slice returned must not be read from and users should /// never write uninitialized bytes to it. unsafe fn spare_mut(&mut self) -> &mut [MaybeUninit]; /// # Safety /// /// The users must have actually initialized at least `n` bytes /// in the uninit slice returned by [`Container::spare_mut`]. unsafe fn advance(&mut self, n: usize); } impl Container for &mut T { fn reserve(&mut self, n: usize) { (**self).reserve(n) } fn len(&self) -> usize { (**self).len() } fn capacity(&self) -> usize { (**self).capacity() } unsafe fn spare_mut(&mut self) -> &mut [MaybeUninit] { (**self).spare_mut() } unsafe fn advance(&mut self, n: usize) { (**self).advance(n) } } impl Container for Vec { fn reserve(&mut self, n: usize) { Vec::reserve_exact(self, n) } fn len(&self) -> usize { Vec::len(self) } fn capacity(&self) -> usize { Vec::capacity(self) } unsafe fn spare_mut(&mut self) -> &mut [MaybeUninit] { self.spare_capacity_mut() } unsafe fn advance(&mut self, n: usize) { let len = self.len(); self.set_len(len + n) } } #[derive(Debug)] pub struct ReadToContainerRngFuture<'a, C: ?Sized, Reader: ?Sized> { reader: &'a mut Reader, container: &'a mut C, min: usize, max: usize, } impl Future for ReadToContainerRngFuture<'_, C, Reader> where C: Container + ?Sized, Reader: AsyncRead + ?Sized + Unpin, { type Output = Result<()>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = &mut *self; let reader = &mut *this.reader; let container = &mut *this.container; let min = &mut this.min; let max = &mut this.max; if *max == 0 { return Poll::Ready(Ok(())); } // Do not test *min here so that if: // // ```rust // read_to_container_rng(r, c, 0..10).await // ``` // // is called, then we would at least try toread in some bytes. loop { // safety: // // We will never read from it and never write uninitialized bytes // to it. let uninit_slice = unsafe { container.spare_mut() }; let len = cmp::min(uninit_slice.len(), *max); let uninit_slice = &mut uninit_slice[..len]; debug_assert_ne!(uninit_slice.len(), 0); let mut read_buf = ReadBuf::uninit(uninit_slice); ready!(Pin::new(&mut *reader).poll_read(cx, &mut read_buf))?; let filled = read_buf.filled().len(); if filled == 0 { return Poll::Ready(Err(ErrorKind::UnexpectedEof.into())); } // safety: // // `read_buf.filled().len()` return number of bytes read in. unsafe { container.advance(filled) }; *min = min.saturating_sub(filled); *max -= filled; if *min == 0 { break; } } Poll::Ready(Ok(())) } } /// * `rng` - The start of the range specify the minimum of bytes to read in, /// while the end of the range specify the maximum of bytes that /// can be read in. /// If the lower bound is not specified, it is default to 0. /// If the upper bound is not specified, it is default to the /// capacity of `bytes`. /// The lower bound must not be larger than the upper bound. /// /// Return [`ErrorKind::UnexpectedEof`] on Eof. /// /// NOTE that this function does not modify any existing data. /// /// # Cancel safety /// /// It is cancel safe and dropping the returned future will not stop the /// wakeup from happening. pub fn read_to_container_rng<'a, C, Reader>( reader: &'a mut Reader, container: &'a mut C, rng: impl std::ops::RangeBounds, ) -> ReadToContainerRngFuture<'a, C, Reader> where C: Container + ?Sized, Reader: AsyncRead + ?Sized + Unpin, { let min = match rng.start_bound().cloned() { Included(val) => val, Excluded(val) => val + 1, Unbounded => 0, }; let max = match rng.end_bound().cloned() { Included(val) => val, Excluded(val) => val - 1, Unbounded => min.max(container.capacity() - container.len()), }; container.reserve(max); assert!(min <= max, "min {min} should be no larger than max {max}"); ReadToContainerRngFuture { reader, container, min, max, } } tokio-io-utility-0.7.6/src/async_read_utility.rs000064400000000000000000000133020072674642500201500ustar 00000000000000use std::{ future::Future, io::Result, marker::Unpin, pin::Pin, task::{Context, Poll}, }; use tokio::io::AsyncRead; mod inner; pub use inner::*; #[cfg(feature = "bytes")] mod bytes_impl; #[cfg(feature = "bytes")] pub use bytes_impl::*; /// Returned future of [`read_to_vec`]. #[derive(Debug)] pub struct ReadToVecFuture<'a, T: ?Sized>(ReadToVecRngFuture<'a, T>); impl Future for ReadToVecFuture<'_, T> { type Output = Result<()>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::new(&mut self.0).poll(cx) } } /// Try to fill data from `reader` into the spare capacity of `vec`. /// /// It can be used to implement buffering. /// /// Return [`std::io::ErrorKind::UnexpectedEof`] on Eof. /// /// NOTE that this function does not modify any existing data. /// /// # Cancel safety /// /// It is cancel safe and dropping the returned future will not stop the /// wakeup from happening. pub fn read_to_vec<'a, T: AsyncRead + ?Sized + Unpin>( reader: &'a mut T, vec: &'a mut Vec, ) -> ReadToVecFuture<'a, T> { ReadToVecFuture(read_to_vec_rng(reader, vec, ..)) } /// Returned future of [`read_exact_to_vec`]. #[derive(Debug)] pub struct ReadExactToVecFuture<'a, T: ?Sized>(ReadToVecRngFuture<'a, T>); impl Future for ReadExactToVecFuture<'_, T> { type Output = Result<()>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::new(&mut self.0).poll(cx) } } /// * `nread` - bytes to read in /// /// Return [`std::io::ErrorKind::UnexpectedEof`] on Eof. /// /// NOTE that this function does not modify any existing data. /// /// # Cancel safety /// /// It is cancel safe and dropping the returned future will not stop the /// wakeup from happening. pub fn read_exact_to_vec<'a, T: AsyncRead + ?Sized + Unpin>( reader: &'a mut T, vec: &'a mut Vec, nread: usize, ) -> ReadExactToVecFuture<'a, T> { ReadExactToVecFuture(read_to_vec_rng(reader, vec, nread..=nread)) } /// Returned future of [`read_to_vec_rng`]. #[derive(Debug)] pub struct ReadToVecRngFuture<'a, Reader: ?Sized>(ReadToContainerRngFuture<'a, Vec, Reader>); impl Future for ReadToVecRngFuture<'_, Reader> { type Output = Result<()>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::new(&mut self.0).poll(cx) } } /// * `rng` - The start of the range specify the minimum of bytes to read in, /// while the end of the range specify the maximum of bytes that /// can be read in. /// If the lower bound is not specified, it is default to 0. /// If the upper bound is not specified, it is default to the /// capacity of `bytes`. /// The lower bound must not be larger than the upper bound. /// /// Return [`std::io::ErrorKind::UnexpectedEof`] on Eof. /// /// NOTE that this function does not modify any existing data. /// /// # Cancel safety /// /// It is cancel safe and dropping the returned future will not stop the /// wakeup from happening. pub fn read_to_vec_rng<'a, T: AsyncRead + ?Sized + Unpin>( reader: &'a mut T, vec: &'a mut Vec, rng: impl std::ops::RangeBounds, ) -> ReadToVecRngFuture<'a, T> { ReadToVecRngFuture(read_to_container_rng(reader, vec, rng)) } #[cfg(test)] mod tests { use super::*; use tokio::io::AsyncWriteExt; #[test] fn test_read_to_vec() { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let (mut r, mut w) = tokio_pipe::pipe().unwrap(); for n in 1..=255 { w.write_u8(n).await.unwrap(); } let mut buffer = Vec::with_capacity(255); read_to_vec(&mut r, &mut buffer).await.unwrap(); assert_eq!(buffer.len(), 255); for (i, each) in buffer.iter().enumerate() { assert_eq!(*each as usize, i + 1); } }); } #[test] fn test_read_exact_to_vec() { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let (mut r, mut w) = tokio_pipe::pipe().unwrap(); let w_task = tokio::spawn(async move { for n in 1..=255 { w.write_u8(n).await.unwrap(); } }); let r_task = tokio::spawn(async move { let mut buffer = vec![0]; read_exact_to_vec(&mut r, &mut buffer, 255).await.unwrap(); for (i, each) in buffer.iter().enumerate() { assert_eq!(*each as usize, i); } }); r_task.await.unwrap(); w_task.await.unwrap(); }); } #[test] fn test_read_to_vec_rng() { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let (mut r, mut w) = tokio_pipe::pipe().unwrap(); for n in 1..=255 { w.write_u8(n).await.unwrap(); } drop(w); let mut buffer = vec![0]; read_to_vec_rng(&mut r, &mut buffer, 1..255).await.unwrap(); assert_eq!(buffer.len(), 255); for (i, each) in buffer.iter().enumerate() { assert_eq!(*each as usize, i); } }); } } tokio-io-utility-0.7.6/src/async_write_utility.rs000064400000000000000000000057750072674642500204060ustar 00000000000000use super::IoSliceExt; use std::io::{self, IoSlice, Result}; use tokio::io::{AsyncWrite, AsyncWriteExt}; /// Return true if the `bufs` contains at least one byte. pub async fn write_vectored_all( writer: &mut Writer, mut bufs: &mut [IoSlice<'_>], ) -> Result<()> { if bufs.is_empty() { return Ok(()); } while bufs[0].is_empty() { bufs = &mut bufs[1..]; if bufs.is_empty() { return Ok(()); } } // Loop Invariant: // - bufs must not be empty; // - bufs contain at least one byte. loop { // bytes must be greater than 0 since bufs contain // at least one byte. let mut bytes = writer.write_vectored(bufs).await?; if bytes == 0 { return Err(io::ErrorKind::WriteZero.into()); } // This loop would also skip all `IoSlice` that is empty // until the first non-empty `IoSlice` is met. while bufs[0].len() <= bytes { bytes -= bufs[0].len(); bufs = &mut bufs[1..]; if bufs.is_empty() { debug_assert_eq!(bytes, 0); return Ok(()); } } bufs[0] = IoSlice::new(&bufs[0].into_inner()[bytes..]); } } #[cfg(test)] mod tests { use super::write_vectored_all; use std::io::IoSlice; use std::slice::from_raw_parts; use tokio::io::AsyncReadExt; fn as_ioslice(v: &[T]) -> IoSlice<'_> { IoSlice::new(unsafe { from_raw_parts(v.as_ptr() as *const u8, v.len() * std::mem::size_of::()) }) } #[test] fn test() { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let (mut r, mut w) = tokio_pipe::pipe().unwrap(); let w_task = tokio::spawn(async move { let buffer0: Vec = (0..1024).collect(); let buffer1: Vec = (1024..2048).collect(); write_vectored_all(&mut w, &mut [as_ioslice(&buffer0), as_ioslice(&buffer1)]) .await .unwrap(); write_vectored_all(&mut w, &mut [as_ioslice(&buffer0), as_ioslice(&buffer1)]) .await .unwrap(); }); let r_task = tokio::spawn(async move { for _ in 0..2 { let mut n = 0u32; let mut buf = [0; 4 * 128]; while n < 2048 { r.read_exact(&mut buf).await.unwrap(); for x in buf.chunks(4) { assert_eq!(x, n.to_ne_bytes()); n += 1; } } } }); r_task.await.unwrap(); w_task.await.unwrap(); }); } } tokio-io-utility-0.7.6/src/init_maybeuninit_io_slice.rs000064400000000000000000000036460072674642500215040ustar 00000000000000use std::{ io::IoSlice, iter::{IntoIterator, Iterator}, mem::MaybeUninit, }; pub fn init_maybeuninit_io_slices_mut<'a, 'io_slice, Iterable>( uninit_io_slices: &'a mut [MaybeUninit>], iterable: Iterable, ) -> &'a mut [IoSlice<'io_slice>] where Iterable: IntoIterator>, { fn inner<'a, 'io_slice>( uninit_io_slices: &'a mut [MaybeUninit>], iterable: &mut dyn Iterator>, ) -> &'a mut [IoSlice<'io_slice>] { let mut cnt = 0; uninit_io_slices .iter_mut() .zip(iterable) .for_each(|(uninit_io_slice, io_slice)| { uninit_io_slice.write(io_slice); cnt += 1; }); // Safety: // // - uninit_io_slices[..cnt] is initialized using iterable. // - MaybeUninit is a transparent type unsafe { &mut *((&mut uninit_io_slices[..cnt]) as *mut _ as *mut [IoSlice<'io_slice>]) } } inner(uninit_io_slices, &mut iterable.into_iter()) } #[cfg(test)] mod tests { use super::*; use crate::IoSliceExt; #[test] fn test() { let mut uninit_io_slices = [MaybeUninit::>::uninit(); 5]; let io_slices = [ IoSlice::new(b"1023x"), IoSlice::new(b"1qwe"), IoSlice::new(b"''weqdq"), IoSlice::new(b"jiasodjx"), IoSlice::new(b"aqw34f"), ]; assert_io_slices_eq( init_maybeuninit_io_slices_mut(&mut uninit_io_slices, io_slices), &io_slices, ); } fn assert_io_slices_eq(x: &[IoSlice<'_>], y: &[IoSlice<'_>]) { assert_eq!(x.len(), y.len()); let x: Vec<&[u8]> = x.iter().copied().map(IoSliceExt::into_inner).collect(); let y: Vec<&[u8]> = y.iter().copied().map(IoSliceExt::into_inner).collect(); assert_eq!(x, y); } } tokio-io-utility-0.7.6/src/io_slice_ext.rs000064400000000000000000000020500072674642500167210ustar 00000000000000use std::{ io::{IoSlice, IoSliceMut}, slice::{from_raw_parts, from_raw_parts_mut}, }; pub trait IoSliceExt<'a> { fn into_inner(self) -> &'a [u8]; } impl<'a> IoSliceExt<'a> for IoSlice<'a> { fn into_inner(self) -> &'a [u8] { let slice = &*self; unsafe { from_raw_parts(slice.as_ptr(), slice.len()) } } } pub trait IoSliceMutExt<'a> { fn into_inner(self) -> &'a mut [u8]; } impl<'a> IoSliceMutExt<'a> for IoSliceMut<'a> { fn into_inner(mut self) -> &'a mut [u8] { let slice = &mut *self; unsafe { from_raw_parts_mut(slice.as_mut_ptr(), slice.len()) } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_io_slice_ext() { let buffer = b"123455"; let inner: &'static [u8] = IoSlice::new(buffer).into_inner(); assert_eq!(inner, buffer); } #[test] fn test_io_slice_mut_ext() { let mut buffer = b"123455".to_vec(); let inner: &mut [u8] = IoSliceMut::new(&mut *buffer).into_inner(); assert_eq!(inner, b"123455"); } } tokio-io-utility-0.7.6/src/lib.rs000064400000000000000000000016140072674642500150260ustar 00000000000000// only enables the nightly `doc_cfg` feature when // the `docsrs` configuration attribute is defined #![cfg_attr(docsrs, feature(doc_cfg))] /// Replacement for [`std::task::ready`]. #[macro_export] macro_rules! ready { ($e:expr) => { match $e { Poll::Ready(t) => t, Poll::Pending => return Poll::Pending, } }; } pub fn assert_send(val: T) -> T where T: Send, { val } mod async_read_utility; pub use async_read_utility::*; mod async_write_utility; pub use async_write_utility::write_vectored_all; mod init_maybeuninit_io_slice; pub use init_maybeuninit_io_slice::init_maybeuninit_io_slices_mut; mod io_slice_ext; pub use io_slice_ext::{IoSliceExt, IoSliceMutExt}; mod reusable_io_slices; pub use reusable_io_slices::ReusableIoSlices; #[cfg(feature = "bytes")] mod write_bytes; #[cfg(feature = "bytes")] pub use write_bytes::write_all_bytes; tokio-io-utility-0.7.6/src/reusable_io_slices.rs000064400000000000000000000054210072674642500201130ustar 00000000000000use std::{ io::IoSlice, mem::{ManuallyDrop, MaybeUninit}, num::NonZeroUsize, ptr::NonNull, slice::{from_raw_parts, from_raw_parts_mut}, }; /// [`Box`]ed [`IoSlice`] that can be reused for different io_slices /// with different lifetime. #[derive(Debug)] pub struct ReusableIoSlices { ptr: NonNull<()>, cap: NonZeroUsize, } unsafe impl Send for ReusableIoSlices {} unsafe impl Sync for ReusableIoSlices {} impl ReusableIoSlices { /// Create new [`ReusableIoSlices`]. pub fn new(cap: NonZeroUsize) -> Self { let mut v = ManuallyDrop::new(Vec::>>::with_capacity(cap.get())); debug_assert_eq!(v.capacity(), cap.get()); debug_assert_eq!(v.len(), 0); let ptr = v.as_mut_ptr(); // Safety: // // - ptr is allocated using Vec::with_capacity, with non-zero cap // - Vec::as_mut_ptr returns a non-null, valid pointer when // the capacity is non-zero. // - It is valid after the vec is dropped since it is wrapped // in ManuallyDrop. let ptr = unsafe { NonNull::new_unchecked(ptr as *mut ()) }; Self { ptr, cap } } /// Return the underlying io_slices pub fn get_mut(&mut self) -> &mut [MaybeUninit>] { unsafe { from_raw_parts_mut( self.ptr.as_ptr() as *mut MaybeUninit>, self.cap.get(), ) } } /// Return the underlying io_slices pub fn get(&self) -> &[MaybeUninit>] { unsafe { from_raw_parts( self.ptr.as_ptr() as *const MaybeUninit>, self.cap.get(), ) } } } impl Drop for ReusableIoSlices { fn drop(&mut self) { // Cast it back to its original type let ptr = self.ptr.as_ptr() as *mut MaybeUninit>; // Safety: // // - ptr are obtained using `Vec::as_mut_ptr` from // a `ManuallyDrop>` with non-zero cap. // - `IoSlice` does not have a `Drop` implementation and is // `Copy`able, so it is ok to pass `0` as length. // - self.cap.get() == v.capacity() let v: Vec>> = unsafe { Vec::from_raw_parts(ptr, 0, self.cap.get()) }; drop(v); } } #[cfg(test)] mod tests { use super::*; #[test] fn test_reusable_io_slices() { let io_slice = IoSlice::new(b"123exr3x"); for size in 1..300 { let cap = NonZeroUsize::new(size).unwrap(); let mut reusable_io_slices = ReusableIoSlices::new(cap); for uninit_io_slice in reusable_io_slices.get_mut() { uninit_io_slice.write(io_slice); } } } } tokio-io-utility-0.7.6/src/write_bytes.rs000064400000000000000000000160760072674642500166300ustar 00000000000000use std::{io, mem, pin::Pin, vec::IntoIter}; use bytes::Bytes; use tokio::io::{AsyncWrite, AsyncWriteExt}; use crate::{init_maybeuninit_io_slices_mut, ReusableIoSlices}; /// * `buffer` - must not contain empty `Bytes`s. #[cfg_attr(docsrs, doc(cfg(feature = "bytes")))] pub async fn write_all_bytes( writer: Pin<&mut (dyn AsyncWrite + Send)>, buffer: &mut Vec, reusable_io_slices: &mut ReusableIoSlices, ) -> io::Result<()> { // `buffer` does not contain any empty `Bytes`s, so: // - We can check for `io::ErrorKind::WriteZero` error easily // - It won't occupy slots in `reusable_io_slices` so that // we can group as many non-zero IoSlice in one write. // - Avoid conserion from/to `VecDeque` unless necessary, // which might allocate. // - Simplify the loop in write_all_bytes_inner. if buffer.is_empty() { return Ok(()); } // This is O(1) let mut iter = mem::take(buffer).into_iter(); let res = write_all_bytes_inner(writer, &mut iter, reusable_io_slices).await; // This is O(1) because of the specailization in std *buffer = Vec::from_iter(iter); res } /// * `buffer` - contains at least one element and must not contain empty /// `Bytes` async fn write_all_bytes_inner( mut writer: Pin<&mut (dyn AsyncWrite + Send)>, iter: &mut IntoIter, reusable_io_slices: &mut ReusableIoSlices, ) -> io::Result<()> { // do-while style loop, because on the first iteration // iter must not be empty 'outer: loop { let uninit_io_slices = reusable_io_slices.get_mut(); // iter must not be empty // io_slices.is_empty() == false since uninit_io_slices also must not // be empty let io_slices = init_maybeuninit_io_slices_mut( uninit_io_slices, // Do not consume the iter yet since write_vectored might // do partial write. iter.as_slice().iter().map(|bytes| io::IoSlice::new(bytes)), ); debug_assert!(!io_slices.is_empty()); let mut n = writer.write_vectored(io_slices).await?; if n == 0 { // Since io_slices is not empty and it does not contain empty // `Bytes`, it must be WriteZero error. return Err(io::Error::from(io::ErrorKind::WriteZero)); } // On first iteration, iter cannot be empty while n >= iter.as_slice()[0].len() { n -= iter.as_slice()[0].len(); // Release `Bytes` so that the memory they occupied // can be reused in `BytesMut`. iter.next().unwrap(); if iter.as_slice().is_empty() { debug_assert_eq!(n, 0); break 'outer; } } if n != 0 { // iter must not be empty let first = &mut iter.as_mut_slice()[0]; // n < buffer[start].len() *first = first.slice(n..); } } Ok(()) } #[cfg(test)] mod tests { use std::{ iter, num::NonZeroUsize, task::{Context, Poll}, }; use bytes::BytesMut; use tokio::io::AsyncWrite; use super::*; use crate::IoSliceExt; /// Limit number of bytes that can be sent for each write. struct WriterRateLimit(usize, Vec); impl AsyncWrite for WriterRateLimit { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { let n = buf.len().min(self.0); let buf = &buf[..n]; Pin::new(&mut self.1).poll_write(cx, buf) } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.1).poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.1).poll_shutdown(cx) } fn poll_write_vectored( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { let mut cnt = 0; let n = self.0; bufs.iter() .copied() .filter_map(|io_slice| { (n > cnt).then(|| { let n = io_slice.len().min(n - cnt); cnt += n; &io_slice.into_inner()[..n] }) }) .for_each(|slice| { self.1.extend(slice); }); Poll::Ready(Ok(cnt)) } fn is_write_vectored(&self) -> bool { true } } #[test] fn test() { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .block_on(async { let bytes: BytesMut = (0..255).collect(); let bytes = bytes.freeze(); let iter = iter::once(bytes).cycle().take(20); let expected_bytes: Vec = iter.clone().flatten().collect(); let mut reusable_io_slices = ReusableIoSlices::new(NonZeroUsize::new(3).unwrap()); // Emulate a pipe where each time only half of the Bytes can be // written. let writer = WriterRateLimit(255 / 2, Vec::new()); tokio::pin!(writer); write_all_bytes( writer.as_mut(), &mut iter.clone().collect(), &mut reusable_io_slices, ) .await .unwrap(); assert_eq!(writer.1, expected_bytes); // Emulate a pipe where each time exactly one Bytes can be // written. writer.0 = 255; writer.1.clear(); write_all_bytes( writer.as_mut(), &mut iter.clone().collect(), &mut reusable_io_slices, ) .await .unwrap(); assert_eq!(writer.1, expected_bytes); // Emulate a pipe where each time one and a half Bytes can be // written. writer.0 = 255 + 255 / 2; writer.1.clear(); write_all_bytes( writer.as_mut(), &mut iter.clone().collect(), &mut reusable_io_slices, ) .await .unwrap(); assert_eq!(writer.1, expected_bytes); // Emulate a pipe where each time one Bytes and a little bit // of the next Bytes can be written. writer.0 = 255 + 5; writer.1.clear(); write_all_bytes( writer.as_mut(), &mut iter.clone().collect(), &mut reusable_io_slices, ) .await .unwrap(); assert_eq!(writer.1, expected_bytes); }); } }