13. 网络与线程池篇¶

use std::io::prelude::*;
use std::net::{TcpListener, TcpStream};
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use std::time::Duration;

// 任务类型：可执行一次的闭包，可跨线程，不借用外部（'static）
type Job = Box<dyn FnOnce() + Send + 'static>;

// 通道消息：区分"任务"和"终止信号"（优雅关闭）
enum Message {
    NewJob(Job),
    Terminate,   // 优雅关闭信号
}

struct ThreadPool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<Message>,   // 任务发送端
}

impl ThreadPool {
    // 创建线程池：size 个 Worker
    fn new(size: usize) -> ThreadPool {
        assert!(size > 0);   // 至少 1 个线程

        // 任务通道：主线程发送，Worker 接收
        let (sender, receiver) = mpsc::channel();
        // receiver 被所有 Worker 共享：Arc（多线程共享）+ Mutex（同一时刻一人接收）
        let receiver = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(size);
        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        ThreadPool { workers, sender }
    }

    // 提交任务：把闭包装箱成 Job 发送给空闲 Worker
    fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,   // 与 Job 相同的约束
    {
        let job = Box::new(f);
        self.sender.send(Message::NewJob(job)).unwrap();
    }
}

// Drop trait：ThreadPool 销毁时优雅关闭所有 Worker
impl Drop for ThreadPool {
    fn drop(&mut self) {
        // 给每个 Worker 发一个 Terminate
        for _ in &self.workers {
            self.sender.send(Message::Terminate).unwrap();
        }
        // 等待每个 Worker 退出
        // worker.thread 是 Option，take() 拿走 JoinHandle 才能 join（否则 E0507）
        for worker in &mut self.workers {
            println!("关闭 Worker {}", worker.id);
            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,   // Option 是为了 take()
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let thread = thread::spawn(move || {
            loop {
                // ⚠️ 关键：先 let 取出任务，锁在语句结束立即释放
                // 若写成 match receiver.lock().unwrap().recv() { ... }，
                // MutexGuard 会活到整个 match 结束，锁在 job() 执行期间一直被持有
                // → 其他 Worker 无法 recv → 所有请求串行！(见注意事项1)
                let job = receiver.lock().unwrap().recv();

                match job {
                    Ok(Message::NewJob(job)) => {
                        println!("Worker {} 执行任务", id);
                        job();   // 此刻锁已释放，其他 Worker 可并发
                    }
                    Ok(Message::Terminate) => {
                        println!("Worker {} 终止", id);
                        break;   // 退出循环，线程结束
                    }
                    Err(_) => break,   // 通道关闭也退出
                }
            }
        });
        Worker { id, thread: Some(thread) }
    }
}

// 处理单个 HTTP 连接：读请求、路由分发、写响应
fn handle_connection(mut stream: TcpStream) {
    let mut buffer = [0; 1024];   // 请求缓冲区（最多 1024 字节）
    stream.read(&mut buffer).unwrap();

    // 字节字符串字面量：buffer 是 [u8]，用 b"..." 匹配请求行
    let get = b"GET / HTTP/1.1";
    let sleep = b"GET /sleep HTTP/1.1";

    // 路由分发：返回 (状态行, 响应体)
    // HTML 带 meta charset="utf-8"，否则浏览器可能把中文显示成乱码
    let (status_line, contents) = if buffer.starts_with(get) {
        (
            "HTTP/1.1 200 OK",
            "<html lang=\"zh\"><head><meta charset=\"utf-8\"><body><h1>首页</h1></body></html>",
        )
    } else if buffer.starts_with(sleep) {
        thread::sleep(Duration::from_secs(5));   // 模拟慢请求
        (
            "HTTP/1.1 200 OK",
            "<html lang=\"zh\"><head><meta charset=\"utf-8\"><body><h1>休眠中...</h1></body></html>",
        )
    } else {
        (
            "HTTP/1.1 404 Not Found",
            "<html lang=\"zh\"><head><meta charset=\"utf-8\"><title>404 Not Found</title></head><body><h1>404 Not Found</h1></body></html>",
        )
    };

    // 拼 HTTP 响应：状态行 + Content-Length 头 + 空行 + 正文
    let response = format!(
        "{}\r\nContent-Length: {}\r\n\r\n{}",
        status_line,
        contents.len(),
        contents
    );

    stream.write(response.as_bytes()).unwrap();   // 写入响应
    stream.flush().unwrap();                      // 强制发送
}

fn main() {
    let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
    let pool = ThreadPool::new(4);   // 固定 4 个 Worker

    for stream in listener.incoming() {
        let stream = stream.unwrap();
        // move：stream 的所有权移入任务闭包（'static 要求）
        pool.execute(move || handle_connection(stream));
    }
    // main 结束时 pool 离开作用域 → Drop 触发 → 优雅关闭
}

// ❌ 错误：match scrutinee 中的 MutexGuard 活到整个 match 结束
//    锁在 job() 执行期间一直被持有 → 其他 Worker 无法 recv → 全部串行
match receiver.lock().unwrap().recv() {
    Ok(Message::NewJob(job)) => { job(); }   // 此刻锁还没释放！
    ...
}

// ✅ 正确：先 let 取出，锁在语句结束立即释放
let job = receiver.lock().unwrap().recv();
match job { ... }