geph5_rt/
bufpool.rs

1//! Buffer-pooled reads: read into a shared thread-local buffer so no per-read
2//! allocation happens until the read actually unblocks. A tokio-io port of the
3//! `async-io-bufpool` crate (which was built on futures-io traits).
4
5use std::cell::RefCell;
6use std::future::Future;
7use std::pin::Pin;
8use std::task::{Context, Poll};
9
10use bytes::Bytes;
11use pin_project_lite::pin_project;
12use tokio::io::{AsyncRead, ReadBuf};
13
14thread_local! {
15    static BUFFER: RefCell<[u8; 65536]> = const { RefCell::new([0u8; 65536]) };
16}
17
18/// Reads from `rdr` into a thread-local buffer, returning the bytes read.
19///
20/// `Ok(None)` indicates EOF.
21pub async fn pooled_read(rdr: impl AsyncRead, limit: usize) -> std::io::Result<Option<Bytes>> {
22    pooled_read_callback(rdr, limit, |b: &[u8]| Bytes::copy_from_slice(b)).await
23}
24
25/// Like [`pooled_read`], but instead of allocating, invokes `resolve` on the
26/// freshly read bytes (still backed by the thread-local buffer).
27///
28/// `Ok(None)` indicates EOF.
29pub async fn pooled_read_callback<T>(
30    rdr: impl AsyncRead,
31    limit: usize,
32    resolve: impl FnMut(&[u8]) -> T,
33) -> std::io::Result<Option<T>> {
34    PooledOnceReader {
35        rdr,
36        resolve,
37        limit,
38    }
39    .await
40}
41
42pin_project! {
43    struct PooledOnceReader<T, F> {
44        #[pin]
45        rdr: T,
46        resolve: F,
47        limit: usize,
48    }
49}
50
51impl<T: AsyncRead, U, F: FnMut(&[u8]) -> U> Future for PooledOnceReader<T, F> {
52    type Output = std::io::Result<Option<U>>;
53
54    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
55        BUFFER.with(|buf| {
56            let mut buf = buf.borrow_mut();
57            let this = self.project();
58            let limit = (*this.limit).min(buf.len());
59            let mut read_buf = ReadBuf::new(&mut buf[..limit]);
60            match this.rdr.poll_read(cx, &mut read_buf) {
61                Poll::Ready(Ok(())) => {
62                    let filled = read_buf.filled();
63                    if filled.is_empty() {
64                        Poll::Ready(Ok(None))
65                    } else {
66                        Poll::Ready(Ok(Some((this.resolve)(filled))))
67                    }
68                }
69                Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
70                Poll::Pending => Poll::Pending,
71            }
72        })
73    }
74}
geph5_rt/bufpool.rs

geph5_rt/
bufpool.rs
