From e713b0c015b660efcde4e7300f700759856d3146 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Mon, 19 Dec 2022 18:03:36 +0100
Subject: [PATCH] fix(net): poll after starting disconnect (#522)

---
 crates/net/network/src/session/active.rs | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/crates/net/network/src/session/active.rs b/crates/net/network/src/session/active.rs
index 2f8ae5cd63..5e7014241f 100644
--- a/crates/net/network/src/session/active.rs
+++ b/crates/net/network/src/session/active.rs
@@ -294,6 +294,16 @@ impl ActiveSession {
         self.conn.inner_mut().start_disconnect(reason);
     }
 
+    /// Flushes the disconnect message and emits the corresponding message
+    fn poll_disconnect(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+        debug_assert!(self.is_disconnecting(), "not disconnecting");
+
+        // try to close the flush out the remaining Disconnect message
+        let _ = ready!(self.conn.poll_close_unpin(cx));
+        self.emit_disconnect();
+        Poll::Ready(())
+    }
+
     /// Removes all timed out requests
     fn evict_timed_out_requests(&mut self, now: Instant) {
         let mut timedout = Vec::new();
@@ -317,10 +327,7 @@ impl Future for ActiveSession {
         let this = self.get_mut();
 
         if this.is_disconnecting() {
-            // try to close the flush out the remaining Disconnect message
-            let _ = ready!(this.conn.poll_close_unpin(cx));
-            this.emit_disconnect();
-            return Poll::Ready(())
+            return this.poll_disconnect(cx)
         }
 
         loop {
@@ -342,6 +349,7 @@ impl Future for ActiveSession {
                                 let reason =
                                     reason.unwrap_or(DisconnectReason::DisconnectRequested);
                                 this.start_disconnect(reason);
+                                return this.poll_disconnect(cx)
                             }
                             SessionCommand::Message(msg) => {
                                 this.on_peer_message(msg);