From 8364de60d1b025b622b9c1a720df1ecb841e3c0f Mon Sep 17 00:00:00 2001
From: "Aode (lion)" <asonix@asonix.dog>
Date: Sat, 5 Mar 2022 13:19:36 -0600
Subject: [PATCH] Add poll_accept method

---
 src/io.rs | 48 ++++++++++++++++++++++++++++--------------------
 1 file changed, 28 insertions(+), 20 deletions(-)

diff --git a/src/io.rs b/src/io.rs
index dac1f96..4b9debe 100644
--- a/src/io.rs
+++ b/src/io.rs
@@ -64,7 +64,7 @@ struct Flush<'a, T: AsFd + 'static> {
 }
 
 struct Accept<'a> {
-    io: &'a Arc<TcpListener>,
+    io: &'a mut Async<TcpListener>,
 }
 
 struct Bind {
@@ -90,13 +90,32 @@ impl Async<TcpListener> {
         .map(Async::new)
     }
 
-    pub async fn accept(&mut self) -> Result<(Async<TcpStream>, Option<SocketAddrAny>)> {
-        let (stream, addr) = Accept {
-            io: self.io.ensure_arc(),
-        }
-        .await?;
+    pub fn poll_accept(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(Async<TcpStream>, Option<SocketAddrAny>)>> {
+        let poll = self
+            .io
+            .as_ref()
+            .try_accept()
+            .map(|nonblock| nonblock.map(|(stream, addr)| (Async::new(stream), addr)));
 
-        Ok((Async::new(stream), addr))
+        poll_nonblocking!(poll);
+
+        ReactorRef::with(|mut reactor| {
+            reactor.register(
+                Arc::clone(self.io.ensure_arc()),
+                cx.waker().clone(),
+                Readiness::read() | Readiness::hangup(),
+            );
+        })
+        .unwrap();
+
+        Poll::Pending
+    }
+
+    pub async fn accept(&mut self) -> Result<(Async<TcpStream>, Option<SocketAddrAny>)> {
+        Accept { io: self }.await
     }
 }
 
@@ -468,21 +487,10 @@ where
 }
 
 impl<'a> Future for Accept<'a> {
-    type Output = Result<(TcpStream, Option<SocketAddrAny>)>;
+    type Output = Result<(Async<TcpStream>, Option<SocketAddrAny>)>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        poll_nonblocking!(self.io.try_accept());
-
-        ReactorRef::with(|mut reactor| {
-            reactor.register(
-                Arc::clone(self.io),
-                cx.waker().clone(),
-                Readiness::read() | Readiness::hangup(),
-            );
-        })
-        .unwrap();
-
-        Poll::Pending
+        self.get_mut().io.poll_accept(cx)
     }
 }