Writing A Threadpool in Rust

文 Akisann@CNblogs / zhaihj@Github

本篇文章同时发布在Github上：https://zhaihj.github.io/writing-a-threadpool-in-rust.html

多线程一直是我相当不想碰的东西，总觉得看起来很棒，用起来却一点都不放心——尤其是过去用Delphi体验了多线程之后。实际上到了多线程里根本就没法定位那里出了错误，因此大部分时间压根不是在“调试”，而是告诉用户怎么用才能避免这个错误。在给OOC写MultiTheard Generating的时候我也紧紧用了最简单的ThreadPool，并且因为互斥锁多线程还没有单线程快。总之，这么多年我一直有意的躲避多线程的问题。

而Rust一直在宣传它在多线程上的优势——所有权如何的好，在Servo里面他们如何用Rust写了个漂亮的Parallel Parser。加上最近一年一直再用Rust写游戏的服务器，于是我决定回头看看Rust下的多线程体验到底如何。

ThreadPool

相比简单的用多线程算个加法，或者写个The Computer Language Benchmarks Game的测试程序来说，ThreadPool可能更适合练手。在这里，我会实现下面这个结构的ThreadPool:

简单来说，主线程通过Channel向ThreadPool发送Job，然后Threadpool里的每一个子线程在空闲时都会尝从Channel里获取Job，然后执行它。执行完毕之后，Job的返回结果会通过另一个Channel传送给主线程。实际上，在C或Delphi里，实现这么一个东西似乎并不是很困难，然而，Rust因为独特的所有权制度，代码比思路要复杂一些，下面，让我们来看看怎么实现这个ThreadPool。

Jobs for the Job

首先，让我们来设计Job，在这篇文章里，Job总是一个没有参数的函数，用代码来说，是这样：

type Job<T> = Box<Fn() -> T + Send + 'static>

由于我们需要把返回值传回给主线程，因此Job的返回值需要Send属性。这里，没有用FnOnce而是Fn仅仅是因为目前Rust的Box并不支持Box()，并且我并不打算用BoxFn来增加文章的复杂度。

对于每一个线程，Job大概是这样使用的：

loop {
    if let Ok(f) = get_a_job() {
        send_to_main_thread(f());
    }
}

我们每个线程都会不停的尝试获取job，执行它，然后获取下一个。相信到这里已经有不少人发现了问题：这个子线程似乎永远都无法结束。因为get_a_job显然只能返回Ok或阻塞——如果当Channel为空它就出错的话，一旦没有新任务，所有的子线程都会退出，之后这个ThreadPool就再也没法用了。于是，为了让这个子线程能退出，我们给Job添加一个用来表示结束的状态：

enum Message<T> {
    Work(Box<Fn() -> T + Send + 'static),
    Terminate,
}

type Job<T> = Message<T>;

这样，一旦Job是Terminate，线程就知道ThreadPool已经准备退出了。于是，我们可以把子线程写成这样：

loop {
    if let Ok(f) = get_a_job() {
        match f {
            Message::Work(f) => f(),
            Message::Terminate => break,
        }
    }
}

有了这个定义之后，让我们来看看ThreadPool该怎么设计。

从前一节的图片里，我们知道这个ThrealPolo需要两个Channel，一个用来接受新Job，另一个用来发送Job的返回结果，因此，ThrealPool可以写成这样：

struct ThreadPool <T>{
    sender: mpsc::Sender<Job<T>>,
    pub result: mpsc::Receiver<T>,
    threads: Vec<Option<thread::JoinHandle<()>>>,
}

这里，我们用了std::sync::mpsc，而threads则是用来储存所有子线程的，毕竟在结束的时候我们还要关闭他们的句柄。下面，我们尝试生成这个ThreadPool。由于子线程数是静态的，几乎所有的工作都可以在new里面完成，我们的大致思路是这样：

• 生成两个Channel, A和B，A用来接受Job，B用来发送结果
• 生成n个线程，每一个线程里：
• 保留A的Receiver和B的Sender
• 通过A的Receiver接受Job，执行，通过Sender发送

看起来并不是很难，让我们把它变成代码：

impl<T> ThreadPool<T> where T: Send + 'static { 

    ....

