From fca73004eb6222d2cc9c8dc9c1146ae1d27196e0 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Thu, 13 Feb 2025 19:57:17 +0200
Subject: [PATCH 1/3] websocket/stream: Avoid memory allocations on flushing

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 src/transport/websocket/stream.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transport/websocket/stream.rs b/src/transport/websocket/stream.rs
index 268d39e18..a108e2395 100644
--- a/src/transport/websocket/stream.rs
+++ b/src/transport/websocket/stream.rs
@@ -127,7 +127,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> futures::AsyncWrite for BufferedStream<S
                         // TODO: optimize
                         self.state = State::ReadyToSend;
                         self.write_ptr = 0;
-                        self.write_buffer = Vec::with_capacity(2000);
+                        self.write_buffer.clear();
                         return Poll::Ready(Ok(()));
                     }
                     Err(_) => return Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into())),

From f2275dd4ed7c69d781ffcf2b245c4da4db9de55f Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Thu, 13 Feb 2025 19:58:36 +0200
Subject: [PATCH 2/3] websocket/stream: Increase buffer capacity from 2000 to
 8192

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 src/transport/websocket/stream.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/transport/websocket/stream.rs b/src/transport/websocket/stream.rs
index a108e2395..fcc453043 100644
--- a/src/transport/websocket/stream.rs
+++ b/src/transport/websocket/stream.rs
@@ -33,6 +33,8 @@ use std::{
 
 // TODO: add tests
 
+const DEFAULT_BUF_SIZE: usize = 8 * 1024;
+
 /// Send state.
 enum State {
     /// State is poisoned.
@@ -70,7 +72,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> BufferedStream<S> {
     /// Create new [`BufferedStream`].
     pub(super) fn new(stream: WebSocketStream<S>) -> Self {
         Self {
-            write_buffer: Vec::with_capacity(2000),
+            write_buffer: Vec::with_capacity(DEFAULT_BUF_SIZE),
             read_buffer: None,
             write_ptr: 0usize,
             stream,
@@ -124,7 +126,6 @@ impl<S: AsyncRead + AsyncWrite + Unpin> futures::AsyncWrite for BufferedStream<S
                 }
                 State::FlushPending => match futures::ready!(self.stream.poll_flush_unpin(cx)) {
                     Ok(_res) => {
-                        // TODO: optimize
                         self.state = State::ReadyToSend;
                         self.write_ptr = 0;
                         self.write_buffer.clear();

From 3244d63ad0deacb9a14dd87f182e9bd3bfd3efd1 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Mon, 17 Feb 2025 19:49:21 +0200
Subject: [PATCH 3/3] websocket/stream: Shrink to fit based on capacity to
 bound growth

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 src/transport/websocket/stream.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/transport/websocket/stream.rs b/src/transport/websocket/stream.rs
index fcc453043..2dd20091c 100644
--- a/src/transport/websocket/stream.rs
+++ b/src/transport/websocket/stream.rs
@@ -129,6 +129,9 @@ impl<S: AsyncRead + AsyncWrite + Unpin> futures::AsyncWrite for BufferedStream<S
                         self.state = State::ReadyToSend;
                         self.write_ptr = 0;
                         self.write_buffer.clear();
+                        // In the unlikely event that the buffer is too large, we need to bound the
+                        // capacity to avoid unbounded memory usage.
+                        self.write_buffer.shrink_to(DEFAULT_BUF_SIZE);
                         return Poll::Ready(Ok(()));
                     }
                     Err(_) => return Poll::Ready(Err(std::io::ErrorKind::UnexpectedEof.into())),