wip

ntex-rs · May 29, 2023 · 8024ded · 8024ded
1 parent 70433e0
commit 8024ded
Showing 1 changed file with 21 additions and 13 deletions.
diff --git a/ntex-glommio/src/io.rs b/ntex-glommio/src/io.rs
@@ -196,14 +196,16 @@ impl Future for WriteTask {
                 // use disconnect timeout, otherwise it could hang forever.
                 loop {
                     match st {
-                        Shutdown::None(ref mut fut) => {
+                        Shutdown::Flush => {
                             // flush write buffer
                             let mut io = this.io.0.borrow_mut();
                             match this.state.with_buf(|buf| flush_io(&mut *io, buf, cx)) {
                                 Poll::Ready(Ok(())) => {
                                     let io = this.io.clone();
                                     let fut = Box::pin(async move {
-                                        io.0.borrow().shutdown(std::net::Shutdown::Write).await
+                                        io.0.borrow()
+                                            .shutdown(std::net::Shutdown::Write)
+                                            .await
                                     });
                                     *st = Shutdown::Close(fut);
                                     continue;
@@ -484,11 +486,7 @@ impl Future for UnixWriteTask {
                             sleep(time)
                         };
 
-                        let io = this.io.clone();
-                        let fut = Box::pin(async move {
-                            io.0.borrow().shutdown(std::net::Shutdown::Write).await
-                        });
-                        this.st = IoWriteState::Shutdown(timeout, Shutdown::None(fut));
+                        this.st = IoWriteState::Shutdown(timeout, Shutdown::Flush);
                         self.poll(cx)
                     }
                     Poll::Ready(WriteStatus::Terminate) => {
@@ -506,16 +504,18 @@ impl Future for UnixWriteTask {
                 // use disconnect timeout, otherwise it could hang forever.
                 loop {
                     match st {
-                        Shutdown::None(ref mut fut) => {
+                        Shutdown::Flush => {
                             // flush write buffer
                             let mut io = this.io.0.borrow_mut();
                             match this.state.with_buf(|buf| flush_io(&mut *io, buf, cx)) {
                                 Poll::Ready(Ok(())) => {
-                                    if ready!(fut.poll(cx)).is_err() {
-                                        this.state.close(None);
-                                        return Poll::Ready(());
-                                    }
-                                    *st = Shutdown::Stopping(0);
+                                    let io = this.io.clone();
+                                    let fut = Box::pin(async move {
+                                        io.0.borrow()
+                                            .shutdown(std::net::Shutdown::Write)
+                                            .await
+                                    });
+                                    *st = Shutdown::Close(fut);
                                     continue;
                                 }
                                 Poll::Ready(Err(err)) => {
@@ -529,6 +529,14 @@ impl Future for UnixWriteTask {
                                 Poll::Pending => (),
                             }
                         }
+                        Shutdown::Close(ref mut fut) => {
+                            if ready!(fut.poll(cx)).is_err() {
+                                this.state.close(None);
+                                return Poll::Ready(());
+                            }
+                            *st = Shutdown::Stopping(0);
+                            continue;
+                        }
                         Shutdown::Stopping(ref mut count) => {
                             // read until 0 or err
                             let mut buf = [0u8; 512];