    // n是线程数，s是每一个线程获取新Job时的等待时间
    fn new(n: usize, s: u64) -> Self {
        // 用来接收Job的channel
        let (tx, rx) = mpsc::channel();
        // 用来发送结果的Channel
        let (tx1, rx1) = mpsc::channel();

        let rx : Arc<Mutex<mpsc::Receiver<Job<T>>>> = Arc::new(Mutex::new(rx));

首先我们定义了必要的channel，需要注意的是，对于channel，通常我们认为Sender可以有多个，但Receiver只有一个，因此在设计时Sender可以自然应对多线程，但Receiver则不。对于我们这里的情况，可以手动对Receiver加一个互斥锁，在Rust里可以使用std::sync::Mutex。

下面就是生成n个线程了：

let v = (0 .. n).map(| _ | {
    let rx = rx.clone();
    let tx1 : mpsc::Sender<T> = tx1.clone();
    Some(thread::spawn(move ||
        loop {
            if let Ok(f) = rx.lock() {
                if let Ok(f) = f.recv() {
                    match f {
                        Message::Work(f) => {
                            let r : T = f();
                            tx1.send(r);
                        },
                        Message::Terminate => break,
                    }
                }
            }
        }
    ))
}).collect::<Vec<_>>();

这里，thread::spawn用来生成新的线程，每个线程所作的事情就如之前描述的一样。最后，把这一切合起来，就可以得到一个ThreadPool了：

ThreadPool {
    sender: tx,
    result: rx1,
    threads: v,
}

有了这个ThreadPool之后，我们还需要一个方法来随时向里添加新Job，不过这件事情非常简单——只需要通过sender发送就够了：

fn add(&self, f: Job<T>) {
    self.sender.send(f).unwrap();
}

为了简单，我们并没有处理send返回的错误，在实际应用里，add可以返回Result。

Finishing and Drop

实际上，到此为止，大部分工作已经结束了，最后还有一个小事情：ThreadPool没法自己结束，因此我们需要手动实现这一部分。在其他语言里析构可能是理所当然的，不过在Rust里，除了C/C++的Wrapper之外，这种事情并不常见。Rust提供了Drop Trait来实现析构，Drop里的drop函数会在内存释放前自动执行。因此，我们只需要这么做：

impl<T> Drop for ThreadPool<T> {
    fn drop(&mut self) {
        for _ in &self.threads{
            self.sender.send(Message::Terminate);
        }
        for t in &mut self.threads {
            if let Some(t) = t.take() {
                t.join().unwrap();
            }
        }
    }
}

在这里，我们先对每一个线程发送结束命令，然后等待他们结束（join）。在实际应用里，线程可能会因为Job比较耗时而无法处理Terminate命令，Drop也会卡在Join的部分，不过可惜的是Rust的Thread并没有提供non-blocking的方法，因此你可能需要Future-rs来实现non-blocking的join。

Try It

到这里，主要部分已经完成了，下面我们来测试一下这个ThreadPool的效果如何：

fn main() {
    let tp = ThreadPool::new(10, 100);
    for i in 0 .. 100 {
        tp.add(
            Message::Work(
            Box::new(move || { println!("Thread: {}", i); i*100 })));
   }
   let mut c = 0;
   loop {
        if let Ok(s) = tp.result.recv() {
            println!("Result: {}", s);
            c += 1;
        }
        if c == 100 {
            break;
        }
    }
}

线程池有10个子线程，每个Job会输出一行文字，并返回一个数字，编译执行后，结果看起来像这样：

Thread: 0
Thread: 1
Result: 0
Result: 100
Thread: 2
Thread: 3
Result: 200
Result: 300
Thread: 4
Thread: 5
Thread: 6
Result: 400
Result: 500
Result: 600
......

Conclusion

在文章里，我尽量平白的描述这个过程，不过如果在没有基础的情况下自己去写的话，可能会遇到很多有趣的问题，比如：实际上thread.join会消耗掉JoinHandle的，因此，如果用Vec<JoinHandle>来储存线程句柄的话，会无法编译通过——因为drop是by ref的。在这篇文章里，我用了跟The Book同样的方法：添加一个Option来解决这个问题。不过，The Boox并不代表着最好的解决办法，比如我们还可以：

while let Some(e) = self.threads.pop() {
    e.join().unwrap();
}

对我来说，这个办法看起来更加简洁和优雅——但不知为什么The Book没有这么做。

本文的源代码可以在我的Github里找到。