// Companion to /blog/virtual-threads-for-rust.
// Put this file at src/main.rs in a Cargo application with edition = "2024".
// Add this dependency to Cargo.toml:
// vthread = "=0.1.0"
// Run with `cargo run`, then open the printed address with `nc HOST PORT`
// in two terminals. Leave one client idle and type in the other.
// The server accepts exactly two connections and exits after both reach EOF.

use vthread::net::{TcpListener, TcpStream};

fn echo(stream: TcpStream) -> vthread::Result<()> {
    let mut buffer = [0; 1024];

    loop {
        let count = stream.read(&mut buffer)?;
        if count == 0 {
            return Ok(());
        }
        stream.write_all(&buffer[..count])?;
    }
}

fn serve(listener: TcpListener) -> vthread::Result<()> {
    vthread::local_scope(|scope| {
        let mut connections = Vec::new();
        for _ in 0..2 {
            let (stream, _) = listener.accept()?;
            connections.push(scope.spawn("connection", move || echo(stream))?);
        }
        for mut connection in connections {
            connection.join()??;
        }
        Ok(())
    })
}

fn main() -> vthread::Result<()> {
    let runtime = vthread::Runtime::builder().carriers(1).build()?;
    let listener = TcpListener::bind(([127, 0, 0, 1], 0).into())?;
    println!("Listening on {}", listener.local_addr()?);

    let result = runtime.run_scope(|scope| {
        let mut server = scope.spawn("server", move || serve(listener))?;
        server.join()?
    });
    let shutdown = runtime.shutdown();
    result?;
    shutdown?;
    Ok(())
}