45 lines
817 B
Rust
45 lines
817 B
Rust
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
use tokio::io::{self, AsyncRead, ReadBuf};
|
|
|
|
pub struct CountingReader<R> {
|
|
reader: R,
|
|
count: usize,
|
|
}
|
|
|
|
impl<R> AsyncRead for CountingReader<R>
|
|
where
|
|
R: AsyncRead + Unpin,
|
|
{
|
|
fn poll_read(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &mut ReadBuf<'_>,
|
|
) -> Poll<io::Result<()>> {
|
|
let before = buf.filled().len();
|
|
let poll = Pin::new(&mut self.reader).poll_read(cx, buf);
|
|
let after = buf.filled().len();
|
|
self.count += after - before;
|
|
poll
|
|
}
|
|
}
|
|
|
|
impl<R> CountingReader<R> {
|
|
pub fn new(reader: R) -> Self {
|
|
Self { reader, count: 0 }
|
|
}
|
|
|
|
pub fn count(&self) -> usize {
|
|
self.count
|
|
}
|
|
}
|
|
|
|
impl<R> AsRef<R> for CountingReader<R>
|
|
where
|
|
R: AsyncRead + Unpin,
|
|
{
|
|
fn as_ref(&self) -> &R {
|
|
&self.reader
|
|
}
|
|
}
|