src/io/polled_fd.rs

   1 use std::io;
   2 use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
   3 use std::task::{Context, Poll};
   4
   5 use mio::event::Evented;
   6 use mio::unix::EventedFd as MioEventedFd;
   7 use mio::Poll as MioPoll;
   8 use mio::{PollOpt, Ready, Token};
   9 use tokio::io::PollEvented;
  10
  11 use crate::tools::Fd;
  12
  13 #[repr(transparent)]
  14 pub struct EventedFd {
  15     fd: Fd,
  16 }
  17
  18 impl EventedFd {
  19     #[inline]
  20     pub fn new(fd: Fd) -> Self {
  21         Self { fd }
  22     }
  23 }
  24
  25 impl AsRawFd for EventedFd {
  26     #[inline]
  27     fn as_raw_fd(&self) -> RawFd {
  28         self.fd.as_raw_fd()
  29     }
  30 }
  31
  32 impl FromRawFd for EventedFd {
  33     #[inline]
  34     unsafe fn from_raw_fd(fd: RawFd) -> Self {
  35         Self::new(Fd::from_raw_fd(fd))
  36     }
  37 }
  38
  39 impl IntoRawFd for EventedFd {
  40     #[inline]
  41     fn into_raw_fd(self) -> RawFd {
  42         self.fd.into_raw_fd()
  43     }
  44 }
  45
  46 impl Evented for EventedFd {
  47     fn register(
  48         &self,
  49         poll: &MioPoll,
  50         token: Token,
  51         interest: Ready,
  52         opts: PollOpt,
  53     ) -> io::Result<()> {
  54         MioEventedFd(self.fd.as_ref()).register(poll, token, interest, opts)
  55     }
  56
  57     fn reregister(
  58         &self,
  59         poll: &MioPoll,
  60         token: Token,
  61         interest: Ready,
  62         opts: PollOpt,
  63     ) -> io::Result<()> {
  64         MioEventedFd(self.fd.as_ref()).reregister(poll, token, interest, opts)
  65     }
  66
  67     fn deregister(&self, poll: &MioPoll) -> io::Result<()> {
  68         MioEventedFd(self.fd.as_ref()).deregister(poll)
  69     }
  70 }
  71
  72 #[repr(transparent)]
  73 pub struct PolledFd {
  74     fd: PollEvented<EventedFd>,
  75 }
  76
  77 impl PolledFd {
  78     pub fn new(fd: Fd) -> tokio::io::Result<Self> {
  79         Ok(Self {
  80             fd: PollEvented::new(EventedFd::new(fd))?,
  81         })
  82     }
  83
  84     pub fn wrap_read<T>(
  85         &self,
  86         cx: &mut Context,
  87         func: impl FnOnce() -> io::Result<T>,
  88     ) -> Poll<io::Result<T>> {
  89         ready!(self.fd.poll_read_ready(cx, mio::Ready::readable()))?;
  90         match func() {
  91             Ok(out) => Poll::Ready(Ok(out)),
  92             Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {
  93                 self.fd.clear_read_ready(cx, mio::Ready::readable())?;
  94                 Poll::Pending
  95             }
  96             Err(err) => Poll::Ready(Err(err)),
  97         }
  98     }
  99
 100     pub fn wrap_write<T>(
 101         &self,
 102         cx: &mut Context,
 103         func: impl FnOnce() -> io::Result<T>,
 104     ) -> Poll<io::Result<T>> {
 105         ready!(self.fd.poll_write_ready(cx))?;
 106         match func() {
 107             Ok(out) => Poll::Ready(Ok(out)),
 108             Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {
 109                 self.fd.clear_write_ready(cx)?;
 110                 Poll::Pending
 111             }
 112             Err(err) => Poll::Ready(Err(err)),
 113         }
 114     }
 115 }
 116
 117 impl AsRawFd for PolledFd {
 118     #[inline]
 119     fn as_raw_fd(&self) -> RawFd {
 120         self.fd.get_ref().as_raw_fd()
 121     }
 122 }
 123
 124 impl IntoRawFd for PolledFd {
 125     #[inline]
 126     fn into_raw_fd(self) -> RawFd {
 127         // for the kind of resource we're managing it should always be possible to extract it from
 128         // its driver
 129         self.fd
 130             .into_inner()
 131             .expect("failed to remove polled file descriptor from reactor")
 132             .into_raw_fd()
 133     }
 134 }