From 75fcf4c75e716f39005d494dc81a47e4a90eb518 Mon Sep 17 00:00:00 2001
From: anitarua <anita@momentohq.com>
Date: Fri, 7 Aug 2026 15:27:07 -0700
Subject: [PATCH 1/2] fix: flush the socket to ensure large messages are fully
 sent

---
 protosocket-connection/src/connection.rs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs
index f3697a9..0a0eade 100644
--- a/protosocket-connection/src/connection.rs
+++ b/protosocket-connection/src/connection.rs
@@ -369,7 +369,18 @@ impl<
             }
             break if self.send_buffer.is_empty() {
                 log::debug!("send buffer is empty");
-                Ok(false)
+                // TLS double-buffers: poll_write reports plaintext accepted into rustls,
+                // not bytes on the wire. Once the socket fills, rustls holds the remainder
+                // and still returns Ok(n) for the whole message, so we drain send_buffer
+                // and arrive here "done" -- parked with no writability interest, tail
+                // stranded, peer blocked on a message that never completes. poll_flush
+                // pushes the remainder, or returns Pending and registers the wake we would
+                // otherwise miss. No-op on an unbuffered stream.
+                match pin!(&mut self.stream).poll_flush(context) {
+                    Poll::Ready(Ok(())) | Poll::Pending => Ok(false),
+                    Poll::Ready(Err(ref e)) if would_block(e) || interrupted(e) => Ok(false),
+                    Poll::Ready(Err(e)) => Err(e),
+                }
             } else {
                 // I need to figure out how to get this from the os rather than hardcoding.
                 // 16 is the lowest I've seen mention of, and I've seen 1024 more commonly.

From c5ce88c88f04c2efdd9e6de4406b278102d7c2a5 Mon Sep 17 00:00:00 2001
From: anitarua <anita@momentohq.com>
Date: Fri, 7 Aug 2026 16:02:35 -0700
Subject: [PATCH 2/2] fix: error branch and comment

---
 protosocket-connection/src/connection.rs | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/protosocket-connection/src/connection.rs b/protosocket-connection/src/connection.rs
index 0a0eade..d13c649 100644
--- a/protosocket-connection/src/connection.rs
+++ b/protosocket-connection/src/connection.rs
@@ -369,16 +369,10 @@ impl<
             }
             break if self.send_buffer.is_empty() {
                 log::debug!("send buffer is empty");
-                // TLS double-buffers: poll_write reports plaintext accepted into rustls,
-                // not bytes on the wire. Once the socket fills, rustls holds the remainder
-                // and still returns Ok(n) for the whole message, so we drain send_buffer
-                // and arrive here "done" -- parked with no writability interest, tail
-                // stranded, peer blocked on a message that never completes. poll_flush
-                // pushes the remainder, or returns Pending and registers the wake we would
-                // otherwise miss. No-op on an unbuffered stream.
+                // Ensure that any remaining data in the TLS buffer is flushed to the socket.
                 match pin!(&mut self.stream).poll_flush(context) {
                     Poll::Ready(Ok(())) | Poll::Pending => Ok(false),
-                    Poll::Ready(Err(ref e)) if would_block(e) || interrupted(e) => Ok(false),
+                    Poll::Ready(Err(ref e)) if would_block(e) || interrupted(e) => continue,
                     Poll::Ready(Err(e)) => Err(e),
                 }
             } else {