Chapter 26: Network Programming

Project 1: Multi-Protocol Chat Server (TCP + WebSocket)

Problem Statement

Build a real-time chat server that evolves from a simple echo server to a production-ready multi-protocol system. You’ll start with basic synchronous TCP, progress through async patterns, add room-based architecture, support WebSocket clients, and finish with production features like health monitoring and graceful shutdown.

Why It Matters

Real-World Impact: Chat servers are the backbone of modern communication:

• Slack/Discord: Handle millions of concurrent WebSocket connections, delivering sub-100ms message latency
• Gaming: Multiplayer games need real-time chat for team coordination (League of Legends: 8M concurrent players)
• Trading platforms: Order updates and price feeds require instant delivery to thousands of traders
• Customer support: Live chat systems handle 100K+ concurrent support conversations

Performance Numbers:

• Synchronous blocking I/O: 1 thread per client = ~1,000 max clients (stack exhaustion at 2MB/thread)
• Async I/O (tokio): 1 thread handles 10,000+ clients (C10K problem solved)
• Message latency: WebSocket ~20ms vs HTTP polling ~500ms (25x faster)
• Memory: Thread-per-connection = 2MB/client, async tasks = 2KB/client (1000x improvement)

Rust-Specific Challenge: Traditional chat servers in languages like Python or Node.js use mutable shared state with locks everywhere. Rust’s ownership system forces us to design better concurrent architectures using channels and message passing. This project teaches you to embrace Rust’s async model and build lock-free broadcast patterns using tokio’s channels.

Use Cases

When you need this pattern:

1. Real-time messaging apps - Slack, Discord, Telegram (persistent connections, instant delivery)
2. Multiplayer games - In-game chat, lobby systems, team communication
3. Live collaboration tools - Code editors (VS Code Live Share), design tools (Figma comments)
4. Financial trading platforms - Order updates, market data feeds, trader chat rooms
5. IoT command and control - Device management consoles, sensor monitoring dashboards
6. Live streaming platforms - Chat alongside video (Twitch, YouTube Live)

Real Examples:

• Discord: Uses WebSocket for real-time chat, handles 19M concurrent connections (2021)
• Slack: WebSocket for messages, falls back to HTTP polling for old clients
• WhatsApp: Custom protocol over TCP for 2B users, E2E encrypted
• IRC servers (UnrealIRCd): Classic TCP-based chat, thousands of channels per server

Learning Goals

• Master async I/O with tokio (TcpListener, TcpStream, spawning tasks)
• Understand broadcast patterns with tokio::sync::broadcast channels
• Learn WebSocket protocol and HTTP upgrade mechanism
• Practice concurrent state management with Arc<RwLock>
• Build production features (graceful shutdown, metrics, health checks)
• Experience the performance difference: sync vs async, thread-per-connection vs task-per-connection

Core Concepts

Before diving into the implementation, let’s understand the fundamental concepts that power modern network servers:

1. TCP Networking Fundamentals

What is TCP? Transmission Control Protocol (TCP) is a reliable, connection-oriented network protocol that guarantees ordered delivery of data between applications over a network.

Key Components:

• Socket: An endpoint for network communication (IP address + port number)
• TcpListener: Listens for incoming connections on a specific port
• TcpStream: Represents an established connection to a remote client
• Buffered I/O: Reading/writing data efficiently using buffers

How it works:

#![allow(unused)]
fn main() {
// Server side
let listener = TcpListener::bind("127.0.0.1:8080")?;  // Listen on port 8080

for stream in listener.incoming() {  // Wait for connections
    let stream = stream?;  // TcpStream to connected client

    // Read from client
    let mut buffer = [0u8; 1024];
    stream.read(&mut buffer)?;

    // Write to client
    stream.write_all(b"Hello, client!")?;
}
}

Client side:

#![allow(unused)]
fn main() {
let mut stream = TcpStream::connect("127.0.0.1:8080")?;
stream.write_all(b"Hello, server!")?;

let mut buffer = [0u8; 1024];
stream.read(&mut buffer)?;
}

Line-Based Protocol: Most text-based protocols (like HTTP, SMTP, chat) use newline-delimited messages:

#![allow(unused)]
fn main() {
// Using BufReader for efficient line reading
let reader = BufReader::new(&stream);
for line in reader.lines() {
    let line = line?;  // Read until '\n'
    println!("Received: {}", line);
}
}

2. Thread-per-Connection vs Task-per-Connection

Thread-per-Connection (Traditional):

#![allow(unused)]
fn main() {
for stream in listener.incoming() {
    std::thread::spawn(move || {
        handle_client(stream);  // Each client gets its own OS thread
    });
}
}

Costs:

• Each thread = ~2MB stack memory (1,000 threads = 2GB RAM)
• Context switching overhead between threads
• Limited by OS thread limits (~10,000 max on Linux)

Task-per-Connection (Async):

#![allow(unused)]
fn main() {
loop {
    let (stream, _) = listener.accept().await?;
    tokio::spawn(async move {
        handle_client(stream).await;  // Lightweight async task
    });
}
}

Benefits:

• Each task = ~2KB memory (1,000x smaller)
• Cooperative multitasking (no context switch overhead)
• Can handle 100,000+ concurrent connections

The C10K Problem: In the 2000s, servers struggled to handle 10,000 concurrent connections due to thread limitations. Async I/O solved this.

3. Async/Await and Tokio Runtime

What is async/await? Async/await is Rust’s way of writing non-blocking concurrent code that looks like synchronous code.

Without async (blocking):

#![allow(unused)]
fn main() {
fn download(url: &str) -> String {
    // Blocks thread for seconds while waiting for network
    http_get(url)
}

// Sequential - takes 6 seconds total
let data1 = download("url1");  // 2 seconds
let data2 = download("url2");  // 2 seconds
let data3 = download("url3");  // 2 seconds
}

With async (non-blocking):

#![allow(unused)]
fn main() {
async fn download(url: &str) -> String {
    // Yields control while waiting, other tasks can run
    http_get(url).await
}

// Concurrent - takes 2 seconds total (all run simultaneously)
let (data1, data2, data3) = tokio::join!(
    download("url1"),
    download("url2"),
    download("url3"),
);
}

Key Concepts:

• Future: A value that will be available in the future (lazy, does nothing until awaited)
• await: Suspends current task until future completes, yields CPU to other tasks
• Runtime: Tokio’s executor that schedules and runs async tasks
• Task: Lightweight unit of execution (like green threads)

Tokio Runtime:

#[tokio::main]  // Creates runtime, runs async main
async fn main() {
    // This code runs on tokio's thread pool
    let task1 = tokio::spawn(async { /* work */ });
    let task2 = tokio::spawn(async { /* work */ });

    // Both tasks run concurrently on the runtime
    task1.await.unwrap();
    task2.await.unwrap();
}

When to await?:

• Blocking operations: listener.accept().await (waits for connection)
• I/O operations: stream.read().await, stream.write().await
• Waiting for tasks: task.await (wait for task to complete)

4. Broadcast Channels (Publish-Subscribe Pattern)

The Problem: How do we send one message to many receivers?

Naive approach (doesn’t work):

#![allow(unused)]
fn main() {
let msg = "Hello".to_string();
for client in clients {
    client.send(msg);  // ERROR: msg moved on first iteration
}
}

Solution: Broadcast Channel:

#![allow(unused)]
fn main() {
use tokio::sync::broadcast;

// Create channel with capacity 100
let (tx, _rx) = broadcast::channel::<String>(100);

// Each subscriber gets their own receiver
let rx1 = tx.subscribe();
let rx2 = tx.subscribe();
let rx3 = tx.subscribe();

// Send message once
tx.send("Hello everyone".to_string()).ok();

// All receivers get the message
assert_eq!(rx1.recv().await.unwrap(), "Hello everyone");
assert_eq!(rx2.recv().await.unwrap(), "Hello everyone");
assert_eq!(rx3.recv().await.unwrap(), "Hello everyone");
}

How it works:

• Sender (tx): Cloneable, can be shared across tasks
• Receiver (rx): Created via tx.subscribe(), each gets a copy of messages
• Broadcasting: tx.send(msg) clones message to all receivers
• Capacity: Older messages are dropped if receivers are slow (configurable)

Perfect for chat servers:

#![allow(unused)]
fn main() {
// When client sends message
tx.send(format!("{}: {}", username, message)).ok();

// All clients receive it via their rx.recv().await
}

5. Shared State with Arc and RwLock

The Problem: Multiple tasks need to access the same data (room list, user list).

Ownership Challenge:

#![allow(unused)]
fn main() {
let rooms = HashMap::new();

tokio::spawn(async move {
    rooms.insert("general", room);  // ERROR: rooms moved here
});

tokio::spawn(async move {
    rooms.insert("random", room);  // ERROR: rooms already moved
});
}

Solution: Arc (Atomic Reference Counting):

#![allow(unused)]
fn main() {
use std::sync::Arc;

let rooms = Arc::new(HashMap::new());

// Clone Arc (cheap - just increments ref count)
let rooms1 = rooms.clone();
let rooms2 = rooms.clone();

tokio::spawn(async move {
    // rooms1 is a smart pointer to shared data
});

tokio::spawn(async move {
    // rooms2 points to same data
});
}

But Arc is immutable! We need interior mutability:

RwLock (Read-Write Lock):

#![allow(unused)]
fn main() {
use tokio::sync::RwLock;

let rooms = Arc::new(RwLock::new(HashMap::new()));

// Many readers simultaneously (doesn't block each other)
{
    let rooms_read = rooms.read().await;  // Shared lock
    let room = rooms_read.get("general");
    // ... read operations ...
}  // Lock released

// Only one writer at a time (blocks readers and other writers)
{
    let mut rooms_write = rooms.write().await;  // Exclusive lock
    rooms_write.insert("new_room".to_string(), room);
}  // Lock released
}

Pattern Summary:

• Arc<T>: Share ownership across tasks
• RwLock<T>: Allow concurrent reads, exclusive writes
• Arc<RwLock<HashMap<K, V>>>: Shared mutable state

Performance:

• Read lock: Multiple simultaneous readers (great for read-heavy workloads)
• Write lock: Exclusive access (blocks everything)
• Choose wisely: Lock only what you need, release quickly

6. WebSocket Protocol

What is WebSocket? WebSocket is a protocol that enables full-duplex (two-way) communication between a client and server over a single TCP connection. Unlike HTTP (request-response), WebSocket keeps the connection open for real-time messaging.

HTTP vs WebSocket:

HTTP (Request-Response):
Client → Server: GET /messages HTTP/1.1
Server → Client: HTTP/1.1 200 OK [messages]
[Connection closes]
[Client must poll repeatedly]

WebSocket (Persistent Connection):
Client → Server: [HTTP Upgrade request]
Server → Client: [HTTP 101 Switching Protocols]
[Connection stays open]
Client ↔ Server: [Messages flow both ways anytime]

HTTP Upgrade Handshake:

Client:
GET /ws HTTP/1.1
Host: localhost:3000
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Sec-WebSocket-Version: 13

Server:
HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=

[Now WebSocket connection is established]

WebSocket Messages:

• Text: String messages (Message::Text("hello"))
• Binary: Raw bytes (Message::Binary(vec![1,2,3]))
• Ping/Pong: Keepalive heartbeat
• Close: Graceful shutdown

Using WebSocket in Rust with Axum:

#![allow(unused)]
fn main() {
use axum::extract::ws::{WebSocket, WebSocketUpgrade};

async fn websocket_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(|socket: WebSocket| async {
        let (mut sender, mut receiver) = socket.split();

        // Send to client
        sender.send(Message::Text("Hello".into())).await.ok();

        // Receive from client
        while let Some(Ok(msg)) = receiver.next().await {
            match msg {
                Message::Text(text) => println!("Got: {}", text),
                Message::Close(_) => break,
                _ => {}
            }
        }
    })
}
}

Why WebSocket for Chat?

• Real-time: Messages delivered instantly (no polling)
• Efficient: One connection, not thousands of HTTP requests
• Browser support: Built into all modern browsers
• Bi-directional: Server can push to clients anytime

7. Graceful Shutdown

The Problem: When server stops (Ctrl+C, deployment, crash), active connections are abruptly closed, losing in-flight messages.

What is Graceful Shutdown? A clean stop where the server:

1. Stops accepting new connections
2. Drains existing connections (let them finish)
3. Waits for in-flight work to complete
4. Shuts down cleanly

Without Graceful Shutdown:

#![allow(unused)]
fn main() {
// Ctrl+C → Process killed → All connections dropped
// Lost messages, clients see connection errors
}

