Limiting the number of open sockets in a tokio-based TCP listener
January 08, 2021 [Async programming, Programming, Rust]I learned quite a bit today about how to think about concurrency in Rust. I was trying to use a Semaphore to limit how many open sockets my TCP listener allowed, and I had real trouble making it work. It either didn't actually work, allowing any number of clients to connect, or the compiler told me I couldn't do what I wanted to do, because the lifetime of my Semaphore was not 'static. Here's the journey I took towards working code that I think is correct (feedback welcome).
Motivation
In the tokio tutorial there is a short section entitled "Backpressure and bounded channels" (partway down the Channels page). It contains this statement:
...take care to ensure total amount of concurrency is bounded. For example, when writing a TCP accept loop, ensure that the total number of open sockets is bounded.
Obviously, when I started work on a TCP accept loop, I wanted to follow this advice.
Like many things in my journey with Rust, it was harder than I expected, and eventually enlightening.
The code
Here is a short Rust program that listens on a TCP port and accepts incoming connections.
Cargo.toml:
[package]
name = "tcp-listener-example"
version = "1.0.0"
edition = "2018"
include = ["src/"]
[dependencies]
tokio = { version = ">=1.0.1", features = ["full"] }
src/main.rs:
use tokio::io::AsyncReadExt;
use tokio::net::TcpListener;
#[tokio::main]
async fn main() {
let listener = TcpListener::bind("0.0.0.0:8080").await.unwrap();
loop {
let (mut tcp_stream, _) = listener.accept().await.unwrap();
tokio::spawn(async move {
let mut buf: [u8; 1024] = [0; 1024];
loop {
let n = tcp_stream.read(&mut buf).await.unwrap();
if n == 0 {
return;
}
print!("{}", String::from_utf8_lossy(&buf[0..n]));
}
});
}
}
This program listens on port 8080, and every time a client connects, it spawns an asynchronous task to deal with it.
If I run it with:
cargo run
It starts, and I can connect to it from multiple other processes like this:
telnet 0.0.0.0 8080
Anything I type into the telnet terminal window gets printed out in the terminal where I ran cargo run. The program works: it listens on TCP port 8080 and prints out all the messages it receives.
So what's the problem?
The problem is that this program can be overwhelmed: if lots of processes connect to it, it will accept all the connections, and eventually run out of sockets. This might prevent other things working right on the computer, or it might crash our program, or something else. We need some kind of sensible limit, as the tokio tutorial mentions.
So how do we limit the number of people allowed to connect at the same time?
Just use a semaphore, dummy
A semaphore does exactly what we need here - it keeps a count of how many people are doing something, and prevents that number getting too big. So all we need to do is restrict the number of clients that we allow to connect using a semaphore.
Here was my first attempt:
use tokio::io::AsyncReadExt; use tokio::net::TcpListener; use tokio::sync::Semaphore; #[tokio::main] async fn main() { let listener = TcpListener::bind("0.0.0.0:8080").await.unwrap(); let sem = Semaphore::new(2); loop { let (mut tcp_stream, _) = listener.accept().await.unwrap(); // Don't copy this code: it doesn't work let aq = sem.try_acquire(); if let Ok(_guard) = aq { tokio::spawn(async move { let mut buf: [u8; 1024] = [0; 1024]; loop { let n = tcp_stream.read(&mut buf).await.unwrap(); if n == 0 { return; } print!("{}", String::from_utf8_lossy(&buf[0..n])); } }); } else { println!("Rejecting client: too many open sockets"); } } }
This compiles fine, but it doesn't do anything! Even though we called Semaphore::new with an argument of 2, intending to allow only 2 clients to connect, in fact I can still connect more times than that. It looks like our code changes had no effect at all.
What we were hoping to happen was that every time a client connected, we created _guard, which is a SemaphoreGuard, that occupies one of the slots in the semaphore. We were expecting that guard to live until the client disconnects, at which point the slot will be released.
Why doesn't it work? It's easy to understand when you think about what tokio::spawn does. It creates a task and asks for it to be executed in the future, but it doesn't actually run it. So tokio::spawn returns immediately, and _guard is dropped, before the code that handles the request is executed. So, obviously, our change doesn't actually restrict how many requests are being handled because the semaphore slot is freed up before the request is processed.
Just hold the guard for longer, dummy
So, let's hold on to the SemaphoreGuard for longer:
use tokio::io::AsyncReadExt; use tokio::net::TcpListener; use tokio::sync::Semaphore; #[tokio::main] async fn main() { let listener = TcpListener::bind("0.0.0.0:8080").await.unwrap(); let sem = Semaphore::new(2); loop { let (mut tcp_stream, _) = listener.accept().await.unwrap(); let aq = sem.try_acquire(); if let Ok(guard) = aq { tokio::spawn(async move { let mut buf: [u8; 1024] = [0; 1024]; loop { let n = tcp_stream.read(&mut buf).await.unwrap(); if n == 0 { drop(guard); return; } print!("{}", String::from_utf8_lossy(&buf[0..n])); } }); } else { println!("Rejecting client: too many open sockets"); } } }
The idea is to pass the SemaphoreGuard object into the code that actually deals with the client request. The way I've attempted that is by referring to guard somewhere within the async move closure. What I've actually done is tell it to drop guard when we are finished with the request, but actually any mention of that variable within the closure would have been enough to tell the compiler we want to move it in, and only drop it when we are done.
It all sounds reasonable, but actually this code doesn't compile. Here's the error I get:
error[E0597]: `sem` does not live long enough --> src/main.rs:12:18 | 12 | let aq = sem.try_acquire(); | ^^^-------------- | | | borrowed value does not live long enough | argument requires that `sem` is borrowed for `'static` ... 29 | } | - `sem` dropped here while still borrowed
What the compiler is saying is that our SemaphoreGuard is referring to sem (the Semaphore object), but that the guard might live longer than the semaphore.
Why? Surely sem is held within a scope that includes the whole of the client-handling code, so it should live long enough?
No. Actually, the async move closure that we are passing to tokio::spawn is being added to a list of tasks to run in the future, so it could live much longer. The fact that we are inside an infinite loop confused me further here, but the principle still remains: whenever we make a closure like this and pass something into it, the closure must own it, or if we are borrowing it, it must live forever (which is what a 'static lifetime means).
The code above passes ownership of guard to the closure, but guard itself is referring to (borrowing) sem. This is why the compiler says that "sem is borrowed for 'static".
Wrong things I tried
Because I didn't understand what I was doing, I tried various other things like making sem an Arc, making guard an Arc, creating guard inside the closure, and even trying to make sem actually have 'static storage by making it a constant. (That last one didn't work because only very simple types like numbers and strings can be constants.)
Solution: Share the Semaphore in an Arc
After what felt like too much thrashing around, I found what I think is the right answer:
use std::sync::Arc; use tokio::io::AsyncReadExt; use tokio::net::TcpListener; use tokio::sync::Semaphore; #[tokio::main] async fn main() { let listener = TcpListener::bind("0.0.0.0:8080").await.unwrap(); let sem = Arc::new(Semaphore::new(2)); loop { let (mut tcp_stream, _) = listener.accept().await.unwrap(); let sem_clone = Arc::clone(&sem); tokio::spawn(async move { let aq = sem_clone.try_acquire(); if let Ok(_guard) = aq { let mut buf: [u8; 1024] = [0; 1024]; loop { let n = tcp_stream.read(&mut buf).await.unwrap(); if n == 0 { return; } print!("{}", String::from_utf8_lossy(&buf[0..n])); } } else { println!("Rejecting client: too many open sockets"); } }); } }
This code:
- Creates a Semaphore and stores it inside an Arc, which is a reference-counting pointer that can be shared between tasks. This means it will live as long as someone holds a reference to it.
- Clones the Arc so we have a copy that can be safely moved into the async move closure. We can't move sem in to the closure because it's going to get used again the next time around the loop. We can move sem_clone in to the closure because it's not used anywhere else. sem and sem_clone both refer to the same Semaphore object, so they agree on the count of clients that are connected, but they are different Arc instances, so one can be moved into the closure.
- Only aquires the SemaphoreGuard once we're inside the closure. This way we're not doing something difficult like borrowing a reference to something that lives outside the closure. Instead, we're borrowing a reference via sem_clone, which is owned by the closure which we are inside, so we know it will live long enough.
It actually works! After two clients are connected, listener.accept actually opens a socket to any new client, but because we return almost immediately from the closure, we only hold it open very briefly before dropping it. This seemed preferable to refusing to open it at all, which I thought would probably leave clients hanging, waiting for a connection that might never come.
Lifetimes are cool, and tricky
Once again, I have learned a lot about what my code is really doing from the Rust compiler. I find this stuff really confusing, but hopefully by writing down my understanding in this post I have helped my current and future selves, and maybe even you, be clearer about how to share a semaphore between multiple asynchronous tasks.
It's really fun and empowering to write code that I am reasonably confident is correct, and also works. The sense that "the compiler has my back" is strong, and I like it.