cookie_factory/
async_bufwriter.rs

1use futures::io::{AsyncWrite, AsyncWriteExt};
2
3pub struct AsyncBufWriter<W: AsyncWrite + Unpin> {
4    inner: W,
5    buf: Vec<u8>,
6    pos: usize,
7}
8
9const DEFAULT_BUF_SIZE: usize = 8 * 1024;
10
11impl<W: AsyncWrite + Unpin> AsyncBufWriter<W> {
12    pub fn new(inner: W) -> AsyncBufWriter<W> {
13        AsyncBufWriter::with_capacity(DEFAULT_BUF_SIZE, inner)
14    }
15
16    pub fn with_capacity(capacity: usize, inner: W) -> AsyncBufWriter<W> {
17        let buf = vec![0; capacity];
18        AsyncBufWriter { inner, buf, pos: 0 }
19    }
20
21    pub async fn flush(&mut self) -> std::io::Result<usize> {
22        let sz = self.inner.write(&self.buf[..self.pos]).await?;
23
24        // if we did not write everything, move data to the beginning of the
25        // buffer
26        if sz < self.pos {
27            for i in 0..(self.pos - sz) {
28                self.buf[i] = self.buf[sz + i]
29            }
30        }
31
32        self.pos -= sz;
33
34        Ok(sz)
35    }
36
37    pub fn remaining(&self) -> usize {
38        self.pos
39    }
40
41    pub fn into_parts(mut self) -> (W, Vec<u8>) {
42        self.buf.truncate(self.pos);
43        (self.inner, self.buf)
44    }
45}
46
47impl<W: AsyncWrite + Unpin> std::io::Write for AsyncBufWriter<W> {
48    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
49        let sz = (&mut self.buf[self.pos..]).write(buf)?;
50        self.pos += sz;
51        Ok(sz)
52    }
53
54    fn flush(&mut self) -> std::io::Result<()> {
55        Ok(())
56    }
57}
58
59pub async fn gen<W: AsyncWrite + Unpin, F: crate::internal::SerializeFn<AsyncBufWriter<W>>>(
60    f: F,
61    w: AsyncBufWriter<W>,
62) -> Result<(AsyncBufWriter<W>, u64), crate::internal::GenError> {
63    match f(crate::internal::WriteContext::from(w)).map(|ctx| ctx.into_inner()) {
64        Err(e) => Err(e),
65        Ok((mut w, _)) => {
66            let sz = w.flush().await?;
67            Ok((w, sz as u64))
68        }
69    }
70}
cookie_factory/async_bufwriter.rs

cookie_factory/
async_bufwriter.rs
