From c991fce9b20c64a700b86b131a8f8f315909f1e2 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Tue, 3 Sep 2024 17:46:45 +0500
Subject: [PATCH] Write handling improvements

---
 ntex-compio/src/io.rs | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/ntex-compio/src/io.rs b/ntex-compio/src/io.rs
index 7f1220c0..aaabbe7d 100644
--- a/ntex-compio/src/io.rs
+++ b/ntex-compio/src/io.rs
@@ -225,21 +225,23 @@ async fn write<T: AsyncWrite>(io: &mut T, state: &WriteContext) -> io::Result<()
                 let BufResult(result, buf1) = io.write(buf).await;
                 buf = buf1;
 
-                match result {
+                return match result {
+                    Ok(0) => Err(io::Error::new(
+                        io::ErrorKind::WriteZero,
+                        "failed to write frame to transport",
+                    )),
                     Ok(size) => {
                         if buf.0.len() == size {
-                            return io.flush().await;
+                            // return io.flush().await;
+                            state.memory_pool().release_write_buf(buf.0);
+                            Ok(())
+                        } else {
+                            buf.0.advance(size);
+                            continue;
                         }
-                        if size == 0 {
-                            return Err(io::Error::new(
-                                io::ErrorKind::WriteZero,
-                                "failed to write frame to transport",
-                            ));
-                        }
-                        buf.0.advance(size);
                     }
-                    Err(e) => return Err(e),
-                }
+                    Err(e) => Err(e),
+                };
             }
         })
         .await