rust 實戰 - 實現一個線程工作池 ThreadPool

如何實現一個線程池

線程池：一種線程使用模式。線程過多會帶來調度開銷，進而影響緩存局部性和整體性能。而線程池維護着多個線程，等待着監督管理者分配可並發執行的任務。這避免了在處理短時間任務時創建與銷毀線程的代價。線程池不僅能夠保證內核的充分利用，還能防止過分調度。可用線程數量應該取決於可用的並發處理器、處理器內核、內存、網絡sockets等的數量。 例如，對於計算密集型任務，線程數一般取cpu數量+2比較合適，線程數過多會導致額外的線程切換開銷。

如何定義線程池Pool呢，首先最大線程數量肯定要作為線程池的一個屬性，並且在new Pool時創建指定的線程。

線程池Pool

pub struct Pool {
  max_workers: usize, // 定義最大線程數
}

impl Pool {
  fn new(max_workers: usize) -> Pool {}
  fn execute<F>(&self, f:F) where F: FnOnce() + 'static + Send {}
}

用execute來執行任務，F: FnOnce() + 'static + Send 是使用thread::spawn線程執行需要滿足的trait, 代表F是一個能在線程里執行的閉包函數。

另一點自然而然會想到在Pool添加一個線程數組， 這個線程數組就是用來執行任務的。比如Vec<Thread> balabala。這里的線程是活的，是一個個不斷接受任務然后執行的實體。
可以看作在一個線程里不斷執行獲取任務並執行的Worker。

struct Worker where
{
    _id: usize, // worker 編號
}

要怎么把任務發送給Worker執行呢？mpsc(multi producer single consumer) 多生產者單消費者可以滿足我們的需求，let (tx, rx) = mpsc::channel() 可以獲取到一對發送端和接收端。
把發送端添加到Pool里面，把接收端添加到Worker里面。Pool通過channel將任務發送給多個worker消費執行。

這里有一點需要特別注意，channel的接收端receiver需要安全的在多個線程間共享，因此需要用Arc<Mutex::<T>>來包裹起來，也就是用鎖來解決並發沖突。

Pool的完整定義

pub struct Pool {
    workers: Vec<Worker>,
    max_workers: usize,
    sender: mpsc::Sender<Message>
}

該是時候定義我們要發給Worker的消息Message了
定義如下的枚舉值

type Job = Box<dyn FnOnce() + 'static + Send>;
enum Message {
    ByeBye,
    NewJob(Job),
}

Job是一個要發送給Worker執行的閉包函數，這里ByeBye用來通知Worker可以終止當前的執行，退出線程。

只剩下實現Worker和Pool的具體邏輯了。

Worker的實現

impl Worker
{
    fn new(id: usize, receiver: Arc::<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let t = thread::spawn( move || {
            loop {
                let receiver = receiver.lock().unwrap();
                let message=  receiver.recv().unwrap();
                match message {
                    Message::NewJob(job) => {
                        println!("do job from worker[{}]", id);
                        job();
                    },
                    Message::ByeBye => {
                        println!("ByeBye from worker[{}]", id);
                        break
                    },
                }  
            }
        });

        Worker {
            _id: id,
            t: Some(t),
        }
    }
}

let message = receiver.lock().unwrap().recv().unwrap(); 這里獲取鎖后從receiver獲取到消息體，然后let message結束后rust的生命周期會自動釋放掉鎖。
但如果寫成

while let message = receiver.lock().unwrap().recv().unwrap() {
};

while let 后面整個括號都是一個作用域，要在這個作用域結束后，鎖才會釋放，比上面let message要鎖定久時間。
rust的mutex鎖沒有對應的unlock方法，由mutex的生命周期管理。

我們給Pool實現Drop trait， 讓Pool被銷毀時，自動暫停掉worker線程的執行。

impl Drop for Pool {
    fn drop(&mut self) {
        for _ in 0..self.max_workers {
            self.sender.send(Message::ByeBye).unwrap();
        }
        for w in self.workers.iter_mut() {
            if let Some(t) = w.t.take() {
                t.join().unwrap();
            }
        }
    }
}

drop方法里面用了兩個循環，而不是在一個循環里做完兩件事?

for w in self.workers.iter_mut() {
    if let Some(t) = w.t.take() {
        self.sender.send(Message::ByeBye).unwrap();
        t.join().unwrap();
    }
}

這里面隱藏了一個會造成死鎖的陷阱，比如兩個Worker, 在單個循環里面迭代所有Worker，再將終止信息發送給通道后，直接調用join，
我們預期是第一個worker要收到消息，並且等他執行完。當情況可能是第二個worker獲取到了消息，第一個worker沒有獲取到，那接下來的join就會阻塞造成死鎖。

注意到沒有，Worker是被包裝在Option內的，這里有兩個點需要注意

1. t.join 需要持有t的所有權
2. 在我們這種情況下，self.workers只能作為引用被for循環迭代。

這里考慮讓Worker持有Option<JoinHandle<()>>，后續可以通過在Option上調用take方法將Some變體的值移出來，並在原來的位置留下None變體。
換而言之，讓運行中的worker持有Some的變體，清理worker時，可以使用None替換掉Some，從而讓Worker失去可以運行的線程

struct Worker where
{
    _id: usize,
    t: Option<JoinHandle<()>>,
}

要點總結

• Mutex依賴於生命周期管理鎖的釋放，使用的時候需要注意是否逾期持有鎖
• Vec<Option<T>> 可以解決某些情況下需要T所有權的場景

完整代碼

use std::thread::{self, JoinHandle};
use std::sync::{Arc, mpsc, Mutex};


type Job = Box<dyn FnOnce() + 'static + Send>;
enum Message {
    ByeBye,
    NewJob(Job),
}

struct Worker where
{
    _id: usize,
    t: Option<JoinHandle<()>>,
}

impl Worker
{
    fn new(id: usize, receiver: Arc::<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let t = thread::spawn( move || {
            loop {
                let message = receiver.lock().unwrap().recv().unwrap();
                match message {
                    Message::NewJob(job) => {
                        println!("do job from worker[{}]", id);
                        job();
                    },
                    Message::ByeBye => {
                        println!("ByeBye from worker[{}]", id);
                        break
                    },
                }  
            }
        });

        Worker {
            _id: id,
            t: Some(t),
        }
    }
}

pub struct Pool {
    workers: Vec<Worker>,
    max_workers: usize,
    sender: mpsc::Sender<Message>
}

impl Pool where {
    pub fn new(max_workers: usize) -> Pool {
        if max_workers == 0 {
            panic!("max_workers must be greater than zero!")
        }
        let (tx, rx) = mpsc::channel();

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

        Pool { workers: workers, max_workers: max_workers, sender: tx }
    }
    
    pub fn execute<F>(&self, f:F) where F: FnOnce() + 'static + Send
    {

        let job = Message::NewJob(Box::new(f));
        self.sender.send(job).unwrap();
    }
}

impl Drop for Pool {
    fn drop(&mut self) {
        for _ in 0..self.max_workers {
            self.sender.send(Message::ByeBye).unwrap();
        }
        for w in self.workers.iter_mut() {
            if let Some(t) = w.t.take() {
                t.join().unwrap();
            }
        }
    }
}


#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn it_works() {
        let p = Pool::new(4);
        p.execute(|| println!("do new job1"));
        p.execute(|| println!("do new job2"));
        p.execute(|| println!("do new job3"));
        p.execute(|| println!("do new job4"));
    }
}