0

I have a case where I have an http server with axum that receive payload in a very high throughput (could get to 20m a second). I need to take those bytes from the request and do some heavy computations with them. The problem that the memory reach unpredictably high (could reach to 5Gb). This is the current setup on how I'm trying to achieve it:

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let (tx, mut rx) = mpsc::channel::<WriteRequest>(32);

    tokio::spawn(async move {
        while let Some(payload) = rx.recv().await {
            tokio::task::spawn_blocking(move || {
                // Run heavy computation here...
                heavy_computation(payload)
            }).await;
        }
    });

    // build our application with a route
    let app = Router::new()
        .route("/write", post(move |req: Bytes| async move {
             let data: WriteRequest = Message::decode(req);
             // send data here
             let _ = tx.send(data).await;

             "ok"
         }));

    let addr = ([0, 0, 0, 0], 8080).into();
    let server = axum::Server::bind(&addr)
        .serve(app.into_make_service());
    
    if let Err(e) = server.await {
        error!("server error: {}", e);
    }
    
    Ok(())
}

I think it's the back-pressure on the bounded channel that keeps the requests piling up until they can be sent to the other task for processing, resulting in the high memory. Because even if I tried to replace the heavy_computation with a simple sleep for about 200ms it ended with the same results. If I eliminate the heavy_computation part the memory stays low.

What is the right way to approach such problem? Or with this kind of high throughput there is nothing can be done here?

kmdreko
  • 42,554
  • 6
  • 57
  • 106
oronsh
  • 13
  • 3
  • _I think it's the backpressure on the bounded channel that keeps the requests piling up until they can be sent to the other task for processing_ Have you tried with 0 buffered channel? Right now it is set to 32 and the channel will not do any backpressure until there is 32 messages in the queue. ( I assume this channel is tokio::sync::mpsc::channel) – Ömer Erden Sep 16 '21 at 07:53
  • Yes it's `tokio::sync::mpsc::channel` isn't `0 buffered` just an unbounded buffer? – oronsh Sep 16 '21 at 07:58
  • No, unbounded channels are defined explicitly as [another type](https://docs.rs/tokio/1.11.0/tokio/sync/mpsc/fn.unbounded_channel.html), but looks like tokio channel doesn't allow `0` buffered channels it requires at least `1`, please try with 1 buffered channel – Ömer Erden Sep 16 '21 at 08:01
  • Will do! I must say though that I've tried to set it to `2` in the past with no success. – oronsh Sep 16 '21 at 08:09
  • Could you share more info to understand this? What is "20m"? Mb of data for each request or number of requests in millions? Are you sure that it is not heavy_computation that allocates? How big is `req: Bytes`? The total memory should be about: 32 * (heavy_computation_size + req_size) + pending_req_count * pending_req_size; To limit pending_req_count it needs to be configured in axum somehow. – battlmonstr Sep 16 '21 at 08:43
  • So each request could be about 200Kb. `heavy_computation` does aggregate the bytes to some extent, then clear it after some interval, lets say 30seconds. I thought it's the cpu aspect of this operation that stuck because even if I replace it with `thread::sleep` it stays the same. If I just passing the data to the channel without doing anything the memory stays low. – oronsh Sep 16 '21 at 08:48

2 Answers2

1

It feels like while the heavy_computation is busy, millions of pending requests pile up. A limit of the number of accepted connections/requests handled at a time is needed. To get to 5Gb usage, it takes just 25K pending requests with 200Kb payloads, not even millions.

axum is based on tower, and tower is based on hyper.

It is a known issue that hyper doesn't have a max connections setting, but people there suggest to use a ConcurrencyLimit middleware from tower, and configure it into the Server, or make a custom accept/handle loop.

It might be possible to pass this middleware into tower via axum too, but otherwise if it is an option for you, you could try go to tower directly or even bare hyper, and implement it using the primitives available there for this workload.

battlmonstr
  • 5,841
  • 1
  • 23
  • 33
  • Yes I think you're right with your analysis of the problem. I tried using the `ConcurrencyLimit` middleware with no help, I think I'll research more in this direction, thanks a lot! – oronsh Sep 16 '21 at 09:25
0

Basically you have too little resources to service all the requests being made. You should start rejecting them at some point. If your server can handle 10 heavy-computations per second but you're receiving 50 requests per second, then your tasks are going to build up and build up until your memory usage explodes.

I think it is the lack of back pressure that is the root of the problem. There's nothing stopping new requests from being made. Your channel has its own bounded constraints, but because it immediately offloads the work onto tokio's blocking thread pool its not effective.

There's a few ways to solve the problem. The simplest would be to use a semaphore like the answer here so that only N computationally-heavy tasks can be running at a time and don't progress your channel until one has been freed up. Then you can actually have back-pressure so your requests can return an error if the channel is full. With this strategy, your memory usage would only scale to the size of your channel.

There are other strategies and crates around for limiting your resources, but whatever the mechanism, you can't keep piling up requests.

kmdreko
  • 42,554
  • 6
  • 57
  • 106