      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
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211

    use crate::io::driver::{Handle, Interest, Registration};

    use mio::event::Source;
    use std::fmt;
    use std::io;
    use std::ops::Deref;

    cfg_io_driver! {
        /// Associates an I/O resource that implements the [`std::io::Read`] and/or
        /// [`std::io::Write`] traits with the reactor that drives it.
        ///
        /// `PollEvented` uses [`Registration`] internally to take a type that
        /// implements [`mio::Evented`] as well as [`std::io::Read`] and or
        /// [`std::io::Write`] and associate it with a reactor that will drive it.
        ///
        /// Once the [`mio::Evented`] type is wrapped by `PollEvented`, it can be
        /// used from within the future's execution model. As such, the
        /// `PollEvented` type provides [`AsyncRead`] and [`AsyncWrite`]
        /// implementations using the underlying I/O resource as well as readiness
        /// events provided by the reactor.
        ///
        /// **Note**: While `PollEvented` is `Sync` (if the underlying I/O type is
        /// `Sync`), the caller must ensure that there are at most two tasks that
        /// use a `PollEvented` instance concurrently. One for reading and one for
        /// writing. While violating this requirement is "safe" from a Rust memory
        /// model point of view, it will result in unexpected behavior in the form
        /// of lost notifications and tasks hanging.
        ///
        /// ## Readiness events
        ///
        /// Besides just providing [`AsyncRead`] and [`AsyncWrite`] implementations,
        /// this type also supports access to the underlying readiness event stream.
        /// While similar in function to what [`Registration`] provides, the
        /// semantics are a bit different.
        ///
        /// Two functions are provided to access the readiness events:
        /// [`poll_read_ready`] and [`poll_write_ready`]. These functions return the
        /// current readiness state of the `PollEvented` instance. If
        /// [`poll_read_ready`] indicates read readiness, immediately calling
        /// [`poll_read_ready`] again will also indicate read readiness.
        ///
        /// When the operation is attempted and is unable to succeed due to the I/O
        /// resource not being ready, the caller must call [`clear_read_ready`] or
        /// [`clear_write_ready`]. This clears the readiness state until a new
        /// readiness event is received.
        ///
        /// This allows the caller to implement additional functions. For example,
        /// [`TcpListener`] implements poll_accept by using [`poll_read_ready`] and
        /// [`clear_read_ready`].
        ///
        /// ## Platform-specific events
        ///
        /// `PollEvented` also allows receiving platform-specific `mio::Ready` events.
        /// These events are included as part of the read readiness event stream. The
        /// write readiness event stream is only for `Ready::writable()` events.
        ///
        /// [`std::io::Read`]: trait@std::io::Read
        /// [`std::io::Write`]: trait@std::io::Write
        /// [`AsyncRead`]: trait@AsyncRead
        /// [`AsyncWrite`]: trait@AsyncWrite
        /// [`mio::Evented`]: trait@mio::Evented
        /// [`Registration`]: struct@Registration
        /// [`TcpListener`]: struct@crate::net::TcpListener
        /// [`clear_read_ready`]: method@Self::clear_read_ready
        /// [`clear_write_ready`]: method@Self::clear_write_ready
        /// [`poll_read_ready`]: method@Self::poll_read_ready
        /// [`poll_write_ready`]: method@Self::poll_write_ready
        pub(crate) struct PollEvented<E: Source> {
            io: Option<E>,
            registration: Registration,
        }
    }

    // ===== impl PollEvented =====

    impl<E: Source> PollEvented<E> {
        /// Creates a new `PollEvented` associated with the default reactor.
        ///
        /// # Panics
        ///
        /// This function panics if thread-local runtime is not set.
        ///
        /// The runtime is usually set implicitly when this function is called
        /// from a future driven by a tokio runtime, otherwise runtime can be set
        /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function.
        #[cfg_attr(feature = "signal", allow(unused))]
        pub(crate) fn new(io: E) -> io::Result<Self> {
            PollEvented::new_with_interest(io, Interest::READABLE | Interest::WRITABLE)
        }

        /// Creates a new `PollEvented` associated with the default reactor, for
        /// specific `Interest` state. `new_with_interest` should be used over `new`
        /// when you need control over the readiness state, such as when a file
        /// descriptor only allows reads. This does not add `hup` or `error` so if
        /// you are interested in those states, you will need to add them to the
        /// readiness state passed to this function.
        ///
        /// # Panics
        ///
        /// This function panics if thread-local runtime is not set.
        ///
        /// The runtime is usually set implicitly when this function is called from
        /// a future driven by a tokio runtime, otherwise runtime can be set
        /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter)
        /// function.
        #[cfg_attr(feature = "signal", allow(unused))]
        pub(crate) fn new_with_interest(io: E, interest: Interest) -> io::Result<Self> {
            Self::new_with_interest_and_handle(io, interest, Handle::current())
        }

        pub(crate) fn new_with_interest_and_handle(
            mut io: E,
            interest: Interest,
            handle: Handle,
        ) -> io::Result<Self> {
            let registration = Registration::new_with_interest_and_handle(&mut io, interest, handle)?;
            Ok(Self {
                io: Some(io),
                registration,
            })
        }

        /// Returns a reference to the registration
        pub(crate) fn registration(&self) -> &Registration {
            &self.registration
        }

        /// Deregister the inner io from the registration and returns a Result containing the inner io
        #[cfg(feature = "net")]
        pub(crate) fn into_inner(mut self) -> io::Result<E> {
            let mut inner = self.io.take().unwrap(); // As io shouldn't ever be None, just unwrap here.
            self.registration.deregister(&mut inner)?;
            Ok(inner)
        }
    }

    feature! {
        #![any(feature = "net", feature = "process")]

        use crate::io::ReadBuf;
        use std::task::{Context, Poll};

        impl<E: Source> PollEvented<E> {
            // Safety: The caller must ensure that `E` can read into uninitialized memory
            pub(crate) unsafe fn poll_read<'a>(
                &'a self,
                cx: &mut Context<'_>,
                buf: &mut ReadBuf<'_>,
            ) -> Poll<io::Result<()>>
            where
                &'a E: io::Read + 'a,
            {
                use std::io::Read;

                let n = ready!(self.registration.poll_read_io(cx, || {
                    let b = &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
                    self.io.as_ref().unwrap().read(b)
                }))?;

                // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the
                // buffer.
                buf.assume_init(n);
                buf.advance(n);
                Poll::Ready(Ok(()))
            }

            pub(crate) fn poll_write<'a>(&'a self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>>
            where
                &'a E: io::Write + 'a,
            {
                use std::io::Write;
                self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write(buf))
            }

            #[cfg(feature = "net")]
            pub(crate) fn poll_write_vectored<'a>(
                &'a self,
                cx: &mut Context<'_>,
                bufs: &[io::IoSlice<'_>],
            ) -> Poll<io::Result<usize>>
            where
                &'a E: io::Write + 'a,
            {
                use std::io::Write;
                self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write_vectored(bufs))
            }
        }
    }

    impl<E: Source> Deref for PollEvented<E> {
        type Target = E;

        fn deref(&self) -> &E {
            self.io.as_ref().unwrap()
        }
    }

    impl<E: Source + fmt::Debug> fmt::Debug for PollEvented<E> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("PollEvented").field("io", &self.io).finish()
        }
    }

    impl<E: Source> Drop for PollEvented<E> {
        fn drop(&mut self) {
            if let Some(mut io) = self.io.take() {
                // Ignore errors
                let _ = self.registration.deregister(&mut io);
            }
        }
    }