With Graceful Shutdown:

#![allow(unused)]
fn main() {
use tokio_util::sync::CancellationToken;

let token = CancellationToken::new();
let token_clone = token.clone();

// Listen for Ctrl+C
tokio::spawn(async move {
    tokio::signal::ctrl_c().await.ok();
    token_clone.cancel();  // Signal shutdown
});

loop {
    tokio::select! {
        // Normal operation
        result = listener.accept() => {
            // Handle new connection
        }

        // Shutdown signal
        _ = token.cancelled() => {
            println!("Stopping new connections...");
            break;  // Stop accepting
        }
    }
}

// Wait for existing connections to finish (timeout after 30s)
println!("Draining connections...");
}

Kubernetes Integration:

1. kubectl delete pod chat-server
2. Kubernetes sends SIGTERM to pod
3. Server receives signal, stops accepting connections
4. Server drains existing connections (30s grace period)
5. Kubernetes sends SIGKILL if still running after 30s

Production Benefits:

• Zero downtime deployments: Old version drains while new version starts
• Data integrity: No lost messages during restart
• Better UX: Clients see clean disconnection, not errors

8. Production Observability (Metrics and Health Checks)

The Problem: Production servers are black boxes. Is it healthy? Overloaded? How many users?

Health Checks:

#![allow(unused)]
fn main() {
// Simple endpoint for load balancers
async fn health_handler() -> &'static str {
    "OK"  // 200 status = healthy
}

// Load balancer uses this:
// - Every 10s: GET /health
// - If 200 OK: Send traffic
// - If error/timeout: Remove from rotation
}

Metrics (Prometheus Format):

#![allow(unused)]
fn main() {
use std::sync::atomic::{AtomicUsize, Ordering};

struct Metrics {
    active_connections: AtomicUsize,
    total_messages: AtomicUsize,
    active_rooms: AtomicUsize,
}

async fn metrics_handler(metrics: Arc<Metrics>) -> String {
    format!(
        "active_connections {}\ntotal_messages {}\nactive_rooms {}\n",
        metrics.active_connections.load(Ordering::Relaxed),
        metrics.total_messages.load(Ordering::Relaxed),
        metrics.active_rooms.load(Ordering::Relaxed),
    )
}
}

Prometheus scrapes this:

GET /metrics every 15 seconds
Stores time-series data
Grafana visualizes it
Alerts fire if thresholds exceeded

Real-World Example:

Dashboard shows:
- Active connections: 15,234 (up from 10k an hour ago)
- Message rate: 2,500 msg/sec
- CPU: 45% (healthy)
- Memory: 2.1GB (healthy)

Alert fires: "Active connections > 20,000" → Scale up!

Why Atomic Types?

#![allow(unused)]
fn main() {
// Lock-free concurrent access
metrics.total_messages.fetch_add(1, Ordering::Relaxed);

// Multiple threads can increment simultaneously
// No mutex, no blocking, extremely fast
}

Connection to This Project

Now let’s see how all these concepts come together to build our chat server:

1. Progressive Architecture Evolution

This project takes you through the evolution of network server architecture:

• Milestone 1 (Blocking I/O): Learn TCP fundamentals with TcpListener and TcpStream. Understand the baseline: one client at a time, completely blocking.

• Milestone 2 (Thread-per-Connection): Scale to multiple clients using std::thread::spawn. Experience the limitation: 1,000 clients = 2GB RAM. This is how Apache and traditional servers work.

• Milestone 3 (Async/Task-per-Connection): Breakthrough to modern async I/O with tokio. One runtime handles 100,000+ clients. This is how Discord, Slack, and all modern chat servers work.

2. Broadcast Pattern for Chat

The core of any chat server is broadcasting messages:

• tokio::sync::broadcast implements the pub/sub pattern perfectly
• When Client A sends “Hello”, the server broadcasts to all clients in the room
• Each client has its own Receiver, all connected to one Sender
• Messages are cloned efficiently to each subscriber

Architecture:

Client A ──> [Reader Task] ──> broadcast::Sender ──┬──> Receiver ──> [Writer Task] ──> Client A
                                                    ├──> Receiver ──> [Writer Task] ──> Client B
                                                    └──> Receiver ──> [Writer Task] ──> Client C

3. Room Isolation with Shared State

Milestone 4 adds rooms using the Arc<RwLock<HashMap>> pattern:

#![allow(unused)]
fn main() {
struct ChatServer {
    rooms: Arc<RwLock<HashMap<String, Room>>>,
}

struct Room {
    tx: broadcast::Sender<String>,   // One broadcast channel per room
    users: HashSet<SocketAddr>,       // Who's in this room?
}
}

Why this structure?

• Arc: Share across all client tasks
• RwLock: Many concurrent readers (checking rooms), few writers (join/leave)
• HashMap<String, Room>: Fast lookup by room name
• Each room has its own broadcast channel for message isolation

4. Multi-Protocol Support

Milestone 5 demonstrates protocol abstraction:

• Same backend (ChatServer) serves both TCP and WebSocket
• TCP clients and WebSocket clients chat together seamlessly
• WebSocket uses HTTP upgrade from axum
• Both protocols share the same room and broadcast infrastructure

Why this matters: Modern apps need multiple protocols. Your API might be HTTP REST, but real-time features need WebSocket. This project shows how to unify them.

5. Production-Ready Features

Milestone 6 adds the polish that separates toys from production systems:

• Graceful Shutdown: CancellationToken signals all tasks to stop accepting work and drain
• Metrics: AtomicUsize for lock-free counters, Prometheus format for observability
• Health Checks: /health endpoint for Kubernetes liveness probes
• Keepalive: WebSocket ping/pong to detect and clean up dead connections

6. Async I/O Patterns

Throughout the project, you’ll master critical async patterns:

Splitting streams:

#![allow(unused)]
fn main() {
let (reader, writer) = stream.into_split();
// Now we can read and write concurrently
}

Select pattern (wait for multiple events):

#![allow(unused)]
fn main() {
tokio::select! {
    msg = client_receiver.recv() => { /* client input */ }
    msg = room_broadcast.recv() => { /* broadcast from room */ }
    _ = shutdown_token.cancelled() => { /* shutdown */ }
}
}

Spawning tasks:

#![allow(unused)]
fn main() {
tokio::spawn(async move {
    // Runs concurrently on tokio runtime
    handle_client(stream).await;
});
}

7. Real-World Performance Gains

By the end, you’ll have built:

8. Why This Architecture Matters

This exact architecture pattern is used by:

• Discord: WebSocket for real-time chat, handles 19M concurrent connections
• Slack: WebSocket primary, HTTP fallback, room-based channels
• Gaming servers: Lobbies and team chat in multiplayer games
• Trading platforms: Real-time order updates to thousands of traders
• IoT platforms: Millions of devices sending sensor data

What You’ll Understand:

After completing this project, when you see “built with async Rust and tokio,” you’ll know exactly what that means:

• Lightweight tasks instead of heavy threads
• Non-blocking I/O that scales to hundreds of thousands of connections
• Broadcast channels for efficient pub/sub
• Shared state with lock-free atomics where possible, RwLock where needed
• Production features like graceful shutdown and observability

This is the foundation of modern high-performance network services in Rust!

Milestone 1: Simple Echo Server (Synchronous, Single-Threaded)

Introduction

Starting Point: Before building a chat server, we need to understand the fundamentals of TCP networking. An echo server is the “Hello World” of network programming—it accepts a connection, reads data, and writes it back.

What We’re Building: A synchronous TCP server that:

• Listens on a port (e.g., 8080)
• Accepts ONE client at a time (blocking)
• Reads lines from the client
• Echoes each line back
• Handles disconnection gracefully

Key Limitation: This server can only handle one client at a time. While Client A is connected, Client B cannot connect—it must wait in the OS accept queue. This is completely unacceptable for production but perfect for learning TCP basics.

Key Concepts

Structs/Types:

• TcpListener - Listens for incoming connections on a port
• TcpStream - Represents a connection to a client
• BufReader<TcpStream> - Buffered reading for line-based protocols
• BufWriter<TcpStream> - Buffered writing for efficiency

Functions and Their Roles:

#![allow(unused)]
fn main() {
// In main.rs or lib.rs

fn run_echo_server(addr: &str) -> io::Result<()>
    // Binds TcpListener to address
    // Loops accepting connections
    // Calls handle_client for each connection

fn handle_client(stream: TcpStream) -> io::Result<()>
    // Wraps stream in BufReader for line reading
    // Loops reading lines until EOF
    // Echoes each line back to client
    // Returns Ok(()) or Err on I/O error
}

Protocol Design:

• Line-based protocol: each message ends with \n
• Client sends: "Hello\n"
• Server echoes: "Hello\n"
• Client closes connection → server sees EOF (0 bytes read)

Checkpoint Tests

#![allow(unused)]
fn main() {
#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::TcpStream;
    use std::thread;
    use std::time::Duration;

    #[test]
    fn test_echo_single_message() {
        // Start server in background thread
        thread::spawn(|| {
            run_echo_server("127.0.0.1:9001").unwrap();
        });
        thread::sleep(Duration::from_millis(100)); // Wait for server to start

        // Connect and send message
        let mut stream = TcpStream::connect("127.0.0.1:9001").unwrap();
        stream.write_all(b"Hello\n").unwrap();

        // Read echo
        let mut buf = [0u8; 1024];
        let n = stream.read(&mut buf).unwrap();
        assert_eq!(&buf[..n], b"Hello\n");
    }

    #[test]
    fn test_echo_multiple_lines() {
        thread::spawn(|| {
            run_echo_server("127.0.0.1:9002").unwrap();
        });
        thread::sleep(Duration::from_millis(100));

        let mut stream = TcpStream::connect("127.0.0.1:9002").unwrap();
        stream.write_all(b"First\nSecond\nThird\n").unwrap();

        let mut reader = BufReader::new(&stream);
        let mut line = String::new();

        reader.read_line(&mut line).unwrap();
        assert_eq!(line, "First\n");

        line.clear();
        reader.read_line(&mut line).unwrap();
        assert_eq!(line, "Second\n");

        line.clear();
        reader.read_line(&mut line).unwrap();
        assert_eq!(line, "Third\n");
    }

    #[test]
    fn test_handles_disconnect() {
        thread::spawn(|| {
            run_echo_server("127.0.0.1:9003").unwrap();
        });
        thread::sleep(Duration::from_millis(100));

        let stream = TcpStream::connect("127.0.0.1:9003").unwrap();
        drop(stream); // Close connection
        // Server should handle gracefully without panicking
        thread::sleep(Duration::from_millis(100));
    }
}
}

Starter Code

use std::io::{self, BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};

fn run_echo_server(addr: &str) -> io::Result<()> {
    // TODO: Create TcpListener and bind to address
    let listener = todo!(); // TcpListener::bind(addr)?

    println!("Echo server listening on {}", addr);

    // TODO: Loop accepting connections
    for stream in listener.incoming() {
        // TODO: Handle connection result
        match stream {
            Ok(stream) => {
                // TODO: Get client address for logging
                let peer = todo!(); // stream.peer_addr()?
                println!("New client: {}", peer);

                // TODO: Handle this client (blocking)
                if let Err(e) = handle_client(stream) {
                    eprintln!("Error handling client: {}", e);
                }
            }
            Err(e) => {
                eprintln!("Connection failed: {}", e);
            }
        }
    }

    Ok(())
}

fn handle_client(stream: TcpStream) -> io::Result<()> {
    // TODO: Create BufReader for line-based reading
    let mut reader = todo!(); // BufReader::new(&stream)

    // TODO: Create writer for sending responses
    let mut writer = todo!(); // &stream or BufWriter::new(&stream)

    let mut line = String::new();

    loop {
        line.clear();

        // TODO: Read a line from client
        let bytes_read = todo!(); // reader.read_line(&mut line)?

        // TODO: Check for EOF (client disconnected)
        if bytes_read == 0 {
            println!("Client disconnected");
            break;
        }

        // TODO: Echo the line back to client
        // writer.write_all(line.as_bytes())?
        // writer.flush()?
        todo!();
    }

    Ok(())
}

fn main() {
    if let Err(e) = run_echo_server("127.0.0.1:8080") {
        eprintln!("Server error: {}", e);
    }
}

Check Your Understanding

