From 72fb249d8fb34037ea6baa7f25f2c2253ac00b61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20CORTIER?= <benoit.cortier@fried-world.eu>
Date: Wed, 2 Mar 2022 18:58:07 -0500
Subject: [PATCH] Refactor similar conditional branches in
 `WebSocketStream::poll_close`

---
 src/lib.rs | 42 ++++++++++++++++--------------------------
 1 file changed, 16 insertions(+), 26 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 7118a01..7af3e08 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -357,34 +357,24 @@ where
     }
 
     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        if self.closing {
+        let res = if self.closing {
             // After queueing it, we call `write_pending` to drive the close handshake to completion.
-            match (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.write_pending()) {
-                Ok(()) => Poll::Ready(Ok(())),
-                Err(::tungstenite::Error::ConnectionClosed) => Poll::Ready(Ok(())),
-                Err(::tungstenite::Error::Io(err))
-                    if err.kind() == std::io::ErrorKind::WouldBlock =>
-                {
-                    trace!("WouldBlock");
-                    Poll::Pending
-                }
-                Err(err) => Poll::Ready(Err(err)),
-            }
+            (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.write_pending())
         } else {
-            match (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.close(None)) {
-                Ok(()) => Poll::Ready(Ok(())),
-                Err(::tungstenite::Error::ConnectionClosed) => Poll::Ready(Ok(())),
-                Err(::tungstenite::Error::Io(err))
-                    if err.kind() == std::io::ErrorKind::WouldBlock =>
-                {
-                    trace!("WouldBlock");
-                    self.closing = true;
-                    Poll::Pending
-                }
-                Err(err) => {
-                    debug!("websocket close error: {}", err);
-                    Poll::Ready(Err(err))
-                }
+            (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.close(None))
+        };
+
+        match res {
+            Ok(()) => Poll::Ready(Ok(())),
+            Err(::tungstenite::Error::ConnectionClosed) => Poll::Ready(Ok(())),
+            Err(::tungstenite::Error::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
+                trace!("WouldBlock");
+                self.closing = true;
+                Poll::Pending
+            }
+            Err(err) => {
+                debug!("websocket close error: {}", err);
+                Poll::Ready(Err(err))
             }
         }
     }