Line data    Source code

       1              : use pin_project_lite::pin_project;
       2              : use std::io::Read;
       3              : use std::pin::Pin;
       4              : use std::{io, task};
       5              : use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
       6              : 
       7              : pin_project! {
       8              :     /// This stream tracks all writes and calls user provided
       9              :     /// callback when the underlying stream is flushed.
      10              :     pub struct MeasuredStream<S, R, W> {
      11              :         #[pin]
      12              :         stream: S,
      13              :         write_count: usize,
      14              :         inc_read_count: R,
      15              :         inc_write_count: W,
      16              :     }
      17              : }
      18              : 
      19              : impl<S, R, W> MeasuredStream<S, R, W> {
      20         3783 :     pub fn new(stream: S, inc_read_count: R, inc_write_count: W) -> Self {
      21         3783 :         Self {
      22         3783 :             stream,
      23         3783 :             write_count: 0,
      24         3783 :             inc_read_count,
      25         3783 :             inc_write_count,
      26         3783 :         }
      27         3783 :     }
      28              : }
      29              : 
      30              : impl<S: AsyncRead + Unpin, R: FnMut(usize), W> AsyncRead for MeasuredStream<S, R, W> {
      31     10553722 :     fn poll_read(
      32     10553722 :         self: Pin<&mut Self>,
      33     10553722 :         context: &mut task::Context<'_>,
      34     10553722 :         buf: &mut ReadBuf<'_>,
      35     10553722 :     ) -> task::Poll<io::Result<()>> {
      36     10553722 :         let this = self.project();
      37     10553722 :         let filled = buf.filled().len();
      38     10553722 :         this.stream.poll_read(context, buf).map_ok(|()| {
      39      3228791 :             let cnt = buf.filled().len() - filled;
      40      3228791 :             // Increment the read count.
      41      3228791 :             (this.inc_read_count)(cnt);
      42     10553722 :         })
      43     10553722 :     }
      44              : }
      45              : 
      46              : impl<S: AsyncWrite + Unpin, R, W: FnMut(usize)> AsyncWrite for MeasuredStream<S, R, W> {
      47      2983808 :     fn poll_write(
      48      2983808 :         self: Pin<&mut Self>,
      49      2983808 :         context: &mut task::Context<'_>,
      50      2983808 :         buf: &[u8],
      51      2983808 :     ) -> task::Poll<io::Result<usize>> {
      52      2983808 :         let this = self.project();
      53      2983808 :         this.stream.poll_write(context, buf).map_ok(|cnt| {
      54      2962027 :             // Increment the write count.
      55      2962027 :             *this.write_count += cnt;
      56      2962027 :             cnt
      57      2983808 :         })
      58      2983808 :     }
      59              : 
      60      2967503 :     fn poll_flush(
      61      2967503 :         self: Pin<&mut Self>,
      62      2967503 :         context: &mut task::Context<'_>,
      63      2967503 :     ) -> task::Poll<io::Result<()>> {
      64      2967503 :         let this = self.project();
      65      2967503 :         this.stream.poll_flush(context).map_ok(|()| {
      66      2967503 :             // Call the user provided callback and reset the write count.
      67      2967503 :             (this.inc_write_count)(*this.write_count);
      68      2967503 :             *this.write_count = 0;
      69      2967503 :         })
      70      2967503 :     }
      71              : 
      72         3093 :     fn poll_shutdown(
      73         3093 :         self: Pin<&mut Self>,
      74         3093 :         context: &mut task::Context<'_>,
      75         3093 :     ) -> task::Poll<io::Result<()>> {
      76         3093 :         self.project().stream.poll_shutdown(context)
      77         3093 :     }
      78              : }
      79              : 
      80              : /// Wrapper for a reader that counts bytes read.
      81              : ///
      82              : /// Similar to MeasuredStream but it's one way and it's sync
      83              : pub struct MeasuredReader<R: Read> {
      84              :     inner: R,
      85              :     byte_count: usize,
      86              : }
      87              : 
      88              : impl<R: Read> MeasuredReader<R> {
      89          656 :     pub fn new(reader: R) -> Self {
      90          656 :         Self {
      91          656 :             inner: reader,
      92          656 :             byte_count: 0,
      93          656 :         }
      94          656 :     }
      95              : 
      96          654 :     pub fn get_byte_count(&self) -> usize {
      97          654 :         self.byte_count
      98          654 :     }
      99              : }
     100              : 
     101              : impl<R: Read> Read for MeasuredReader<R> {
     102       943526 :     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
     103       943526 :         let result = self.inner.read(buf);
     104       943526 :         if let Ok(n_bytes) = result {
     105       943524 :             self.byte_count += n_bytes
     106            2 :         }
     107       943526 :         result
     108       943526 :     }
     109              : }