38. Project: a multi-room chat server with Tokio

The C guide builds this with select() and Go §38 with goroutines and channels. In Tokio each connection is a task, and each room is a broadcast channel.

// src/bin/chat.rs    connect with: nc 127.0.0.1 6000
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::broadcast;

type Rooms = Arc<Mutex<HashMap<String, broadcast::Sender<(u64, String)>>>>;

fn room(rooms: &Rooms, name: &str) -> broadcast::Sender<(u64, String)> {
    let mut map = rooms.lock().unwrap();      // std Mutex is fine: never held across .await
    map.entry(name.to_string()).or_insert_with(|| broadcast::channel(256).0).clone()
}

async fn client(stream: TcpStream, id: u64, rooms: Rooms) -> std::io::Result<()> {
    let (reader, mut writer) = stream.into_split();
    let mut lines = BufReader::new(reader).lines();
    let mut nick = format!("guest{id}");
    let mut tx = room(&rooms, "lobby");
    let mut rx = tx.subscribe();
    writer.write_all(b"welcome. commands: /nick NAME, /join ROOM, /quit\n").await?;

    loop {
        tokio::select! {
            line = lines.next_line() => {
                let Some(line) = line? else { break };            // EOF: client left
                let line = line.trim();
                if line.len() > 512 { writer.write_all(b"line too long\n").await?; continue; }
                match line.split_once(' ') {
                    Some(("/nick", n)) if !n.is_empty() => nick = n.chars().take(24).collect(),
                    Some(("/join", r)) if !r.is_empty() => {
                        tx = room(&rooms, r);
                        rx = tx.subscribe();                      // the old receiver is dropped
                        writer.write_all(format!("joined {r}\n").as_bytes()).await?;
                    }
                    _ if line == "/quit" => break,
                    _ if !line.is_empty() => { let _ = tx.send((id, format!("{nick}: {line}"))); }
                    _ => {}
                }
            }
            msg = rx.recv() => match msg {
                Ok((from, text)) if from != id => writer.write_all(format!("{text}\n").as_bytes()).await?,
                Ok(_) => {}                                       // our own message
                Err(broadcast::error::RecvError::Lagged(n)) => {  // slow client: skip, do not block others
                    writer.write_all(format!("[missed {n} messages]\n").as_bytes()).await?;
                }
                Err(broadcast::error::RecvError::Closed) => break,
            },
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:6000").await?;
    let rooms: Rooms = Arc::default();
    println!("chat on {}", listener.local_addr()?);
    let mut next_id = 0u64;
    loop {
        let (stream, peer) = listener.accept().await?;
        next_id += 1;
        let (id, rooms) = (next_id, Arc::clone(&rooms));
        tokio::spawn(async move {
            if let Err(e) = client(stream, id, rooms).await { eprintln!("{peer}: {e}"); }
        });
    }
}

Design notes that carry over to every real server:

• Slow consumers. broadcast has a fixed buffer. A client that cannot keep up gets Lagged and skips messages, so one slow client never stalls the room. Go chat servers often get this wrong with unbuffered channels, and one stuck client then blocks everyone.
• Input limits. Line length and nickname length are capped. Note that next_line still reads a whole line into memory before checking it, so a client can send a 1 GB line. For a hostile network, use tokio_util::codec::LinesCodec::new_with_max_length.
• Locks and .await. The std::sync::Mutex is held only inside room(), which never awaits. That is why a std mutex is correct here (§31.5).