• Why does listener.incoming() block the entire program? Because it’s synchronous—the thread waits until a connection arrives.
• What happens if Client A is connected and Client B tries to connect? Client B waits in the OS accept queue until Client A disconnects.
• Why use BufReader instead of reading bytes directly? Efficiency—buffering reduces system calls. read_line reads until \n efficiently.
• What does read_line return when the client closes the connection? Ok(0) (EOF signal).
• Why is write_all necessary instead of just write? write might not write all bytes (short write), write_all loops until everything is written.

Why Milestone 1 Isn’t Enough → Moving to Milestone 2

Critical Limitation: The server handles one client at a time. While one client is typing, all other clients are blocked. In production, this is unacceptable:

• Scenario: Client A connects and goes idle (reading messages). Client B cannot connect at all.
• Scale: Single-threaded = 1 client max active, unusable for chat

What We’re Adding:

• Thread-per-connection model: Spawn a new thread for each client
• Concurrent clients: 10-1000 clients can connect simultaneously (limited by thread stack memory)
• Independence: Slow/idle clients don’t block others

Improvement:

• Concurrency: 1 client → ~1,000 concurrent clients (before hitting OS limits)
• Responsiveness: Fast client not blocked by slow client
• Cost: Each thread = ~2MB stack memory (1000 threads = 2GB RAM)
• Limitation: Still not production-scale (can’t handle 10K+ clients)

Why This Matters: Most servers in the 1990s-2000s used this model (Apache MPM prefork). It works for moderate loads but doesn’t scale to modern requirements (C10K problem).

Milestone 2: Multi-Threaded Echo Server

Introduction

The Problem with Milestone 1: One client blocks all others. We need concurrent handling.

The Solution: Spawn a thread for each connection using std::thread::spawn. Now each client gets independent execution—one slow client doesn’t affect others.

New Concepts:

• Thread-per-connection architecture
• Moving ownership into threads (move closure)
• Error handling across thread boundaries

Limitation Preview: Threads are expensive (2MB stack each). We can handle ~1,000 clients but not 10,000+. Milestone 3 will solve this with async I/O.

Key Concepts

New Patterns:

• thread::spawn(move || { ... }) - Move ownership of TcpStream into thread
• Each client handled in isolation
• Main thread only accepts connections, doesn’t handle them

Functions:

#![allow(unused)]
fn main() {
fn run_threaded_echo_server(addr: &str) -> io::Result<()>
    // Binds listener
    // For each connection: spawns thread with handle_client
    // Main thread continues accepting

// handle_client stays the same from Milestone 1
}

Checkpoint Tests

#![allow(unused)]
fn main() {
#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{BufRead, BufReader, Write};
    use std::net::TcpStream;
    use std::thread;
    use std::time::Duration;

    #[test]
    fn test_concurrent_clients() {
        // Start server
        thread::spawn(|| {
            run_threaded_echo_server("127.0.0.1:9004").unwrap();
        });
        thread::sleep(Duration::from_millis(100));

        // Connect 3 clients concurrently
        let mut clients: Vec<TcpStream> = (0..3)
            .map(|_| TcpStream::connect("127.0.0.1:9004").unwrap())
            .collect();

        // All should be able to send/receive simultaneously
        for (i, client) in clients.iter_mut().enumerate() {
            let msg = format!("Client {}\n", i);
            client.write_all(msg.as_bytes()).unwrap();

            let mut reader = BufReader::new(&*client);
            let mut line = String::new();
            reader.read_line(&mut line).unwrap();
            assert_eq!(line, msg);
        }
    }

    #[test]
    fn test_slow_client_doesnt_block() {
        thread::spawn(|| {
            run_threaded_echo_server("127.0.0.1:9005").unwrap();
        });
        thread::sleep(Duration::from_millis(100));

        // Client 1: connects but doesn't send (idle)
        let _slow_client = TcpStream::connect("127.0.0.1:9005").unwrap();

        // Client 2: should still work instantly
        let mut fast_client = TcpStream::connect("127.0.0.1:9005").unwrap();
        fast_client.write_all(b"Fast\n").unwrap();

        let mut buf = [0u8; 1024];
        let n = fast_client.read(&mut buf).unwrap();
        assert_eq!(&buf[..n], b"Fast\n");
    }

    #[test]
    fn test_many_connections() {
        thread::spawn(|| {
            run_threaded_echo_server("127.0.0.1:9006").unwrap();
        });
        thread::sleep(Duration::from_millis(100));

        // Spawn 50 client threads
        let handles: Vec<_> = (0..50)
            .map(|i| {
                thread::spawn(move || {
                    let mut stream = TcpStream::connect("127.0.0.1:9006").unwrap();
                    let msg = format!("Thread {}\n", i);
                    stream.write_all(msg.as_bytes()).unwrap();

                    let mut buf = [0u8; 1024];
                    let n = stream.read(&mut buf).unwrap();
                    assert_eq!(&buf[..n], msg.as_bytes());
                })
            })
            .collect();

        for h in handles {
            h.join().unwrap();
        }
    }
}
}

Starter Code

use std::io::{self, BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
use std::thread;

fn run_threaded_echo_server(addr: &str) -> io::Result<()> {
    let listener = TcpListener::bind(addr)?;
    println!("Multi-threaded echo server listening on {}", addr);

    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                let peer = stream.peer_addr()?;
                println!("New client: {}", peer);

                // TODO: Spawn a thread to handle this client
                // Use `move` to transfer ownership of `stream` into the thread
                // thread::spawn(move || {
                //     if let Err(e) = handle_client(stream) {
                //         eprintln!("Error with {}: {}", peer, e);
                //     }
                // });
                todo!();
            }
            Err(e) => {
                eprintln!("Connection failed: {}", e);
            }
        }
    }

    Ok(())
}

fn handle_client(stream: TcpStream) -> io::Result<()> {
    // Same as Milestone 1 - copy your implementation
    let reader = BufReader::new(&stream);
    let mut writer = &stream;
    let mut line = String::new();

    for line_result in reader.lines() {
        let line = line_result?;
        writer.write_all(line.as_bytes())?;
        writer.write_all(b"\n")?;
        writer.flush()?;
    }

    Ok(())
}

fn main() {
    if let Err(e) = run_threaded_echo_server("127.0.0.1:8080") {
        eprintln!("Server error: {}", e);
    }
}

Check Your Understanding

• Why do we need move in thread::spawn(move || ...)? Because the thread needs ownership of stream—it outlives the loop iteration.
• How many clients can this server handle? Limited by OS thread limits and memory (~1,000 threads = 2GB stack memory).
• What happens if a thread panics while handling a client? That client’s connection is dropped, but other clients and the server continue normally.
• Why is the main thread freed up now? It only accepts connections and spawns threads—it doesn’t wait for clients to finish.
• What’s the memory cost of 1,000 clients? ~2GB (2MB stack per thread) plus heap allocations.

Why Milestone 2 Isn’t Enough → Moving to Milestone 3

Limitation: Thread Scalability Crisis

• Thread-per-connection = 2MB stack/client
• 1,000 threads = 2GB memory (just stack!)
• 10,000 threads = 20GB memory + OS scheduler thrashing
• Modern servers need 100K+ concurrent connections (C10K problem)

What We’re Adding:

• Async I/O with tokio: One thread handles 10,000+ clients via async tasks
• Broadcast channel: Send messages to all connected clients
• Chat functionality: Transform echo server into actual chat

Improvement:

• Memory: 2MB/client → 2KB/client (1000x reduction)
• Scalability: 1,000 clients → 100,000+ clients on same hardware
• Features: Echo → broadcast chat (all clients see all messages)
• Architecture: Thread-per-connection → task-per-connection (async)

Performance Numbers:

• Threads: Context switch = 1-10μs, 1,000 threads max
• Async tasks: Await yield = 10-100ns, 100K+ tasks possible
• Real example: Discord handles 19M concurrent WebSocket connections (impossible with threads)

Milestone 3: Async TCP Chat with Broadcast

Introduction

The Async Revolution: Instead of blocking threads, we use async/await. When waiting for I/O (reading from socket), the task yields control to tokio’s runtime, which runs other tasks. This is cooperative multitasking—10,000 tasks share a few OS threads.

From Echo to Chat: Instead of echoing back to the sender, we broadcast each message to ALL connected clients. This is the core of a chat server.

Architecture:

• tokio::sync::broadcast channel: All clients subscribe to receive messages
• Each client has 2 tasks:
  • Reader task: Reads from socket, sends to broadcast channel
  • Writer task: Receives from broadcast channel, writes to socket
• Main loop: Accepts connections, spawns client handler tasks

Key Concepts

Structs/Types:

• tokio::net::TcpListener - Async version of std::net::TcpListener
• tokio::net::TcpStream - Async TCP connection
• tokio::sync::broadcast::Sender<String> - Broadcast channel for messages
• tokio::sync::broadcast::Receiver<String> - Subscriber to broadcast channel

Functions and Roles:

#![allow(unused)]
fn main() {
async fn run_chat_server(addr: &str) -> io::Result<()>
    // Creates broadcast channel
    // Binds TcpListener
    // Loops accepting connections
    // Spawns handle_client for each connection

async fn handle_client(
    stream: TcpStream,
    tx: broadcast::Sender<String>,
    addr: SocketAddr
)
    // Splits stream into read/write halves
    // Spawns reader task (reads lines, broadcasts)
    // Spawns writer task (receives broadcasts, writes)
    // Uses tokio::select! to cancel both when one ends
}

Key Pattern: Split read/write

┌─────────────┐
│  TcpStream  │
└──────┬──────┘
       │ split()
   ┌───┴───┐
   │       │
┌──▼─┐  ┌─▼───┐
│Read│  │Write│
└─┬──┘  └──┬──┘
  │        │
  ▼        ▼
Reader   Writer
 Task     Task

Checkpoint Tests

