3 files changed, 42 insertions, 80 deletions

diff --git a/tokio/src/net/unix/datagram/socket.rs b/tokio/src/net/unix/datagram/socket.rs
index 6215b579..f9e9321b 100644
--- a/tokio/src/net/unix/datagram/socket.rs
+++ b/tokio/src/net/unix/datagram/socket.rs
@@ -1,4 +1,4 @@
-use crate::io::PollEvented;
+use crate::io::{Interest, PollEvented};
 use crate::net::unix::SocketAddr;
 
 use std::convert::TryFrom;
@@ -270,7 +270,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn connect<P: AsRef<Path>>(&self, path: P) -> io::Result<()> {
-        self.io.get_ref().connect(path)
+        self.io.connect(path)
     }
 
     /// Sends data on the socket to the socket's peer.
@@ -301,7 +301,8 @@ impl UnixDatagram {
     /// ```
     pub async fn send(&self, buf: &[u8]) -> io::Result<usize> {
         self.io
-            .async_io(mio::Interest::WRITABLE, |sock| sock.send(buf))
+            .registration()
+            .async_io(Interest::WRITABLE, || self.io.send(buf))
             .await
     }
 
@@ -330,7 +331,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn try_send(&self, buf: &[u8]) -> io::Result<usize> {
-        self.io.get_ref().send(buf)
+        self.io.send(buf)
     }
 
     /// Try to send a datagram to the peer without waiting.
@@ -369,7 +370,7 @@ impl UnixDatagram {
     where
         P: AsRef<Path>,
     {
-        self.io.get_ref().send_to(buf, target)
+        self.io.send_to(buf, target)
     }
 
     /// Receives data from the socket.
@@ -400,7 +401,8 @@ impl UnixDatagram {
     /// ```
     pub async fn recv(&self, buf: &mut [u8]) -> io::Result<usize> {
         self.io
-            .async_io(mio::Interest::READABLE, |sock| sock.recv(buf))
+            .registration()
+            .async_io(Interest::READABLE, || self.io.recv(buf))
             .await
     }
 
@@ -429,7 +431,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn try_recv(&self, buf: &mut [u8]) -> io::Result<usize> {
-        self.io.get_ref().recv(buf)
+        self.io.recv(buf)
     }
 
     /// Sends data on the socket to the specified address.
@@ -470,9 +472,8 @@ impl UnixDatagram {
         P: AsRef<Path>,
     {
         self.io
-            .async_io(mio::Interest::WRITABLE, |sock| {
-                sock.send_to(buf, target.as_ref())
-            })
+            .registration()
+            .async_io(Interest::WRITABLE, || self.io.send_to(buf, target.as_ref()))
             .await
     }
 
@@ -512,7 +513,8 @@ impl UnixDatagram {
     pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
         let (n, addr) = self
             .io
-            .async_io(mio::Interest::READABLE, |sock| sock.recv_from(buf))
+            .registration()
+            .async_io(Interest::READABLE, || self.io.recv_from(buf))
             .await?;
 
         Ok((n, SocketAddr(addr)))
@@ -551,7 +553,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn try_recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
-        let (n, addr) = self.io.get_ref().recv_from(buf)?;
+        let (n, addr) = self.io.recv_from(buf)?;
         Ok((n, SocketAddr(addr)))
     }
 
@@ -596,7 +598,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn local_addr(&self) -> io::Result<SocketAddr> {
-        self.io.get_ref().local_addr().map(SocketAddr)
+        self.io.local_addr().map(SocketAddr)
     }
 
     /// Returns the address of this socket's peer.
@@ -645,7 +647,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn peer_addr(&self) -> io::Result<SocketAddr> {
-        self.io.get_ref().peer_addr().map(SocketAddr)
+        self.io.peer_addr().map(SocketAddr)
     }
 
     /// Returns the value of the `SO_ERROR` option.
@@ -668,7 +670,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn take_error(&self) -> io::Result<Option<io::Error>> {
-        self.io.get_ref().take_error()
+        self.io.take_error()
     }
 
     /// Shuts down the read, write, or both halves of this connection.
@@ -704,7 +706,7 @@ impl UnixDatagram {
     /// # }
     /// ```
     pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
-        self.io.get_ref().shutdown(how)
+        self.io.shutdown(how)
     }
 }
 
@@ -722,12 +724,12 @@ impl TryFrom<std::os::unix::net::UnixDatagram> for UnixDatagram {
 
 impl fmt::Debug for UnixDatagram {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.io.get_ref().fmt(f)
+        self.io.fmt(f)
     }
 }
 
 impl AsRawFd for UnixDatagram {
     fn as_raw_fd(&self) -> RawFd {
-        self.io.get_ref().as_raw_fd()
+        self.io.as_raw_fd()
     }
 }
diff --git a/tokio/src/net/unix/listener.rs b/tokio/src/net/unix/listener.rs
index 8f0d4c0b..b1da0e3c 100644
--- a/tokio/src/net/unix/listener.rs
+++ b/tokio/src/net/unix/listener.rs
@@ -1,4 +1,4 @@
-use crate::io::PollEvented;
+use crate::io::{Interest, PollEvented};
 use crate::net::unix::{SocketAddr, UnixStream};
 
 use std::convert::TryFrom;
@@ -90,19 +90,20 @@ impl UnixListener {
 
     /// Returns the local socket address of this listener.
     pub fn local_addr(&self) -> io::Result<SocketAddr> {
-        self.io.get_ref().local_addr().map(SocketAddr)
+        self.io.local_addr().map(SocketAddr)
     }
 
     /// Returns the value of the `SO_ERROR` option.
     pub fn take_error(&self) -> io::Result<Option<io::Error>> {
-        self.io.get_ref().take_error()
+        self.io.take_error()
     }
 
     /// Accepts a new incoming connection to this listener.
     pub async fn accept(&self) -> io::Result<(UnixStream, SocketAddr)> {
         let (mio, addr) = self
             .io
-            .async_io(mio::Interest::READABLE, |sock| sock.accept())
+            .registration()
+            .async_io(Interest::READABLE, || self.io.accept())
             .await?;
 
         let addr = SocketAddr(addr);
@@ -119,21 +120,10 @@ impl UnixListener {
     /// The caller is responsible to ensure that `poll_accept` is called from a
     /// single task. Failing to do this could result in tasks hanging.
     pub fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<io::Result<(UnixStream, SocketAddr)>> {
-        loop {
-            let ev = ready!(self.io.poll_read_ready(cx))?;
-
-            match self.io.get_ref().accept() {
-                Ok((sock, addr)) => {
-                    let addr = SocketAddr(addr);
-                    let sock = UnixStream::new(sock)?;
-                    return Poll::Ready(Ok((sock, addr)));
-                }
-                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {
-                    self.io.clear_readiness(ev);
-                }
-                Err(err) => return Err(err).into(),
-            }
-        }
+        let (sock, addr) = ready!(self.io.registration().poll_read_io(cx, || self.io.accept()))?;
+        let addr = SocketAddr(addr);
+        let sock = UnixStream::new(sock)?;
+        Poll::Ready(Ok((sock, addr)))
     }
 }
 
@@ -161,12 +151,12 @@ impl TryFrom<std::os::unix::net::UnixListener> for UnixListener {
 
 impl fmt::Debug for UnixListener {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.io.get_ref().fmt(f)
+        self.io.fmt(f)
     }
 }
 
 impl AsRawFd for UnixListener {
     fn as_raw_fd(&self) -> RawFd {
-        self.io.get_ref().as_raw_fd()
+        self.io.as_raw_fd()
     }
 }
diff --git a/tokio/src/net/unix/stream.rs b/tokio/src/net/unix/stream.rs
index 2f3dd128..1d840926 100644
--- a/tokio/src/net/unix/stream.rs
+++ b/tokio/src/net/unix/stream.rs
@@ -7,7 +7,7 @@ use crate::net::unix::SocketAddr;
 
 use std::convert::TryFrom;
 use std::fmt;
-use std::io::{self, Read, Write};
+use std::io;
 use std::net::Shutdown;
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::os::unix::net;
@@ -39,7 +39,7 @@ impl UnixStream {
         let stream = mio::net::UnixStream::connect(path)?;
         let stream = UnixStream::new(stream)?;
 
-        poll_fn(|cx| stream.io.poll_write_ready(cx)).await?;
+        poll_fn(|cx| stream.io.registration().poll_write_ready(cx)).await?;
         Ok(stream)
     }
 
@@ -84,12 +84,12 @@ impl UnixStream {
 
     /// Returns the socket address of the local half of this connection.
     pub fn local_addr(&self) -> io::Result<SocketAddr> {
-        self.io.get_ref().local_addr().map(SocketAddr)
+        self.io.local_addr().map(SocketAddr)
     }
 
     /// Returns the socket address of the remote half of this connection.
     pub fn peer_addr(&self) -> io::Result<SocketAddr> {
-        self.io.get_ref().peer_addr().map(SocketAddr)
+        self.io.peer_addr().map(SocketAddr)
     }
 
     /// Returns effective credentials of the process which called `connect` or `pair`.
@@ -99,7 +99,7 @@ impl UnixStream {
 
     /// Returns the value of the `SO_ERROR` option.
     pub fn take_error(&self) -> io::Result<Option<io::Error>> {
-        self.io.get_ref().take_error()
+        self.io.take_error()
     }
 
     /// Shuts down the read, write, or both halves of this connection.
@@ -108,7 +108,7 @@ impl UnixStream {
     /// specified portions to immediately return with an appropriate value
     /// (see the documentation of `Shutdown`).
     pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
-        self.io.get_ref().shutdown(how)
+        self.io.shutdown(how)
     }
 
     // These lifetime markers also appear in the generated documentation, and make
@@ -199,29 +199,8 @@ impl UnixStream {
         cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<io::Result<()>> {
-        loop {
-            let ev = ready!(self.io.poll_read_ready(cx))?;
-
-            // Safety: `UnixStream::read` will not peek at the maybe uinitialized bytes.
-            let b = unsafe {
-                &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8])
-            };
-            match self.io.get_ref().read(b) {
-                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                    self.io.clear_readiness(ev);
-                }
-                Ok(n) => {
-                    // Safety: We trust `UnixStream::read` to have filled up `n` bytes
-                    // in the buffer.
-                    unsafe {
-                        buf.assume_init(n);
-                    }
-                    buf.advance(n);
-                    return Poll::Ready(Ok(()));
-                }
-                Err(e) => return Poll::Ready(Err(e)),
-            }
-        }
+        // Safety: `UdpStream::read` correctly handles reads into uninitialized memory
+        unsafe { self.io.poll_read(cx, buf) }
     }
 
     pub(crate) fn poll_write_priv(
@@ -229,27 +208,18 @@ impl UnixStream {
         cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<io::Result<usize>> {
-        loop {
-            let ev = ready!(self.io.poll_write_ready(cx))?;
-
-            match self.io.get_ref().write(buf) {
-                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                    self.io.clear_readiness(ev);
-                }
-                x => return Poll::Ready(x),
-            }
-        }
+        self.io.poll_write(cx, buf)
     }
 }
 
 impl fmt::Debug for UnixStream {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.io.get_ref().fmt(f)
+        self.io.fmt(f)
     }
 }
 
 impl AsRawFd for UnixStream {
     fn as_raw_fd(&self) -> RawFd {
-        self.io.get_ref().as_raw_fd()
+        self.io.as_raw_fd()
     }
 }