orb 0.12.1

An abstraction for writing runtime agnostic async code. Orb provides interfaces to adapt different async runtimes like smol and tokio.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use super::{AsyncRead, AsyncWrite};
use std::{fmt, io};

/// A buffered reader that wraps an `AsyncRead` trait object and a buffer.
pub struct AsyncBufRead {
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
}

impl AsyncBufRead {
    /// Creates a new `AsyncBufRead` with the given reader and buffer capacity.
    #[inline]
    pub fn new(capacity: usize) -> Self {
        assert!(capacity > 0, "capacity {} must > 0", capacity);
        AsyncBufRead { buf: vec![0; capacity], pos: 0, cap: 0 }
    }

    #[inline]
    pub async fn read_buffered<T: AsyncRead>(
        &mut self, reader: &mut T, buf: &mut [u8],
    ) -> io::Result<usize> {
        // If we have bytes in our buffer, copy them to `buf`.
        if self.pos < self.cap {
            let n = std::cmp::min(buf.len(), self.cap - self.pos);
            buf[..n].copy_from_slice(&self.buf[self.pos..self.pos + n]);
            self.pos += n;
            return Ok(n);
        }

        // If the request is larger than our buffer, read directly into `buf`.
        // This avoids extra copying.
        if buf.len() >= self.buf.len() {
            return reader.read(buf).await;
        }

        // Otherwise, fill our buffer and then copy to `buf`.
        self.cap = reader.read(&mut self.buf).await?;
        self.pos = 0;
        let n = std::cmp::min(buf.len(), self.cap);
        buf[..n].copy_from_slice(&self.buf[..n]);
        self.pos += n;
        Ok(n)
    }
}

/// A buffered writer that wraps an `AsyncWrite` trait object and a buffer.
pub struct AsyncBufWrite {
    buf: Vec<u8>,
    pos: usize,
}

impl AsyncBufWrite {
    /// Creates a new `AsyncBufWrite` with the given writer and buffer capacity.
    #[inline]
    pub fn new(capacity: usize) -> Self {
        assert!(capacity > 0, "capacity {} must > 0", capacity);
        AsyncBufWrite { buf: vec![0; capacity], pos: 0 }
    }

    /// Flushes the buffered data to the underlying writer.
    #[inline]
    pub async fn flush<W: AsyncWrite>(&mut self, writer: &mut W) -> io::Result<()> {
        if self.pos > 0 {
            writer.write_all(&self.buf[..self.pos]).await?;
            self.pos = 0;
        }
        Ok(())
    }

    #[inline]
    pub async fn write_buffered<W: AsyncWrite>(
        &mut self, writer: &mut W, buf: &[u8],
    ) -> io::Result<usize> {
        // If the incoming buffer is larger than our internal buffer's capacity,
        // flush our buffer and write the incoming buffer directly.
        if buf.len() >= self.buf.len() {
            self.flush(writer).await?;
            return writer.write(buf).await;
        }

        // If the incoming buffer doesn't fit in the remaining space in our buffer,
        // flush our buffer.
        if self.buf.len() - self.pos < buf.len() {
            self.flush(writer).await?;
        }
        // Copy the incoming buffer into our internal buffer.
        let n = buf.len();
        self.buf[self.pos..self.pos + n].copy_from_slice(buf);
        self.pos += n;
        Ok(n)
    }
}

pub struct AsyncBufStream<T: AsyncRead + AsyncWrite> {
    read_buf: AsyncBufRead,
    write_buf: AsyncBufWrite,
    inner: T,
}

impl<T: AsyncRead + AsyncWrite> AsyncBufStream<T> {
    #[inline]
    pub fn new(stream: T, buf_size: usize) -> Self {
        Self {
            read_buf: AsyncBufRead::new(buf_size),
            write_buf: AsyncBufWrite::new(buf_size),
            inner: stream,
        }
    }

    #[inline(always)]
    pub async fn flush(&mut self) -> io::Result<()> {
        self.write_buf.flush(&mut self.inner).await
    }

    #[inline(always)]
    pub fn get_inner(&mut self) -> &mut T {
        &mut self.inner
    }
}

impl<T: AsyncRead + AsyncWrite + fmt::Debug> fmt::Debug for AsyncBufStream<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.inner.fmt(f)
    }
}

impl<T: AsyncRead + AsyncWrite + fmt::Display> fmt::Display for AsyncBufStream<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.inner.fmt(f)
    }
}

impl<T: AsyncRead + AsyncWrite> AsyncRead for AsyncBufStream<T> {
    /// Async version of read function
    ///
    /// On ok, return the bytes read
    #[inline(always)]
    async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.read_buf.read_buffered(&mut self.inner, buf).await
    }
}

impl<T: AsyncRead + AsyncWrite> AsyncWrite for AsyncBufStream<T> {
    /// Async version of write function
    ///
    /// On ok, return the bytes written
    #[inline(always)]
    async fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.write_buf.write_buffered(&mut self.inner, buf).await
    }
}