#![allow(unused)]
fn main() {
#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
    use tokio::net::TcpStream;
    use tokio::time::{sleep, Duration};

    #[tokio::test]
    async fn test_broadcast_to_all_clients() {
        // Start server
        tokio::spawn(async {
            run_chat_server("127.0.0.1:9007").await.unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        // Connect 3 clients
        let mut client1 = TcpStream::connect("127.0.0.1:9007").await.unwrap();
        let mut client2 = TcpStream::connect("127.0.0.1:9007").await.unwrap();
        let mut client3 = TcpStream::connect("127.0.0.1:9007").await.unwrap();

        // Client1 sends message
        client1.write_all(b"Hello from Client1\n").await.unwrap();

        // All clients (including sender) should receive it
        let mut reader1 = BufReader::new(&mut client1);
        let mut reader2 = BufReader::new(&mut client2);
        let mut reader3 = BufReader::new(&mut client3);

        let mut line = String::new();

        reader1.read_line(&mut line).await.unwrap();
        assert!(line.contains("Hello from Client1"));

        line.clear();
        reader2.read_line(&mut line).await.unwrap();
        assert!(line.contains("Hello from Client1"));

        line.clear();
        reader3.read_line(&mut line).await.unwrap();
        assert!(line.contains("Hello from Client1"));
    }

    #[tokio::test]
    async fn test_many_concurrent_clients() {
        tokio::spawn(async {
            run_chat_server("127.0.0.1:9008").await.unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        // Connect 100 clients concurrently
        let clients: Vec<_> = futures::future::join_all(
            (0..100).map(|_| TcpStream::connect("127.0.0.1:9008"))
        ).await;

        assert_eq!(clients.len(), 100);
        for client in clients {
            assert!(client.is_ok());
        }
    }

    #[tokio::test]
    async fn test_client_disconnect_graceful() {
        tokio::spawn(async {
            run_chat_server("127.0.0.1:9009").await.unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        let stream = TcpStream::connect("127.0.0.1:9009").await.unwrap();
        drop(stream); // Disconnect

        // Server should handle gracefully
        sleep(Duration::from_millis(100)).await;

        // New client can still connect
        let _new_client = TcpStream::connect("127.0.0.1:9009").await.unwrap();
    }
}
}

Starter Code

use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::broadcast;
use std::net::SocketAddr;

#[tokio::main]
async fn main() {
    if let Err(e) = run_chat_server("127.0.0.1:8080").await {
        eprintln!("Server error: {}", e);
    }
}

async fn run_chat_server(addr: &str) -> tokio::io::Result<()> {
    // TODO: Create broadcast channel with capacity 100
    let (tx, _rx) = todo!(); // broadcast::channel(100)

    // TODO: Bind TcpListener
    let listener = todo!(); // TcpListener::bind(addr).await?

    println!("Chat server listening on {}", addr);

    loop {
        // TODO: Accept connection (this is async now: .await)
        let (stream, addr) = todo!(); // listener.accept().await?

        println!("Client connected: {}", addr);

        // Clone sender for this client
        let tx = tx.clone();

        // TODO: Spawn async task to handle client
        tokio::spawn(async move {
            if let Err(e) = handle_client(stream, tx, addr).await {
                eprintln!("Error with {}: {}", addr, e);
            }
        });
    }
}

async fn handle_client(
    stream: TcpStream,
    tx: broadcast::Sender<String>,
    addr: SocketAddr,
) -> tokio::io::Result<()> {
    // TODO: Split stream into read and write halves
    let (reader, mut writer) = todo!(); // stream.into_split()
    let mut reader = BufReader::new(reader);

    // TODO: Subscribe to broadcast channel
    let mut rx = todo!(); // tx.subscribe()

    // Spawn task to read from client and broadcast
    let tx_clone = tx.clone();
    let addr_clone = addr;
    let mut read_task = tokio::spawn(async move {
        let mut line = String::new();
        loop {
            line.clear();
            // TODO: Read line from client (async)
            match reader.read_line(&mut line).await {
                Ok(0) => break, // EOF
                Ok(_) => {
                    // TODO: Broadcast message to all clients
                    let msg = format!("[{}] {}", addr_clone, line.trim());
                    // tx_clone.send(msg).ok();
                    todo!();
                }
                Err(e) => {
                    eprintln!("Read error: {}", e);
                    break;
                }
            }
        }
    });

    // Spawn task to receive broadcasts and write to client
    let mut write_task = tokio::spawn(async move {
        loop {
            // TODO: Receive message from broadcast channel
            match rx.recv().await {
                Ok(msg) => {
                    // TODO: Write message to client
                    // writer.write_all(msg.as_bytes()).await?
                    // writer.write_all(b"\n").await?
                    todo!();
                }
                Err(_) => break,
            }
        }
        Ok::<_, tokio::io::Error>(())
    });

    // TODO: Wait for either task to finish, then cancel the other
    tokio::select! {
        _ = &mut read_task => write_task.abort(),
        _ = &mut write_task => read_task.abort(),
    }

    println!("Client disconnected: {}", addr);
    Ok(())
}

Check Your Understanding

• Why split the stream into read/write halves? To have concurrent reading and writing—one task reads, one writes.
• What happens when a client sends a message? Reader task broadcasts it; all clients’ writer tasks receive and send to their sockets.
• Why does each client subscribe separately? Each needs their own receiver to get broadcast messages independently.
• What’s the difference between tokio::spawn and std::thread::spawn? tokio spawns async task (lightweight), thread spawns OS thread (heavyweight).
• Why use tokio::select!? To cancel both tasks when one finishes (e.g., client disconnects).
• Memory usage for 10,000 clients? ~20MB (2KB/task) vs 20GB with threads (2MB/thread).

Why Milestone 3 Isn’t Enough → Moving to Milestone 4

Limitation: No Isolation Between Conversations

• All clients in one giant room
• Can’t have private conversations or topic-specific channels
• No way to organize discussions (like Slack channels or Discord servers)
• Spam in one room affects everyone

What We’re Adding:

• Room-based architecture: Clients join specific rooms (e.g., “general”, “random”, “gaming”)
• Isolated broadcasts: Messages only go to clients in the same room
• Room management: JOIN, LEAVE, LIST commands
• Scalability: 1,000 rooms × 100 clients/room = better organization

Improvement:

• Organization: One global chat → multiple isolated rooms
• Privacy: Private rooms possible
• Scalability: Broadcast overhead reduced (send to 10 room members vs 10,000 global)
• Features: Like IRC channels or Slack workspaces

Real-World Pattern: Every chat platform uses rooms:

• Discord: Servers → Channels
• Slack: Workspaces → Channels
• IRC: Networks → Channels (#general, #help)

Milestone 4: Room-Based Architecture

Introduction

The Problem: Broadcasting to all clients doesn’t scale organizationally. Users need separate conversations.

The Solution:

• Store a HashMap<RoomId, Room> where each room has its own broadcast channel
• Clients join rooms with JOIN room_name
• Messages only broadcast within the current room
• New commands: JOIN, LEAVE, LIST (list rooms)

Architecture:

ChatServer
  └─ rooms: Arc<RwLock<HashMap<String, Room>>>
       ├─ "general": Room { tx: broadcast::Sender, users: HashSet }
       ├─ "gaming": Room { tx: broadcast::Sender, users: HashSet }
       └─ "random": Room { tx: broadcast::Sender, users: HashSet }

Key Concepts

Structs:

#![allow(unused)]
fn main() {
struct ChatServer {
    rooms: Arc<RwLock<HashMap<String, Room>>>,
}

struct Room {
    tx: broadcast::Sender<String>,
    users: HashSet<SocketAddr>,
}
}

Functions:

#![allow(unused)]
fn main() {
impl ChatServer {
    fn new() -> Self
        // Initialize with empty rooms HashMap

    async fn join_room(&self, room_id: String, user: SocketAddr)
        -> broadcast::Receiver<String>
        // Get or create room
        // Add user to room's user set
        // Return receiver for room's broadcast channel

    async fn leave_room(&self, room_id: &str, user: &SocketAddr)
        // Remove user from room
        // Delete room if empty (cleanup)

    async fn list_rooms(&self) -> Vec<String>
        // Return list of active room names
}
}

Protocol Commands:

• JOIN room_name - Join a room
• LEAVE - Leave current room
• LIST - List all rooms
• Any other text - Broadcast to current room

Checkpoint Tests

#![allow(unused)]
fn main() {
#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_room_isolation() {
        let server = ChatServer::new();

        let addr1 = "127.0.0.1:1111".parse().unwrap();
        let addr2 = "127.0.0.1:2222".parse().unwrap();

        // User 1 joins "general"
        let mut rx1 = server.join_room("general".to_string(), addr1).await;

        // User 2 joins "gaming"
        let mut rx2 = server.join_room("gaming".to_string(), addr2).await;

        // Get room references to send messages
        let rooms = server.rooms.read().await;
        let general = rooms.get("general").unwrap();
        let gaming = rooms.get("gaming").unwrap();

        // Send to "general"
        general.tx.send("Hello general".to_string()).ok();

        // User 1 receives, User 2 doesn't
        assert_eq!(rx1.try_recv().unwrap(), "Hello general");
        assert!(rx2.try_recv().is_err()); // No message in gaming room
    }

    #[tokio::test]
    async fn test_join_multiple_rooms() {
        let server = ChatServer::new();
        let addr = "127.0.0.1:3333".parse().unwrap();

        let _rx1 = server.join_room("general".to_string(), addr).await;
        let _rx2 = server.join_room("gaming".to_string(), addr).await;

        let rooms = server.list_rooms().await;
        assert_eq!(rooms.len(), 2);
        assert!(rooms.contains(&"general".to_string()));
        assert!(rooms.contains(&"gaming".to_string()));
    }

    #[tokio::test]
    async fn test_leave_room_cleanup() {
        let server = ChatServer::new();
        let addr = "127.0.0.1:4444".parse().unwrap();

        server.join_room("temp".to_string(), addr).await;
        assert_eq!(server.list_rooms().await.len(), 1);

        server.leave_room("temp", &addr).await;
        assert_eq!(server.list_rooms().await.len(), 0); // Room deleted when empty
    }

    #[tokio::test]
    async fn test_broadcast_within_room() {
        tokio::spawn(async {
            run_room_chat_server("127.0.0.1:9010").await.unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        let mut client1 = TcpStream::connect("127.0.0.1:9010").await.unwrap();
        let mut client2 = TcpStream::connect("127.0.0.1:9010").await.unwrap();
        let mut client3 = TcpStream::connect("127.0.0.1:9010").await.unwrap();

        // Client 1 and 2 join "general"
        client1.write_all(b"JOIN general\n").await.unwrap();
        client2.write_all(b"JOIN general\n").await.unwrap();

        // Client 3 joins "other"
        client3.write_all(b"JOIN other\n").await.unwrap();

        sleep(Duration::from_millis(50)).await;

        // Client 1 sends message
        client1.write_all(b"Hello room\n").await.unwrap();

        // Client 2 should receive, Client 3 should not
        let mut reader2 = BufReader::new(&mut client2);
        let mut reader3 = BufReader::new(&mut client3);

        let mut line2 = String::new();
        reader2.read_line(&mut line2).await.unwrap();
        assert!(line2.contains("Hello room"));

        // Client 3 shouldn't receive (different room)
        // (Test with timeout to avoid blocking)
    }
}
}

Starter Code

use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{broadcast, RwLock};

struct ChatServer {
    rooms: Arc<RwLock<HashMap<String, Room>>>,
}

struct Room {
    tx: broadcast::Sender<String>,
    users: HashSet<SocketAddr>,
}

impl ChatServer {
    fn new() -> Self {
        // TODO: Initialize ChatServer with empty rooms HashMap
        todo!()
    }

    async fn join_room(
        &self,
        room_id: String,
        user: SocketAddr,
    ) -> broadcast::Receiver<String> {
        // TODO: Acquire write lock on rooms
        let mut rooms = todo!(); // self.rooms.write().await

        // TODO: Get or create room
        let room = rooms.entry(room_id.clone()).or_insert_with(|| {
            // Create new room with broadcast channel
            todo!()
        });

        // TODO: Add user to room's user set
        todo!();

        // TODO: Return receiver for room's broadcast channel
        todo!()
    }

    async fn leave_room(&self, room_id: &str, user: &SocketAddr) {
        // TODO: Acquire write lock
        let mut rooms = todo!();

        // TODO: Remove user from room
        if let Some(room) = rooms.get_mut(room_id) {
            room.users.remove(user);

            // TODO: Delete room if empty
            if room.users.is_empty() {
                // rooms.remove(room_id);
                todo!();
            }
        }
    }

    async fn list_rooms(&self) -> Vec<String> {
        // TODO: Return list of room names
        todo!()
    }

    async fn broadcast_to_room(&self, room_id: &str, msg: String) {
        // TODO: Read lock, get room, send message
        let rooms = self.rooms.read().await;
        if let Some(room) = rooms.get(room_id) {
            // Ignore send errors (no receivers)
            let _ = room.tx.send(msg);
        }
    }
}

#[tokio::main]
async fn main() {
    if let Err(e) = run_room_chat_server("127.0.0.1:8080").await {
        eprintln!("Server error: {}", e);
    }
}

async fn run_room_chat_server(addr: &str) -> tokio::io::Result<()> {
    let server = Arc::new(ChatServer::new());
    let listener = TcpListener::bind(addr).await?;
    println!("Room-based chat server listening on {}", addr);

    loop {
        let (stream, addr) = listener.accept().await?;
        let server = server.clone();

        tokio::spawn(async move {
            if let Err(e) = handle_room_client(stream, server, addr).await {
                eprintln!("Error with {}: {}", addr, e);
            }
        });
    }
}

async fn handle_room_client(
    stream: TcpStream,
    server: Arc<ChatServer>,
    addr: SocketAddr,
) -> tokio::io::Result<()> {
    let (reader, mut writer) = stream.into_split();
    let mut reader = BufReader::new(reader);
    let mut line = String::new();

    // Track which room user is in
    let mut current_room: Option<(String, broadcast::Receiver<String>)> = None;

    writer.write_all(b"Welcome! Commands: JOIN <room>, LEAVE, LIST\n").await?;

    loop {
        line.clear();

        // TODO: Use tokio::select! to either:
        // 1. Read command from client
        // 2. Receive broadcast message from room (if in a room)

        tokio::select! {
            // Read from client
            result = reader.read_line(&mut line) => {
                match result {
                    Ok(0) => break, // EOF
                    Ok(_) => {
                        let trimmed = line.trim();

                        // TODO: Parse commands
                        if trimmed.starts_with("JOIN ") {
                            // Extract room name
                            // Leave current room if any
                            // Join new room
                            // Update current_room
                            todo!()
                        } else if trimmed == "LEAVE" {
                            // Leave current room
                            todo!()
                        } else if trimmed == "LIST" {
                            // List all rooms
                            todo!()
                        } else {
                            // Broadcast message to current room
                            if let Some((room_id, _)) = &current_room {
                                let msg = format!("[{}] {}", addr, trimmed);
                                server.broadcast_to_room(room_id, msg).await;
                            } else {
                                writer.write_all(b"Join a room first!\n").await?;
                            }
                        }
                    }
                    Err(e) => {
                        eprintln!("Read error: {}", e);
                        break;
                    }
                }
            }

            // Receive from broadcast (if in a room)
            msg = async {
                match &mut current_room {
                    Some((_, rx)) => rx.recv().await,
                    None => std::future::pending().await, // Never completes
                }
            } => {
                if let Ok(msg) = msg {
                    writer.write_all(msg.as_bytes()).await?;
                    writer.write_all(b"\n").await?;
                }
            }
        }
    }

    // Cleanup: leave room on disconnect
    if let Some((room_id, _)) = current_room {
        server.leave_room(&room_id, &addr).await;
    }

    Ok(())
}

Check Your Understanding

• Why use Arc<RwLock<HashMap>>? Arc for shared ownership across tasks, RwLock for concurrent read/write access.
• When do we use read lock vs write lock? Read for listing/broadcasting (many concurrent), write for join/leave (modify state).
• Why delete empty rooms? Memory cleanup—don’t keep rooms with no users.
• What’s the advantage of per-room broadcast channels? Scalability—broadcasting to 10 users instead of 10,000.
• How does tokio::select! help here? Concurrently wait for either: client input OR broadcast message from room.

Why Milestone 4 Isn’t Enough → Moving to Milestone 5

Limitation: TCP-Only Protocol

• Modern clients expect WebSocket (browsers can’t do raw TCP)
• No web-based clients possible
• Limited to terminal/native apps
• Missing the most common chat protocol today

What We’re Adding:

• WebSocket support: HTTP upgrade from axum server
• Multi-protocol: Same backend serves both TCP and WebSocket clients
• Unified architecture: Both protocols use same room system
• Browser compatibility: Can build web UI for chat

Improvement:

• Accessibility: Terminal clients → Web browsers + terminals
• Modern protocol: WebSocket is standard for real-time web apps
• Flexibility: Choose protocol based on client type
• Real-world: Discord uses WebSocket, Slack uses WebSocket, everyone uses WebSocket for web clients

Why This Matters: WebSocket is the de facto standard for browser-based real-time communication. Supporting it makes your chat server accessible to the widest audience.

Milestone 5: WebSocket Support (Multi-Protocol)

Introduction

The Gap: Our TCP chat works great for terminal clients, but modern users expect web-based chat (like Discord, Slack). Browsers can’t make raw TCP connections—they need WebSocket.

The Solution:

• Run an HTTP server (axum) alongside TCP server
• HTTP /ws endpoint upgrades to WebSocket
• WebSocket clients join the same rooms as TCP clients
• Both protocols share the same ChatServer backend

Architecture:

          ChatServer (shared)
           /          \
          /            \
    TCP Server    HTTP Server (axum)
         |              |
         |              └─ /ws → WebSocket
         |
    TCP Clients    WebSocket Clients

Key Concepts

New Dependencies:

[dependencies]
tokio = { version = "1", features = ["full"] }
axum = "0.7"
tower-http = { version = "0.5", features = ["cors"] }
futures-util = "0.3"

Structs/Types:

• axum::Router - HTTP route configuration
• axum::extract::ws::WebSocket - WebSocket connection
• axum::extract::ws::WebSocketUpgrade - HTTP upgrade request
• axum::extract::State<Arc<ChatServer>> - Shared server state

Functions:

#![allow(unused)]
fn main() {
async fn websocket_handler(
    ws: WebSocketUpgrade,
    State(server): State<Arc<ChatServer>>,
) -> impl IntoResponse
    // Handles HTTP upgrade to WebSocket
    // Returns response that upgrades connection

async fn handle_websocket_client(
    socket: WebSocket,
    server: Arc<ChatServer>,
)
    // Similar to handle_room_client but for WebSocket
    // Reads WebSocket messages (Message::Text)
    // Sends to room broadcast
}

Key Difference: WebSocket Messages:

• TCP: reader.read_line() → String
• WebSocket: socket.recv() → Message::Text(string) or Message::Close

Checkpoint Tests

#![allow(unused)]
fn main() {
#[cfg(test)]
mod tests {
    use super::*;
    use tokio_tungstenite::{connect_async, tungstenite::Message};
    use futures_util::{SinkExt, StreamExt};

    #[tokio::test]
    async fn test_websocket_connection() {
        // Start server
        tokio::spawn(async {
            run_multiprotocol_server("127.0.0.1:9011", "127.0.0.1:9111")
                .await
                .unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        // Connect via WebSocket
        let (ws_stream, _) = connect_async("ws://127.0.0.1:9111/ws")
            .await
            .unwrap();

        assert!(ws_stream.is_ok());
    }

    #[tokio::test]
    async fn test_tcp_and_websocket_interop() {
        tokio::spawn(async {
            run_multiprotocol_server("127.0.0.1:9012", "127.0.0.1:9112")
                .await
                .unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        // TCP client
        let mut tcp_client = TcpStream::connect("127.0.0.1:9012").await.unwrap();
        tcp_client.write_all(b"JOIN general\n").await.unwrap();

        // WebSocket client
        let (ws_stream, _) = connect_async("ws://127.0.0.1:9112/ws")
            .await
            .unwrap();
        let (mut ws_write, mut ws_read) = ws_stream.split();

        ws_write.send(Message::Text("JOIN general".to_string()))
            .await
            .unwrap();

        sleep(Duration::from_millis(50)).await;

        // TCP client sends message
        tcp_client.write_all(b"Hello from TCP\n").await.unwrap();

        // WebSocket client should receive it
        let msg = ws_read.next().await.unwrap().unwrap();
        assert!(matches!(msg, Message::Text(text) if text.contains("Hello from TCP")));
    }

    #[tokio::test]
    async fn test_websocket_commands() {
        tokio::spawn(async {
            run_multiprotocol_server("127.0.0.1:9013", "127.0.0.1:9113")
                .await
                .unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        let (ws_stream, _) = connect_async("ws://127.0.0.1:9113/ws")
            .await
            .unwrap();
        let (mut write, mut read) = ws_stream.split();

        // Send LIST command
        write.send(Message::Text("LIST".to_string())).await.unwrap();

        // Should receive room list
        let response = read.next().await.unwrap().unwrap();
        assert!(matches!(response, Message::Text(_)));
    }
}
}

Starter Code

use axum::{
    extract::{ws::WebSocket, ws::WebSocketUpgrade, State},
    response::IntoResponse,
    routing::get,
    Router,
};
use futures_util::{SinkExt, StreamExt};
use tokio::net::TcpListener;
use std::sync::Arc;

// ChatServer and Room structs from Milestone 4 (unchanged)

#[tokio::main]
async fn main() {
    run_multiprotocol_server("127.0.0.1:8080", "127.0.0.1:3000")
        .await
        .unwrap();
}

async fn run_multiprotocol_server(
    tcp_addr: &str,
    http_addr: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let server = Arc::new(ChatServer::new());

    // Spawn TCP server
    let tcp_server = server.clone();
    let tcp_addr = tcp_addr.to_string();
    tokio::spawn(async move {
        run_tcp_server(&tcp_addr, tcp_server).await
    });

    // Run HTTP/WebSocket server
    run_http_server(http_addr, server).await?;

    Ok(())
}

async fn run_tcp_server(
    addr: &str,
    server: Arc<ChatServer>,
) -> tokio::io::Result<()> {
    // Same as Milestone 4's run_room_chat_server
    let listener = TcpListener::bind(addr).await?;
    println!("TCP chat server listening on {}", addr);

    loop {
        let (stream, addr) = listener.accept().await?;
        let server = server.clone();

        tokio::spawn(async move {
            if let Err(e) = handle_room_client(stream, server, addr).await {
                eprintln!("TCP client error: {}", e);
            }
        });
    }
}

async fn run_http_server(
    addr: &str,
    server: Arc<ChatServer>,
) -> Result<(), Box<dyn std::error::Error>> {
    // TODO: Create axum Router with WebSocket route
    let app = Router::new()
        .route("/ws", get(websocket_handler))
        .with_state(server);

    println!("HTTP/WebSocket server listening on {}", addr);

    // TODO: Bind and serve
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}

async fn websocket_handler(
    ws: WebSocketUpgrade,
    State(server): State<Arc<ChatServer>>,
) -> impl IntoResponse {
    // TODO: Upgrade HTTP connection to WebSocket
    // ws.on_upgrade(|socket| handle_websocket_client(socket, server))
    todo!()
}

async fn handle_websocket_client(socket: WebSocket, server: Arc<ChatServer>) {
    // Generate fake address for WebSocket client
    use std::sync::atomic::{AtomicUsize, Ordering};
    static NEXT_ID: AtomicUsize = AtomicUsize::new(1);
    let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    let addr: SocketAddr = format!("0.0.0.0:{}", 10000 + id).parse().unwrap();

    // TODO: Split socket into sender/receiver
    let (mut sender, mut receiver) = todo!(); // socket.split()

    let mut current_room: Option<(String, broadcast::Receiver<String>)> = None;

    // Send welcome message
    sender
        .send(axum::extract::ws::Message::Text(
            "Welcome! Commands: JOIN <room>, LEAVE, LIST".to_string(),
        ))
        .await
        .ok();

    loop {
        tokio::select! {
            // Receive from WebSocket
            msg = receiver.next() => {
                match msg {
                    Some(Ok(axum::extract::ws::Message::Text(text))) => {
                        let trimmed = text.trim();

                        // TODO: Handle commands (same logic as TCP)
                        // JOIN <room>, LEAVE, LIST, or broadcast message

                        if trimmed.starts_with("JOIN ") {
                            // Parse room name, join room
                            todo!()
                        } else if trimmed == "LIST" {
                            // List rooms and send back
                            todo!()
                        } else if trimmed == "LEAVE" {
                            // Leave current room
                            todo!()
                        } else {
                            // Broadcast to room
                            if let Some((room_id, _)) = &current_room {
                                let msg = format!("[WS-{}] {}", id, trimmed);
                                server.broadcast_to_room(room_id, msg).await;
                            }
                        }
                    }
                    Some(Ok(axum::extract::ws::Message::Close(_))) | None => {
                        break;
                    }
                    _ => {}
                }
            }

            // Receive from room broadcast
            msg = async {
                match &mut current_room {
                    Some((_, rx)) => rx.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                if let Ok(msg) = msg {
                    // TODO: Send to WebSocket client
                    // sender.send(Message::Text(msg)).await.ok();
                    todo!();
                }
            }
        }
    }

    // Cleanup
    if let Some((room_id, _)) = current_room {
        server.leave_room(&room_id, &addr).await;
    }
}

// handle_room_client from Milestone 4 (unchanged)

Check Your Understanding

• Why do we need WebSocketUpgrade? WebSocket starts as HTTP GET request with special headers, then upgrades.
• What’s the difference between Message::Text and raw strings? WebSocket protocol has framing—messages can be text, binary, ping, pong, or close.
• Can TCP and WebSocket clients chat together? Yes! They both use the same ChatServer backend and rooms.
• Why generate a fake SocketAddr for WebSocket clients? The room system uses SocketAddr as user ID, WebSocket doesn’t have one from TCP layer.
• How does CORS affect WebSocket? If serving web UI from different origin, need CORS headers (tower-http).

Why Milestone 5 Isn’t Enough → Moving to Milestone 6

Limitation: Not Production-Ready

• No health monitoring: Can’t tell if server is healthy, degraded, or overloaded
• No graceful shutdown: Ctrl+C kills connections abruptly, loses messages
• No connection keepalive: Dead connections stay open (zombie clients)
• No observability: Can’t measure performance or debug issues

What We’re Adding:

• Metrics endpoint: Prometheus-style metrics (active connections, messages/sec, rooms)
• Graceful shutdown: Drain connections cleanly on SIGTERM/SIGINT
• Ping/Pong keepalive: Detect and close dead WebSocket connections
• Health check endpoint: /health for load balancers

Improvement:

• Observability: Blind operation → full metrics (Grafana dashboards possible)
• Reliability: Abrupt shutdown → graceful drain (zero lost messages)
• Resource cleanup: Zombie connections → automatic cleanup via keepalive
• Production-ready: Development toy → deployable service

Real-World Importance:

• Kubernetes: Needs /health endpoint for liveness/readiness probes
• Load balancers: Send traffic based on health checks
• Monitoring: Prometheus scrapes /metrics every 15s
• Graceful shutdown: Critical for zero-downtime deploys

Milestone 6: Production Features (Metrics, Graceful Shutdown, Keepalive)

Introduction

From Development to Production: The server works, but production requires:

1. Observability: What’s happening? (metrics)
2. Reliability: Clean shutdowns (graceful stop)
3. Resource management: Clean up dead connections (keepalive)

What We’re Adding:

• Metrics: /metrics endpoint exposing active connections, rooms, message rate
• Health check: /health endpoint (returns 200 OK if healthy)
• Graceful shutdown: CancellationToken to stop accepting new connections, drain existing
• WebSocket ping/pong: Periodic pings, close connection if no pong

Key Concepts

New Dependencies:

tokio-util = "0.7"

Structs:

#![allow(unused)]
fn main() {
struct Metrics {
    active_connections: AtomicUsize,
    total_messages: AtomicUsize,
    active_rooms: AtomicUsize,
}
}

Functions:

#![allow(unused)]
fn main() {
async fn metrics_handler(State(metrics): State<Arc<Metrics>>) -> String
    // Returns Prometheus-format metrics text
    // Example: "active_connections 42\ntotal_messages 1337\n"

async fn health_handler() -> &'static str
    // Returns "OK" with 200 status

async fn websocket_with_keepalive(socket: WebSocket, ...)
    // Spawns ping task (sends ping every 30s)
    // Main loop checks for pong responses
    // Closes connection if no pong received
}

Graceful Shutdown Pattern:

#![allow(unused)]
fn main() {
use tokio_util::sync::CancellationToken;

let token = CancellationToken::new();

// Spawn signal handler
tokio::spawn(async move {
    tokio::signal::ctrl_c().await.ok();
    token.cancel();
});

// In accept loop:
tokio::select! {
    result = listener.accept() => { /* handle */ }
    _ = token.cancelled() => {
        println!("Shutting down...");
        break;
    }
}
}

Checkpoint Tests

#![allow(unused)]
fn main() {
#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_metrics_endpoint() {
        tokio::spawn(async {
            run_production_server("127.0.0.1:9014", "127.0.0.1:9114")
                .await
                .unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        // Connect some clients
        let _client1 = TcpStream::connect("127.0.0.1:9014").await.unwrap();
        let _client2 = TcpStream::connect("127.0.0.1:9014").await.unwrap();

        // Query metrics
        let response = reqwest::get("http://127.0.0.1:9114/metrics")
            .await
            .unwrap()
            .text()
            .await
            .unwrap();

        assert!(response.contains("active_connections"));
        assert!(response.contains("2")); // 2 active connections
    }

    #[tokio::test]
    async fn test_health_endpoint() {
        tokio::spawn(async {
            run_production_server("127.0.0.1:9015", "127.0.0.1:9115")
                .await
                .unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        let response = reqwest::get("http://127.0.0.1:9115/health")
            .await
            .unwrap();

        assert_eq!(response.status(), 200);
        assert_eq!(response.text().await.unwrap(), "OK");
    }

    #[tokio::test]
    async fn test_websocket_keepalive() {
        tokio::spawn(async {
            run_production_server("127.0.0.1:9016", "127.0.0.1:9116")
                .await
                .unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        let (ws_stream, _) = connect_async("ws://127.0.0.1:9116/ws")
            .await
            .unwrap();
        let (mut write, mut read) = ws_stream.split();

        // Server should send ping after 30s (test with shorter interval in dev)
        // For testing, modify ping interval to 100ms
        sleep(Duration::from_millis(200)).await;

        // Should receive at least one ping
        let mut received_ping = false;
        while let Ok(Some(Ok(msg))) = tokio::time::timeout(
            Duration::from_millis(100),
            read.next()
        ).await {
            if matches!(msg, Message::Ping(_)) {
                received_ping = true;
                break;
            }
        }

        assert!(received_ping);
    }

    #[tokio::test]
    async fn test_graceful_shutdown() {
        let token = CancellationToken::new();
        let token_clone = token.clone();

        let server_handle = tokio::spawn(async move {
            run_production_server_with_token(
                "127.0.0.1:9017",
                "127.0.0.1:9117",
                token_clone,
            )
            .await
        });

        sleep(Duration::from_millis(100)).await;

        // Connect client
        let _client = TcpStream::connect("127.0.0.1:9017").await.unwrap();

        // Trigger shutdown
        token.cancel();

        // Server should stop accepting new connections
        sleep(Duration::from_millis(100)).await;

        // Server task should complete
        assert!(server_handle.is_finished() ||
                tokio::time::timeout(Duration::from_secs(2), server_handle)
                    .await
                    .is_ok());
    }
}
}

Starter Code

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::time::{interval, Duration};
use tokio_util::sync::CancellationToken;
use axum::{routing::get, Router, extract::State, response::IntoResponse};

struct Metrics {
    active_connections: AtomicUsize,
    total_messages: AtomicUsize,
    active_rooms: AtomicUsize,
}

impl Metrics {
    fn new() -> Self {
        // TODO: Initialize with zeros
        todo!()
    }

    fn connection_opened(&self) {
        // TODO: Increment active_connections
        todo!()
    }

    fn connection_closed(&self) {
        // TODO: Decrement active_connections
        todo!()
    }

    fn message_sent(&self) {
        // TODO: Increment total_messages
        todo!()
    }

    fn format_prometheus(&self) -> String {
        // TODO: Format as Prometheus metrics
        // Format: "metric_name value\n"
        format!(
            "active_connections {}\ntotal_messages {}\nactive_rooms {}\n",
            self.active_connections.load(Ordering::Relaxed),
            self.total_messages.load(Ordering::Relaxed),
            self.active_rooms.load(Ordering::Relaxed),
        )
    }
}

#[tokio::main]
async fn main() {
    let token = CancellationToken::new();
    let token_clone = token.clone();

    // TODO: Spawn signal handler
    tokio::spawn(async move {
        tokio::signal::ctrl_c().await.ok();
        println!("Shutdown signal received");
        token_clone.cancel();
    });

    if let Err(e) = run_production_server_with_token(
        "127.0.0.1:8080",
        "127.0.0.1:3000",
        token,
    )
    .await
    {
        eprintln!("Server error: {}", e);
    }
}

async fn run_production_server_with_token(
    tcp_addr: &str,
    http_addr: &str,
    shutdown_token: CancellationToken,
) -> Result<(), Box<dyn std::error::Error>> {
    let server = Arc::new(ChatServer::new());
    let metrics = Arc::new(Metrics::new());

    // Spawn TCP server with shutdown token
    let tcp_server = server.clone();
    let tcp_metrics = metrics.clone();
    let tcp_token = shutdown_token.clone();
    let tcp_addr = tcp_addr.to_string();

    tokio::spawn(async move {
        run_tcp_server_with_shutdown(&tcp_addr, tcp_server, tcp_metrics, tcp_token)
            .await
    });

    // Run HTTP server with metrics
    run_http_server_with_metrics(http_addr, server, metrics, shutdown_token).await?;

    Ok(())
}

async fn run_tcp_server_with_shutdown(
    addr: &str,
    server: Arc<ChatServer>,
    metrics: Arc<Metrics>,
    shutdown: CancellationToken,
) -> tokio::io::Result<()> {
    let listener = TcpListener::bind(addr).await?;
    println!("TCP server listening on {}", addr);

    loop {
        // TODO: Use tokio::select! to either accept connection or shutdown
        tokio::select! {
            result = listener.accept() => {
                match result {
                    Ok((stream, addr)) => {
                        metrics.connection_opened();

                        let server = server.clone();
                        let metrics = metrics.clone();

                        tokio::spawn(async move {
                            if let Err(e) = handle_room_client_with_metrics(
                                stream, server, addr, metrics.clone()
                            ).await {
                                eprintln!("Client error: {}", e);
                            }
                            metrics.connection_closed();
                        });
                    }
                    Err(e) => eprintln!("Accept error: {}", e),
                }
            }
            _ = shutdown.cancelled() => {
                println!("TCP server shutting down");
                break;
            }
        }
    }

    Ok(())
}

async fn run_http_server_with_metrics(
    addr: &str,
    server: Arc<ChatServer>,
    metrics: Arc<Metrics>,
    shutdown: CancellationToken,
) -> Result<(), Box<dyn std::error::Error>> {
    // TODO: Create router with health, metrics, and WebSocket endpoints
    let app = Router::new()
        .route("/health", get(health_handler))
        .route("/metrics", get(metrics_handler))
        .route("/ws", get(websocket_handler_with_keepalive))
        .with_state((server, metrics));

    println!("HTTP server listening on {}", addr);

    let listener = tokio::net::TcpListener::bind(addr).await?;

    // TODO: Serve with graceful shutdown
    axum::serve(listener, app)
        .with_graceful_shutdown(async move {
            shutdown.cancelled().await;
        })
        .await?;

    Ok(())
}

async fn health_handler() -> &'static str {
    // TODO: Return "OK"
    "OK"
}

async fn metrics_handler(
    State((_, metrics)): State<(Arc<ChatServer>, Arc<Metrics>)>,
) -> String {
    // TODO: Return Prometheus-formatted metrics
    metrics.format_prometheus()
}

async fn websocket_handler_with_keepalive(
    ws: WebSocketUpgrade,
    State((server, metrics)): State<(Arc<ChatServer>, Arc<Metrics>)>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| {
        handle_websocket_with_keepalive(socket, server, metrics)
    })
}

async fn handle_websocket_with_keepalive(
    socket: WebSocket,
    server: Arc<ChatServer>,
    metrics: Arc<Metrics>,
) {
    metrics.connection_opened();

    // TODO: Implement keepalive ping/pong
    // 1. Split socket
    // 2. Spawn ping task (send ping every 30s)
    // 3. Main loop: handle messages + check for pong
    // 4. Close if no pong received

    let (mut sender, mut receiver) = socket.split();

    // Spawn ping task
    let ping_task = tokio::spawn(async move {
        let mut interval = interval(Duration::from_secs(30));
        loop {
            interval.tick().await;
            if sender.send(axum::extract::ws::Message::Ping(vec![])).await.is_err() {
                break;
            }
        }
    });

    // Main message handling loop (similar to Milestone 5)
    // Add: track last_pong time, timeout if too old
    // ... (rest of WebSocket client handling) ...

    ping_task.abort();
    metrics.connection_closed();
}

// handle_room_client_with_metrics: same as handle_room_client but calls metrics.message_sent()

Check Your Understanding

• Why use AtomicUsize for metrics? Lock-free atomic operations—multiple threads can increment without locks.
• What’s the Prometheus format? Text format: metric_name value\n (e.g., active_connections 42\n)
• How does CancellationToken enable graceful shutdown? Tokio select! waits for either new connection or token.cancel(), breaks loop on shutdown.
• Why send ping messages? Detect dead connections (network failure, client crash) so server can clean up.
• What happens if client doesn’t respond to ping? After timeout (e.g., 60s no pong), close the connection.
• Why is graceful shutdown important? Kubernetes sends SIGTERM, waits 30s, then SIGKILL. Graceful shutdown drains connections cleanly.

Complete Working Example

Below is a fully functional multi-protocol chat server with all 6 milestones integrated:

use std::collections::{HashMap, HashSet};
use std::io::{self, BufRead, BufReader, Write};
use std::net::{SocketAddr, TcpListener as StdTcpListener, TcpStream as StdTcpStream};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

use axum::{
    extract::{ws::WebSocket, ws::WebSocketUpgrade, State},
    response::IntoResponse,
    routing::get,
    Router,
};
use futures_util::{SinkExt, StreamExt};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader as TokioBufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{broadcast, RwLock};
use tokio::time::interval;
use tokio_util::sync::CancellationToken;

//================================================
// Milestone 1: Simple Echo Server (Synchronous, Single-Threaded)
//================================================

fn run_echo_server(addr: &str) -> io::Result<()> {
    let listener = StdTcpListener::bind(addr)?;
    println!("Echo server listening on {}", addr);

    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                let peer = stream.peer_addr()?;
                println!("New client: {}", peer);

                if let Err(e) = handle_client_sync(stream) {
                    eprintln!("Error handling client: {}", e);
                }
            }
            Err(e) => {
                eprintln!("Connection failed: {}", e);
            }
        }
    }

    Ok(())
}

fn handle_client_sync(stream: StdTcpStream) -> io::Result<()> {
    let reader = BufReader::new(&stream);
    let mut writer = &stream;

    for line_result in reader.lines() {
        let line = line_result?;
        if line.is_empty() {
            continue;
        }
        writer.write_all(line.as_bytes())?;
        writer.write_all(b"\n")?;
        writer.flush()?;
    }

    Ok(())
}

//================================================
// Milestone 2: Multi-Threaded Echo Server
//================================================

fn run_threaded_echo_server(addr: &str) -> io::Result<()> {
    let listener = StdTcpListener::bind(addr)?;
    println!("Multi-threaded echo server listening on {}", addr);

    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                let peer = stream.peer_addr()?;
                println!("New client: {}", peer);

                thread::spawn(move || {
                    if let Err(e) = handle_client_sync(stream) {
                        eprintln!("Error with {}: {}", peer, e);
                    }
                });
            }
            Err(e) => {
                eprintln!("Connection failed: {}", e);
            }
        }
    }

    Ok(())
}

//================================================
// Milestone 3: Async TCP Chat with Broadcast
//================================================

async fn run_chat_server(addr: &str) -> tokio::io::Result<()> {
    let (tx, _rx) = broadcast::channel(100);
    let listener = TcpListener::bind(addr).await?;

    println!("Chat server listening on {}", addr);

    loop {
        let (stream, addr) = listener.accept().await?;
        println!("Client connected: {}", addr);

        let tx = tx.clone();

        tokio::spawn(async move {
            if let Err(e) = handle_chat_client(stream, tx, addr).await {
                eprintln!("Error with {}: {}", addr, e);
            }
        });
    }
}

async fn handle_chat_client(
    stream: TcpStream,
    tx: broadcast::Sender<String>,
    addr: SocketAddr,
) -> tokio::io::Result<()> {
    let (reader, mut writer) = stream.into_split();
    let mut reader = TokioBufReader::new(reader);
    let mut rx = tx.subscribe();

    let tx_clone = tx.clone();
    let addr_clone = addr;
    let mut read_task = tokio::spawn(async move {
        let mut line = String::new();
        loop {
            line.clear();
            match reader.read_line(&mut line).await {
                Ok(0) => break,
                Ok(_) => {
                    let msg = format!("[{}] {}", addr_clone, line.trim());
                    tx_clone.send(msg).ok();
                }
                Err(e) => {
                    eprintln!("Read error: {}", e);
                    break;
                }
            }
        }
    });

    let mut write_task = tokio::spawn(async move {
        loop {
            match rx.recv().await {
                Ok(msg) => {
                    if writer.write_all(msg.as_bytes()).await.is_err() {
                        break;
                    }
                    if writer.write_all(b"\n").await.is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
        Ok::<_, tokio::io::Error>(())
    });

    tokio::select! {
        _ = &mut read_task => write_task.abort(),
        _ = &mut write_task => read_task.abort(),
    }

    println!("Client disconnected: {}", addr);
    Ok(())
}

//================================================
// Milestone 4: Room-Based Architecture
//================================================

struct ChatServer {
    rooms: Arc<RwLock<HashMap<String, Room>>>,
}

struct Room {
    tx: broadcast::Sender<String>,
    users: HashSet<SocketAddr>,
}

impl ChatServer {
    fn new() -> Self {
        ChatServer {
            rooms: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    async fn join_room(
        &self,
        room_id: String,
        user: SocketAddr,
    ) -> broadcast::Receiver<String> {
        let mut rooms = self.rooms.write().await;

        let room = rooms.entry(room_id.clone()).or_insert_with(|| {
            let (tx, _) = broadcast::channel(100);
            Room {
                tx,
                users: HashSet::new(),
            }
        });

        room.users.insert(user);
        room.tx.subscribe()
    }

    async fn leave_room(&self, room_id: &str, user: &SocketAddr) {
        let mut rooms = self.rooms.write().await;

        if let Some(room) = rooms.get_mut(room_id) {
            room.users.remove(user);

            if room.users.is_empty() {
                rooms.remove(room_id);
            }
        }
    }

    async fn list_rooms(&self) -> Vec<String> {
        self.rooms.read().await.keys().cloned().collect()
    }

    async fn broadcast_to_room(&self, room_id: &str, msg: String) {
        let rooms = self.rooms.read().await;
        if let Some(room) = rooms.get(room_id) {
            let _ = room.tx.send(msg);
        }
    }
}

async fn run_room_chat_server(addr: &str) -> tokio::io::Result<()> {
    let server = Arc::new(ChatServer::new());
    let listener = TcpListener::bind(addr).await?;
    println!("Room-based chat server listening on {}", addr);

    loop {
        let (stream, addr) = listener.accept().await?;
        let server = server.clone();

        tokio::spawn(async move {
            if let Err(e) = handle_room_client(stream, server, addr).await {
                eprintln!("Error with {}: {}", addr, e);
            }
        });
    }
}

async fn handle_room_client(
    stream: TcpStream,
    server: Arc<ChatServer>,
    addr: SocketAddr,
) -> tokio::io::Result<()> {
    let (reader, mut writer) = stream.into_split();
    let mut reader = TokioBufReader::new(reader);
    let mut line = String::new();

    let mut current_room: Option<(String, broadcast::Receiver<String>)> = None;

    writer.write_all(b"Welcome! Commands: JOIN <room>, LEAVE, LIST\n").await?;

    loop {
        line.clear();

        tokio::select! {
            result = reader.read_line(&mut line) => {
                match result {
                    Ok(0) => break,
                    Ok(_) => {
                        let trimmed = line.trim();

                        if trimmed.starts_with("JOIN ") {
                            let room_name = trimmed[5..].trim().to_string();
                            if let Some((old_room, _)) = &current_room {
                                server.leave_room(old_room, &addr).await;
                            }
                            let rx = server.join_room(room_name.clone(), addr).await;
                            current_room = Some((room_name.clone(), rx));
                            writer.write_all(format!("Joined room: {}\n", room_name).as_bytes()).await?;
                        } else if trimmed == "LEAVE" {
                            if let Some((room_id, _)) = current_room.take() {
                                server.leave_room(&room_id, &addr).await;
                                writer.write_all(b"Left room\n").await?;
                            } else {
                                writer.write_all(b"Not in a room\n").await?;
                            }
                        } else if trimmed == "LIST" {
                            let rooms = server.list_rooms().await;
                            let list = if rooms.is_empty() {
                                "No rooms available\n".to_string()
                            } else {
                                format!("Rooms: {}\n", rooms.join(", "))
                            };
                            writer.write_all(list.as_bytes()).await?;
                        } else {
                            if let Some((room_id, _)) = &current_room {
                                let msg = format!("[{}] {}", addr, trimmed);
                                server.broadcast_to_room(room_id, msg).await;
                            } else {
                                writer.write_all(b"Join a room first!\n").await?;
                            }
                        }
                    }
                    Err(e) => {
                        eprintln!("Read error: {}", e);
                        break;
                    }
                }
            }

            msg = async {
                match &mut current_room {
                    Some((_, rx)) => rx.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                if let Ok(msg) = msg {
                    writer.write_all(msg.as_bytes()).await?;
                    writer.write_all(b"\n").await?;
                }
            }
        }
    }

    if let Some((room_id, _)) = current_room {
        server.leave_room(&room_id, &addr).await;
    }

    Ok(())
}

//================================================
// Milestone 5: WebSocket Support (Multi-Protocol)
//================================================

async fn run_multiprotocol_server(
    tcp_addr: &str,
    http_addr: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let server = Arc::new(ChatServer::new());

    let tcp_server = server.clone();
    let tcp_addr = tcp_addr.to_string();
    tokio::spawn(async move {
        run_tcp_server(&tcp_addr, tcp_server).await
    });

    run_http_server(http_addr, server).await?;

    Ok(())
}

async fn run_tcp_server(
    addr: &str,
    server: Arc<ChatServer>,
) -> tokio::io::Result<()> {
    let listener = TcpListener::bind(addr).await?;
    println!("TCP chat server listening on {}", addr);

    loop {
        let (stream, addr) = listener.accept().await?;
        let server = server.clone();

        tokio::spawn(async move {
            if let Err(e) = handle_room_client(stream, server, addr).await {
                eprintln!("TCP client error: {}", e);
            }
        });
    }
}

async fn run_http_server(
    addr: &str,
    server: Arc<ChatServer>,
) -> Result<(), Box<dyn std::error::Error>> {
    let app = Router::new()
        .route("/ws", get(websocket_handler))
        .with_state(server);

    println!("HTTP/WebSocket server listening on {}", addr);

    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}

async fn websocket_handler(
    ws: WebSocketUpgrade,
    State(server): State<Arc<ChatServer>>,
) -> impl IntoResponse {
    ws.on_upgrade(|socket| handle_websocket_client(socket, server))
}

async fn handle_websocket_client(socket: WebSocket, server: Arc<ChatServer>) {
    static NEXT_ID: AtomicUsize = AtomicUsize::new(1);
    let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    let addr: SocketAddr = format!("0.0.0.0:{}", 10000 + id).parse().unwrap();

    let (mut sender, mut receiver) = socket.split();

    let mut current_room: Option<(String, broadcast::Receiver<String>)> = None;

    let _ = sender
        .send(axum::extract::ws::Message::Text(
            "Welcome! Commands: JOIN <room>, LEAVE, LIST".to_string(),
        ))
        .await;

    loop {
        tokio::select! {
            msg = receiver.next() => {
                match msg {
                    Some(Ok(axum::extract::ws::Message::Text(text))) => {
                        let trimmed = text.trim();

                        if trimmed.starts_with("JOIN ") {
                            let room_name = trimmed[5..].trim().to_string();
                            if let Some((old_room, _)) = &current_room {
                                server.leave_room(old_room, &addr).await;
                            }
                            let rx = server.join_room(room_name.clone(), addr).await;
                            current_room = Some((room_name.clone(), rx));
                            let _ = sender.send(axum::extract::ws::Message::Text(
                                format!("Joined room: {}", room_name)
                            )).await;
                        } else if trimmed == "LIST" {
                            let rooms = server.list_rooms().await;
                            let list = if rooms.is_empty() {
                                "No rooms available".to_string()
                            } else {
                                format!("Rooms: {}", rooms.join(", "))
                            };
                            let _ = sender.send(axum::extract::ws::Message::Text(list)).await;
                        } else if trimmed == "LEAVE" {
                            if let Some((room_id, _)) = current_room.take() {
                                server.leave_room(&room_id, &addr).await;
                                let _ = sender.send(axum::extract::ws::Message::Text(
                                    "Left room".to_string()
                                )).await;
                            }
                        } else {
                            if let Some((room_id, _)) = &current_room {
                                let msg = format!("[WS-{}] {}", id, trimmed);
                                server.broadcast_to_room(room_id, msg).await;
                            }
                        }
                    }
                    Some(Ok(axum::extract::ws::Message::Close(_))) | None => {
                        break;
                    }
                    _ => {}
                }
            }

            msg = async {
                match &mut current_room {
                    Some((_, rx)) => rx.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                if let Ok(msg) = msg {
                    let _ = sender.send(axum::extract::ws::Message::Text(msg)).await;
                }
            }
        }
    }

    if let Some((room_id, _)) = current_room {
        server.leave_room(&room_id, &addr).await;
    }
}

//================================================
// Milestone 6: Production Features
//================================================

struct Metrics {
    active_connections: AtomicUsize,
    total_messages: AtomicUsize,
    active_rooms: AtomicUsize,
}

impl Metrics {
    fn new() -> Self {
        Metrics {
            active_connections: AtomicUsize::new(0),
            total_messages: AtomicUsize::new(0),
            active_rooms: AtomicUsize::new(0),
        }
    }

    fn connection_opened(&self) {
        self.active_connections.fetch_add(1, Ordering::Relaxed);
    }

    fn connection_closed(&self) {
        self.active_connections.fetch_sub(1, Ordering::Relaxed);
    }

    fn message_sent(&self) {
        self.total_messages.fetch_add(1, Ordering::Relaxed);
    }

    fn format_prometheus(&self) -> String {
        format!(
            "active_connections {}\ntotal_messages {}\nactive_rooms {}\n",
            self.active_connections.load(Ordering::Relaxed),
            self.total_messages.load(Ordering::Relaxed),
            self.active_rooms.load(Ordering::Relaxed),
        )
    }
}

async fn run_production_server_with_token(
    tcp_addr: &str,
    http_addr: &str,
    shutdown_token: CancellationToken,
) -> Result<(), Box<dyn std::error::Error>> {
    let server = Arc::new(ChatServer::new());
    let metrics = Arc::new(Metrics::new());

    let tcp_server = server.clone();
    let tcp_metrics = metrics.clone();
    let tcp_token = shutdown_token.clone();
    let tcp_addr = tcp_addr.to_string();

    tokio::spawn(async move {
        run_tcp_server_with_shutdown(&tcp_addr, tcp_server, tcp_metrics, tcp_token)
            .await
    });

    run_http_server_with_metrics(http_addr, server, metrics, shutdown_token).await?;

    Ok(())
}

async fn run_tcp_server_with_shutdown(
    addr: &str,
    server: Arc<ChatServer>,
    metrics: Arc<Metrics>,
    shutdown: CancellationToken,
) -> tokio::io::Result<()> {
    let listener = TcpListener::bind(addr).await?;
    println!("TCP server listening on {}", addr);

    loop {
        tokio::select! {
            result = listener.accept() => {
                match result {
                    Ok((stream, addr)) => {
                        metrics.connection_opened();

                        let server = server.clone();
                        let metrics = metrics.clone();

                        tokio::spawn(async move {
                            if let Err(e) = handle_room_client_with_metrics(
                                stream, server, addr, metrics.clone()
                            ).await {
                                eprintln!("Client error: {}", e);
                            }
                            metrics.connection_closed();
                        });
                    }
                    Err(e) => eprintln!("Accept error: {}", e),
                }
            }
            _ = shutdown.cancelled() => {
                println!("TCP server shutting down");
                break;
            }
        }
    }

    Ok(())
}

async fn handle_room_client_with_metrics(
    stream: TcpStream,
    server: Arc<ChatServer>,
    addr: SocketAddr,
    metrics: Arc<Metrics>,
) -> tokio::io::Result<()> {
    let (reader, mut writer) = stream.into_split();
    let mut reader = TokioBufReader::new(reader);
    let mut line = String::new();

    let mut current_room: Option<(String, broadcast::Receiver<String>)> = None;

    writer.write_all(b"Welcome! Commands: JOIN <room>, LEAVE, LIST\n").await?;

    loop {
        line.clear();

        tokio::select! {
            result = reader.read_line(&mut line) => {
                match result {
                    Ok(0) => break,
                    Ok(_) => {
                        let trimmed = line.trim();

                        if trimmed.starts_with("JOIN ") {
                            let room_name = trimmed[5..].trim().to_string();
                            if let Some((old_room, _)) = &current_room {
                                server.leave_room(old_room, &addr).await;
                            }
                            let rx = server.join_room(room_name.clone(), addr).await;
                            current_room = Some((room_name.clone(), rx));
                            writer.write_all(format!("Joined room: {}\n", room_name).as_bytes()).await?;
                        } else if trimmed == "LEAVE" {
                            if let Some((room_id, _)) = current_room.take() {
                                server.leave_room(&room_id, &addr).await;
                                writer.write_all(b"Left room\n").await?;
                            }
                        } else if trimmed == "LIST" {
                            let rooms = server.list_rooms().await;
                            let list = format!("Rooms: {}\n", rooms.join(", "));
                            writer.write_all(list.as_bytes()).await?;
                        } else {
                            if let Some((room_id, _)) = &current_room {
                                let msg = format!("[{}] {}", addr, trimmed);
                                server.broadcast_to_room(room_id, msg).await;
                                metrics.message_sent();
                            } else {
                                writer.write_all(b"Join a room first!\n").await?;
                            }
                        }
                    }
                    Err(e) => {
                        eprintln!("Read error: {}", e);
                        break;
                    }
                }
            }

            msg = async {
                match &mut current_room {
                    Some((_, rx)) => rx.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                if let Ok(msg) = msg {
                    writer.write_all(msg.as_bytes()).await?;
                    writer.write_all(b"\n").await?;
                }
            }
        }
    }

    if let Some((room_id, _)) = current_room {
        server.leave_room(&room_id, &addr).await;
    }

    Ok(())
}

async fn run_http_server_with_metrics(
    addr: &str,
    server: Arc<ChatServer>,
    metrics: Arc<Metrics>,
    shutdown: CancellationToken,
) -> Result<(), Box<dyn std::error::Error>> {
    let app = Router::new()
        .route("/health", get(health_handler))
        .route("/metrics", get(metrics_handler))
        .route("/ws", get(websocket_handler_with_keepalive))
        .with_state((server, metrics));

    println!("HTTP server listening on {}", addr);

    let listener = tokio::net::TcpListener::bind(addr).await?;

    axum::serve(listener, app)
        .with_graceful_shutdown(async move {
            shutdown.cancelled().await;
        })
        .await?;

    Ok(())
}

async fn health_handler() -> &'static str {
    "OK"
}

async fn metrics_handler(
    State((_, metrics)): State<(Arc<ChatServer>, Arc<Metrics>)>,
) -> String {
    metrics.format_prometheus()
}

async fn websocket_handler_with_keepalive(
    ws: WebSocketUpgrade,
    State((server, metrics)): State<(Arc<ChatServer>, Arc<Metrics>)>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| {
        handle_websocket_with_keepalive(socket, server, metrics)
    })
}

async fn handle_websocket_with_keepalive(
    socket: WebSocket,
    server: Arc<ChatServer>,
    metrics: Arc<Metrics>,
) {
    metrics.connection_opened();

    static NEXT_ID: AtomicUsize = AtomicUsize::new(1);
    let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    let addr: SocketAddr = format!("0.0.0.0:{}", 10000 + id).parse().unwrap();

    let (mut sender, mut receiver) = socket.split();

    let mut ping_interval = interval(Duration::from_secs(30));
    let mut current_room: Option<(String, broadcast::Receiver<String>)> = None;

    let _ = sender
        .send(axum::extract::ws::Message::Text(
            "Welcome! Commands: JOIN <room>, LEAVE, LIST".to_string(),
        ))
        .await;

    loop {
        tokio::select! {
            _ = ping_interval.tick() => {
                if sender.send(axum::extract::ws::Message::Ping(vec![])).await.is_err() {
                    break;
                }
            }

            msg = receiver.next() => {
                match msg {
                    Some(Ok(axum::extract::ws::Message::Text(text))) => {
                        let trimmed = text.trim();

                        if trimmed.starts_with("JOIN ") {
                            let room_name = trimmed[5..].trim().to_string();
                            if let Some((old_room, _)) = &current_room {
                                server.leave_room(old_room, &addr).await;
                            }
                            let rx = server.join_room(room_name.clone(), addr).await;
                            current_room = Some((room_name.clone(), rx));
                            let _ = sender.send(axum::extract::ws::Message::Text(
                                format!("Joined room: {}", room_name)
                            )).await;
                        } else if trimmed == "LIST" {
                            let rooms = server.list_rooms().await;
                            let list = format!("Rooms: {}", rooms.join(", "));
                            let _ = sender.send(axum::extract::ws::Message::Text(list)).await;
                        } else if trimmed == "LEAVE" {
                            if let Some((room_id, _)) = current_room.take() {
                                server.leave_room(&room_id, &addr).await;
                                let _ = sender.send(axum::extract::ws::Message::Text(
                                    "Left room".to_string()
                                )).await;
                            }
                        } else {
                            if let Some((room_id, _)) = &current_room {
                                let msg = format!("[WS-{}] {}", id, trimmed);
                                server.broadcast_to_room(room_id, msg).await;
                                metrics.message_sent();
                            }
                        }
                    }
                    Some(Ok(axum::extract::ws::Message::Pong(_))) => {
                        // Received pong response
                    }
                    Some(Ok(axum::extract::ws::Message::Close(_))) | None => {
                        break;
                    }
                    _ => {}
                }
            }

            msg = async {
                match &mut current_room {
                    Some((_, rx)) => rx.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                if let Ok(msg) = msg {
                    let _ = sender.send(axum::extract::ws::Message::Text(msg)).await;
                }
            }
        }
    }

    if let Some((room_id, _)) = current_room {
        server.leave_room(&room_id, &addr).await;
    }

    metrics.connection_closed();
}

//================================================
// Tests
//================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Read;
    use tokio::time::sleep;

    // Milestone 1 Tests
    #[test]
    fn test_echo_single_message() {
        thread::spawn(|| {
            run_echo_server("127.0.0.1:9001").unwrap();
        });
        thread::sleep(Duration::from_millis(100));

        let mut stream = StdTcpStream::connect("127.0.0.1:9001").unwrap();
        stream.write_all(b"Hello\n").unwrap();
        stream.flush().unwrap();

        // Give server time to process and respond
        thread::sleep(Duration::from_millis(50));

        let mut buf = [0u8; 1024];
        let n = stream.read(&mut buf).unwrap();
        assert_eq!(&buf[..n], b"Hello\n");
    }

    // Milestone 2 Tests
    #[test]
    fn test_concurrent_clients() {
        thread::spawn(|| {
            run_threaded_echo_server("127.0.0.1:9004").unwrap();
        });
        thread::sleep(Duration::from_millis(100));

        let mut clients: Vec<StdTcpStream> = (0..3)
            .map(|_| StdTcpStream::connect("127.0.0.1:9004").unwrap())
            .collect();

        for (i, client) in clients.iter_mut().enumerate() {
            let msg = format!("Client {}\n", i);
            client.write_all(msg.as_bytes()).unwrap();

            let mut reader = BufReader::new(&*client);
            let mut line = String::new();
            reader.read_line(&mut line).unwrap();
            assert_eq!(line, msg);
        }
    }

    // Milestone 3 Tests
    #[tokio::test]
    async fn test_broadcast_chat() {
        tokio::spawn(async {
            run_chat_server("127.0.0.1:9007").await.unwrap();
        });
        sleep(Duration::from_millis(100)).await;

        let mut client1 = TcpStream::connect("127.0.0.1:9007").await.unwrap();
        let mut client2 = TcpStream::connect("127.0.0.1:9007").await.unwrap();

        sleep(Duration::from_millis(50)).await;

        client1.write_all(b"Hello from client1\n").await.unwrap();

        let mut reader2 = TokioBufReader::new(&mut client2);
        let mut line = String::new();
        reader2.read_line(&mut line).await.unwrap();
        assert!(line.contains("Hello from client1"));
    }

    // Milestone 4 Tests
    #[tokio::test]
    async fn test_room_isolation() {
        let server = ChatServer::new();

        let addr1 = "127.0.0.1:1111".parse().unwrap();
        let addr2 = "127.0.0.1:2222".parse().unwrap();

        let mut rx1 = server.join_room("general".to_string(), addr1).await;
        let mut rx2 = server.join_room("gaming".to_string(), addr2).await;

        let rooms = server.rooms.read().await;
        let general = rooms.get("general").unwrap();

        general.tx.send("Hello general".to_string()).ok();

        assert_eq!(rx1.try_recv().unwrap(), "Hello general");
        assert!(rx2.try_recv().is_err());
    }

    #[tokio::test]
    async fn test_leave_room_cleanup() {
        let server = ChatServer::new();
        let addr = "127.0.0.1:4444".parse().unwrap();

        server.join_room("temp".to_string(), addr).await;
        assert_eq!(server.list_rooms().await.len(), 1);

        server.leave_room("temp", &addr).await;
        assert_eq!(server.list_rooms().await.len(), 0);
    }

    // Milestone 6 Tests
    #[test]
    fn test_metrics_initialization() {
        let metrics = Metrics::new();
        assert_eq!(metrics.active_connections.load(Ordering::Relaxed), 0);
        assert_eq!(metrics.total_messages.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn test_metrics_operations() {
        let metrics = Metrics::new();

        metrics.connection_opened();
        assert_eq!(metrics.active_connections.load(Ordering::Relaxed), 1);

        metrics.message_sent();
        assert_eq!(metrics.total_messages.load(Ordering::Relaxed), 1);

        metrics.connection_closed();
        assert_eq!(metrics.active_connections.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn test_prometheus_format() {
        let metrics = Metrics::new();
        metrics.connection_opened();
        metrics.message_sent();

        let output = metrics.format_prometheus();
        assert!(output.contains("active_connections 1"));
        assert!(output.contains("total_messages 1"));
    }
}

#[tokio::main]
async fn main() {
    println!("Multi-Protocol Chat Server - All Milestones");
    println!("===========================================");
    println!("Run with `cargo test --bin complete_26_network_chat_server` to test all milestones");
    println!("\nTo run production server with all features:");
    println!("  TCP server on 127.0.0.1:8080");
    println!("  HTTP/WebSocket server on 127.0.0.1:3000");
    println!("  Health check: http://127.0.0.1:3000/health");
    println!("  Metrics: http://127.0.0.1:3000/metrics");
    println!("  WebSocket: ws://127.0.0.1:3000/ws");

    let token = CancellationToken::new();
    let token_clone = token.clone();

    tokio::spawn(async move {
        tokio::signal::ctrl_c().await.ok();
        println!("\nShutdown signal received");
        token_clone.cancel();
    });

    if let Err(e) = run_production_server_with_token(
        "127.0.0.1:8080",
        "127.0.0.1:3000",
        token,
    )
    .await
    {
        eprintln!("Server error: {}", e);
    }

    println!("Server stopped gracefully");
}