path: root/src/loggable.rs

use std::{fmt::Display, mem, pin::pin, task::Poll};

use futures::ready;
use pin_project::pin_project;
pub use tokio::io::AsyncWrite;

#[pin_project]
#[derive(Debug)]
pub struct Loggable<W> {
    log_buffer: Vec<u8>,
    #[pin]
    writer: W,
}

impl<W> Loggable<W> {
    pub fn new(writer: W) -> Self {
        Self {
            log_buffer: Vec::new(),
            writer,
        }
    }

    pub fn into_inner(self) -> W {
        self.writer
    }

    pub fn take_log(&mut self) -> Vec<u8> {
        let log: Vec<u8> = mem::replace(&mut self.log_buffer, Vec::new());
        log
    }
}

impl<W: AsyncWrite + Unpin + Send> Display for Loggable<W> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let str = str::from_utf8(&self.log_buffer).unwrap_or("buffer to string conversion failed");
        f.write_str(str)
    }
}

impl<W: AsyncWrite + Unpin + Send> AsyncWrite for Loggable<W> {
    fn poll_write(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<Result<usize, std::io::Error>> {
        let this = self.as_mut().project();
        let ready = ready!(this.writer.poll_write(cx, buf));
        self.log_buffer.extend_from_slice(buf);
        Poll::Ready(ready)
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        let this = self.project();
        Poll::Ready(ready!(this.writer.poll_flush(cx)))
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        let this = self.project();
        Poll::Ready(ready!(this.writer.poll_shutdown(cx)))
    }
}