nio 0.1.4

Async runtime for Rust
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
153
154
155
156
157
158
159

use crate::{driver::IoWaker, rt::context::LocalContext};
use mio::{Interest, event::Source};
use std::{
    future::{PollFn, poll_fn},
    io::{ErrorKind, Read, Result, Write},
    task::{Context, Poll},
};

#[derive(Debug)]
pub struct AsyncIO<Io: Source> {
    pub(crate) io: Io,
    waker: Box<IoWaker>,
}

impl<Io: Source> std::panic::UnwindSafe for AsyncIO<Io> {}
impl<Io: Source> std::panic::RefUnwindSafe for AsyncIO<Io> {}

impl<Io: Source> AsyncIO<Io> {
    pub fn new(io: Io) -> Result<Self> {
        Self::with_interest(io, Interest::READABLE | Interest::WRITABLE)
    }

    pub fn with_interest(mut io: Io, interests: Interest) -> Result<Self> {
        let waker = IoWaker::new();
        let token = mio::Token(waker.addr());

        LocalContext::with(|ctx| ctx.io_registry.register(&mut io, token, interests))?;

        Ok(Self { io, waker })
    }

    pub fn io_read<'a, F, T>(
        &'a self,
        mut f: F,
    ) -> PollFn<impl FnMut(&mut Context) -> Poll<Result<T>> + use<'a, F, Io, T>>
    where
        F: FnMut(&'a Io) -> Result<T>,
    {
        poll_fn(move |cx| {
            self.waker.reader.register(cx);

            let readiness = self.waker.readiness();
            if readiness.is_readable() {
                match f(&self.io) {
                    Err(err) if err.kind() == ErrorKind::WouldBlock => {
                        self.waker.clear_read(readiness)
                    }
                    res => return Poll::Ready(res),
                }
            }
            Poll::Pending
        })
    }

    pub fn io_write<'a, F, T>(
        &'a self,
        mut f: F,
    ) -> PollFn<impl FnMut(&mut Context) -> Poll<Result<T>> + use<'a, F, Io, T>>
    where
        F: FnMut(&'a Io) -> Result<T>,
    {
        poll_fn(move |cx| {
            self.waker.writer.register(cx);

            let readiness = self.waker.readiness();
            if readiness.is_writable() {
                match f(&self.io) {
                    Err(err) if err.kind() == ErrorKind::WouldBlock => {
                        self.waker.clear_write(readiness)
                    }
                    res => return Poll::Ready(res),
                }
            }
            Poll::Pending
        })
    }

    pub fn io_writable(&self) -> PollFn<impl FnMut(&mut Context) -> Poll<()> + use<'_, Io>> {
        poll_fn(move |cx| {
            self.waker.writer.register(cx);
            if self.waker.readiness().is_writable() {
                return Poll::Ready(());
            }
            Poll::Pending
        })
    }
}

impl<Io: Source> AsyncIO<Io> {
    pub fn poll_read<'a>(&'a self, cx: &mut Context, buf: &mut [u8]) -> Poll<Result<usize>>
    where
        &'a Io: Read,
    {
        self.waker.reader.register(cx);

        let readiness = self.waker.readiness();
        if readiness.is_readable() {
            match Read::read(&mut &self.io, buf) {
                Ok(nbytes) => {
                    // When mio is using the epoll or kqueue selector, reading a partially full
                    // buffer is sufficient to show that the socket buffer has been drained.
                    //
                    // This optimization does not work for level-triggered selectors such as
                    // windows or when poll is used.
                    //
                    // Read more:
                    // https://github.com/tokio-rs/tokio/issues/5866
                    #[cfg(all(not(mio_unsupported_force_poll_poll), not(windows)))]
                    if 0 < nbytes && nbytes < buf.len() {
                        // same as: nbytes in 1..buf.len()
                        self.waker.clear_read(readiness);
                    }
                    return Poll::Ready(Ok(nbytes));
                }
                Err(err) if err.kind() == ErrorKind::WouldBlock => self.waker.clear_read(readiness),
                Err(err) => return Poll::Ready(Err(err)),
            }
        }
        Poll::Pending
    }

    pub fn poll_write<'a>(&'a self, cx: &mut Context, buf: &[u8]) -> Poll<Result<usize>>
    where
        &'a Io: Write,
    {
        self.waker.writer.register(cx);

        let readiness = self.waker.readiness();
        if readiness.is_writable() {
            match Write::write(&mut &self.io, buf) {
                Ok(nbytes) => {
                    // if we write only part of our buffer, this is sufficient on unix to show
                    // that the socket buffer is full.  Unfortunately this assumption
                    // fails for level-triggered selectors (like on Windows or poll even for
                    // UNIX): https://github.com/tokio-rs/tokio/issues/5866
                    #[cfg(all(not(mio_unsupported_force_poll_poll), not(windows)))]
                    if 0 < nbytes && nbytes < buf.len() {
                        self.waker.clear_write(readiness);
                    }
                    return Poll::Ready(Ok(nbytes));
                }
                Err(err) if err.kind() == ErrorKind::WouldBlock => {
                    self.waker.clear_write(readiness)
                }
                Err(err) => return Poll::Ready(Err(err)),
            }
        }
        Poll::Pending
    }
}

impl<Io: Source> Drop for AsyncIO<Io> {
    fn drop(&mut self) {
        LocalContext::with(|ctx| {
            let _ = ctx.io_registry.deregister(&mut self.io);
        });
        self.waker.drop_waker();
    }
}