From dc960fba63470c0fa7e8de47c329431e95f1b528 Mon Sep 17 00:00:00 2001
From: Mauro D <mauro@stalw.art>
Date: Wed, 16 Aug 2023 16:22:09 +0200
Subject: [PATCH] Flush tcp streap after write

---
 crates/imap/src/core/session.rs |  2 ++
 crates/imap/src/core/writer.rs  | 22 ++++++++++++++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/crates/imap/src/core/session.rs b/crates/imap/src/core/session.rs
index 8ee7088e..ab78ac60 100644
--- a/crates/imap/src/core/session.rs
+++ b/crates/imap/src/core/session.rs
@@ -118,6 +118,7 @@ impl Session<TcpStream> {
             tracing::debug!(parent: &session.span, event = "error", reason = %err, "Failed to write greeting.");
             return Err(());
         }
+        let _ = session.stream.flush().await;
 
         // Split stream into read and write halves
         let (stream_rx, stream_tx) = tokio::io::split(session.stream);
@@ -206,6 +207,7 @@ impl Session<TlsStream<TcpStream>> {
             tracing::debug!(parent: &span, event = "error", reason = %err, "Failed to write greeting.");
             return Err(());
         }
+        let _ = stream.flush().await;
 
         // Spit stream into read and write halves
         let (stream_rx, stream_tx) = tokio::io::split(stream);
diff --git a/crates/imap/src/core/writer.rs b/crates/imap/src/core/writer.rs
index b97d7f7b..42f4114e 100644
--- a/crates/imap/src/core/writer.rs
+++ b/crates/imap/src/core/writer.rs
@@ -64,9 +64,14 @@ pub fn spawn_writer(mut stream: Event, span: tracing::Span) -> mpsc::Sender<Even
                                     )
                                 );*/
 
-                                if let Err(err) = stream_tx.write_all(bytes.as_ref()).await {
-                                    debug!("Failed to write to stream: {}", err);
-                                    break 'outer;
+                                match stream_tx.write_all(bytes.as_ref()).await {
+                                    Ok(_) => {
+                                        let _ = stream_tx.flush().await;
+                                    }
+                                    Err(err) => {
+                                        debug!("Failed to write to stream: {}", err);
+                                        break 'outer;
+                                    }
                                 }
                             }
                             Event::Upgrade(channel) => {
@@ -93,9 +98,14 @@ pub fn spawn_writer(mut stream: Event, span: tracing::Span) -> mpsc::Sender<Even
                     while let Some(event) = rx.recv().await {
                         match event {
                             Event::Bytes(bytes) => {
-                                if let Err(err) = stream_tx.write_all(bytes.as_ref()).await {
-                                    debug!("Failed to write to stream: {}", err);
-                                    break 'outer;
+                                match stream_tx.write_all(bytes.as_ref()).await {
+                                    Ok(_) => {
+                                        let _ = stream_tx.flush().await;
+                                    }
+                                    Err(err) => {
+                                        debug!("Failed to write to stream: {}", err);
+                                        break 'outer;
+                                    }
                                 }
                             }
                             _ => {