ws: add buffer reclamation to prevent memory bloat on long-lived connections

BytesMut retains its allocation even after being fully drained, which
means a single large WebSocket message permanently inflates memory for
the lifetime of the connection. This adds a configurable
`max_buffer_size` threshold (default 128 KiB) — when the read buffer
is empty and its capacity exceeds this limit, it gets replaced with a
fresh allocation.

New API: `Codec::max_buffer_size(size)` to configure the threshold.

Fixes #2075
This commit is contained in:
Varun Chawla 2026-02-22 20:17:51 -08:00
parent 96a4964c1b
commit 480ee0eef7
No known key found for this signature in database
1 changed files with 27 additions and 0 deletions

View File

@ -71,6 +71,10 @@ pub enum Item {
pub struct Codec { pub struct Codec {
flags: Flags, flags: Flags,
max_size: usize, max_size: usize,
/// When the read buffer is fully drained and its capacity exceeds this
/// threshold, it is replaced with a fresh allocation to reclaim memory.
/// Defaults to 128 KiB.
max_buffer_size: usize,
} }
bitflags! { bitflags! {
@ -87,6 +91,7 @@ impl Codec {
pub const fn new() -> Codec { pub const fn new() -> Codec {
Codec { Codec {
max_size: 65_536, max_size: 65_536,
max_buffer_size: 128 * 1024,
flags: Flags::SERVER, flags: Flags::SERVER,
} }
} }
@ -100,6 +105,20 @@ impl Codec {
self self
} }
/// Set the max buffer capacity threshold for read buffer reclamation.
///
/// When the read buffer is fully drained (empty) and its allocated capacity
/// exceeds this threshold, it is replaced with a fresh, smaller allocation.
/// This prevents long-lived connections from permanently holding memory
/// allocated for occasional large messages.
///
/// By default this is set to 128 KiB.
#[must_use = "This returns the a new Codec, without modifying the original."]
pub fn max_buffer_size(mut self, size: usize) -> Self {
self.max_buffer_size = size;
self
}
/// Set decoder to client mode. /// Set decoder to client mode.
/// ///
/// By default decoder works in server mode. /// By default decoder works in server mode.
@ -220,6 +239,14 @@ impl Decoder for Codec {
type Error = ProtocolError; type Error = ProtocolError;
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> { fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
// Reclaim memory: if the buffer has been fully drained and its capacity
// exceeds the threshold, replace it with a fresh allocation. This prevents
// peak memory usage from becoming permanent on long-lived connections
// (e.g., WebSockets that occasionally receive large payloads).
if src.is_empty() && src.capacity() > self.max_buffer_size {
*src = BytesMut::new();
}
match Parser::parse(src, self.flags.contains(Flags::SERVER), self.max_size) { match Parser::parse(src, self.flags.contains(Flags::SERVER), self.max_size) {
Ok(Some((finished, opcode, payload))) => { Ok(Some((finished, opcode, payload))) => {
// continuation is not supported // continuation is not supported