High memory usage / large websocket buffer per client remains in memory after message is sent

[RandomGHUser]

Issue: the websocket buffer per client will remain in memory and potentially continue expanding until stopping/closing the connection. This becames a major memory usage issue when sending a large amount of data to client(s).

Related to #3198 and #1967

I would expect the in-memory buffer to flush after a message is sent to a client (please correct me if I'm mistaken with the expectation).

Example below creates a websocket connection, connect to it and send a message "[x]M" (i.e. 5M or 10M) to generate a client message of x million random characters.

Memory starts around 4MB in the example video. Then the memory jumps to expand the in-memory buffer per client to accomodate the message size, but doesn't flush after being sent.

use std::{env, thread};
use ::actix::{Actor, StreamHandler};
use rand::distributions::{Alphanumeric, DistString};
use actix_web_actors::ws::{self};
use actix_web::{web, App, HttpResponse, HttpServer, HttpRequest, Error, get};

pub fn get_address() -> String {
    match env::var("LISTEN_ADDRESS") {
        Ok(e) =>e.to_owned(),
        Err(_e) => "127.0.0.1:8080".to_string()
    }
}
pub fn assert_address() {
    let addr = get_address();
    std::net::TcpListener::bind(&addr).expect(&format!("Could not open port on {addr}"));
}
fn main() {
    assert_address();
    let actix = thread::spawn(|| {
        start().unwrap();
        std::process::exit(0);
    });
    actix.join().unwrap();
}

// actix

#[get("/ws")]
async fn websocket(req: HttpRequest, stream: web::Payload) -> Result<HttpResponse, Error> {
    let resp = ws::start(MyWs{}, &req, stream);
    resp
}
#[actix_web::main]
pub async fn start() -> std::io::Result<()> {
    let address = get_address();
    println!("Serving on {}", address);
    println!("websocat ws://{address}/ws -S -B 1000000000");
    HttpServer::new(|| {
        App::new()
        .service(websocket)
    })
    .bind(address)?
    .run()
    .await
}

// websocket

pub struct MyWs{}
impl Actor for MyWs {
    type Context = ws::WebsocketContext<Self>;
}
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for MyWs {
    fn handle(&mut self, _msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
        if let Ok(ws::Message::Text(text)) = _msg{
            match text.to_string().as_str().split("\n").collect::<Vec<&str>>()[0]{
                "close" => {
                    ctx.text("closing.");
                    ctx.close(None);
                },
                _=> { 
                    let splits:Vec<&str> = text.split("M").collect();
                    if splits.len() < 2{
                        ctx.text("Generate x million characters, usage: [number]M");
                    }
                    else {
                        let c = splits[0];
                        println!("Generating String.");
                        let str = Alphanumeric.sample_string(&mut rand::thread_rng(), 1000000*c.parse::<usize>().unwrap());
                        println!("Done");
                        if splits[1].len() > 1 {
                            println!("Not sending to client.");
                        } else {
                            println!("Sending to client.");
                            ctx.text(str);    
                            println!("Done.");
                        }
                    }
                }
            }
        }
    }
}

example.webm

Edit: updated example along with video.

• Rust Version (I.e, output of rustc -V):
  rustc 1.80.0-nightly (6e1d94708 2024-05-10)
• Actix Web Version:
  4.5.1

[asonix]

I've narrowed the problem down to two buffers:

• write_buf in actix-http's h1 dispatcher
• buf in actix-web-actors' WebsocketContextFut

These buffers each use BytesMut, which is designed to keep its capacity around as long as possible to optimize for bytes moving in and out.

Unfortunately, this means that once a BytesMut grows to fit a set of continuous bytes, it won't shrink back down. In order to send a message over a websocket, these two buffers are filled with the entire contents of the file, the first after passing through the websocket encoder, and the second after passing through the http encoder.

I verified that if I std::mem::take these buffers rather than write_buf.clear() and buf.split() them, memory use returns to basically nothing after the large file is sent.

Since the default behavior is to reuse existing allocations, i'm hesitant to actually mem::take them as a fix here.

It does seem somewhat unfortunate that these buffers are needed at all when it seems that, at least for the http encoder, the actual message contents are not modified at all, and are simply copied verbatim. The ws encoder sometimes copies verbatim although sometimes there is a 4 byte mask applied at the start.

I don't know if I'd be confident enough to try to eliminate these extra buffers myself, though.

I'd be interested to hear @robjtede 's thoughts on potential solutions here