From 56917cd21e38c8e45d2315f852b57d76038a8480 Mon Sep 17 00:00:00 2001
From: Daniel Abramov <daniel.abramov@element.io>
Date: Tue, 9 May 2023 14:22:03 +0200
Subject: [PATCH] Remove boilerplace in `poll_flush()`

---
 src/lib.rs | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 7e6bb9b..4403b99 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -388,19 +388,13 @@ where
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        match (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.write_pending()) {
-            Ok(()) => Poll::Ready(Ok(())),
-            Err(WsError::ConnectionClosed) => {
-                // WebSocket is closing and there is nothing to send anymore.
-                // Not an failure, the flush operation is a success.
-                Poll::Ready(Ok(()))
-            }
-            Err(WsError::Io(ref e)) if e.kind() == std::io::ErrorKind::WouldBlock => {
-                trace!("WouldBlock");
-                Poll::Pending
+        (*self).with_context(Some((ContextWaker::Write, cx)), |s| cvt(s.write_pending())).map(|r| {
+            // WebSocket connection has just been closed. Flushing completed, not an error.
+            match r {
+                Err(WsError::ConnectionClosed) => Ok(()),
+                other => other,
             }
-            Err(e) => Poll::Ready(Err(e)),
-        }
+        })
     }
 
     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {