Fixed issues with #61

zonyitoo · Nov 1, 2016 · 93fef09 · 93fef09
1 parent 55a20e0
commit 93fef09
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 16 deletions.
diff --git a/src/net/tcp.rs b/src/net/tcp.rs
@@ -20,7 +20,8 @@ use std::os::unix::io::{FromRawFd, RawFd};
 use mio::Ready;
 use mio::tcp::{TcpListener as MioTcpListener, TcpStream as MioTcpStream};
 
-use super::{each_addr, GenericEvented};
+use scheduler::ReadyType;
+use super::{each_addr, make_timeout, GenericEvented, SyncGuard};
 
 macro_rules! create_tcp_listener {
     ($inner:expr) => (TcpListener::new($inner, Ready::readable()));
@@ -41,14 +42,33 @@ impl TcpListener {
     }
 
     pub fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> {
-        match self.get_inner().accept() {
-            Ok((stream, addr)) => {
-                trace!("TcpListener({:?}): accept() => Ok(..)", self.token);
-                return create_tcp_stream!(stream).map(|stream| (stream, addr));
+        let mut sync_guard = SyncGuard::new();
+
+        loop {
+            match self.get_inner().accept() {
+                Ok((stream, addr)) => {
+                    trace!("TcpListener({:?}): accept() => Ok(..)", self.token);
+                    return create_tcp_stream!(stream).map(|stream| (stream, addr));
+                }
+                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {
+                    trace!("TcpListener({:?}): accept() => WouldBlock", self.token);
+                }
+                Err(err) => {
+                    trace!("TcpListener({:?}): accept() => Err(..)", self.token);
+                    return Err(err);
+                }
             }
-            Err(err) => {
-                trace!("TcpListener({:?}): accept() => Err(..)", self.token);
-                return Err(err);
+
+            trace!("TcpListener({:?}): wait(Readable)", self.token);
+            sync_guard.disarm();
+
+            match *self.read_timeout.lock() {
+                None => self.ready_states.wait(ReadyType::Readable),
+                Some(t) => {
+                    if self.ready_states.wait_timeout(ReadyType::Readable, t) {
+                        return Err(make_timeout());
+                    }
+                }
             }
         }
     }

diff --git a/src/net/unix.rs b/src/net/unix.rs
@@ -19,7 +19,8 @@ use mio::deprecated::unix::UnixListener as MioUnixListener;
 use mio::deprecated::unix::UnixSocket as MioUnixSocket;
 use mio::deprecated::unix::UnixStream as MioUnixStream;
 
-use super::GenericEvented;
+use scheduler::ReadyType;
+use super::{GenericEvented, SyncGuard};
 
 macro_rules! create_unix_listener {
     ($inner:expr) => (UnixListener::new($inner, Ready::readable()));
@@ -86,14 +87,33 @@ impl UnixListener {
     }
 
     pub fn accept(&self) -> io::Result<UnixStream> {
-        match self.get_inner().accept() {
-            Ok(stream) => {
-                trace!("UnixListener({:?}): accept() => Ok(..)", self.token);
-                return create_unix_stream!(stream);
+        let mut sync_guard = SyncGuard::new();
+
+        loop {
+            match self.get_inner().accept() {
+                Ok(stream) => {
+                    trace!("UnixListener({:?}): accept() => Ok(..)", self.token);
+                    return create_unix_stream!(stream);
+                }
+                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {
+                    trace!("UnixListener({:?}): accept() => WouldBlock", self.token);
+                }
+                Err(err) => {
+                    trace!("UnixListener({:?}): accept() => Err(..)", self.token);
+                    return Err(err);
+                }
             }
-            Err(err) => {
-                trace!("UnixListener({:?}): accept() => Err(..)", self.token);
-                return Err(err);
+
+            trace!("UnixListener({:?}): wait(Readable)", self.token);
+            sync_guard.disarm();
+
+            match *self.read_timeout.lock() {
+                None => self.ready_states.wait(ReadyType::Readable),
+                Some(t) => {
+                    if self.ready_states.wait_timeout(ReadyType::Readable, t) {
+                        return Err(make_timeout());
+                    }
+                }
             }
         }
     }