(index<- ) ./libsync/task_pool.rs
git branch: * master 5200215 auto merge of #14035 : alexcrichton/rust/experimental, r=huonw
modified: Wed Apr 9 17:27:03 2014
2 // file at the top-level directory of this distribution and at
3 // http://rust-lang.org/COPYRIGHT.
4 //
5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8 // option. This file may not be copied, modified, or distributed
9 // except according to those terms.
10
11 #![allow(missing_doc)]
12
13 /// A task pool abstraction. Useful for achieving predictable CPU
14 /// parallelism.
15
16 use std::task;
17
18 enum Msg<T> {
19 Execute(proc(&T):Send),
20 Quit
21 }
22
23 pub struct TaskPool<T> {
24 channels: Vec<Sender<Msg<T>>>,
25 next_index: uint,
26 }
27
28 #[unsafe_destructor]
29 impl<T> Drop for TaskPool<T> {
30 fn drop(&mut self) {
31 for channel in self.channels.mut_iter() {
32 channel.send(Quit);
33 }
34 }
35 }
36
37 impl<T> TaskPool<T> {
38 /// Spawns a new task pool with `n_tasks` tasks. If the `sched_mode`
39 /// is None, the tasks run on this scheduler; otherwise, they run on a
40 /// new scheduler with the given mode. The provided `init_fn_factory`
41 /// returns a function which, given the index of the task, should return
42 /// local data to be kept around in that task.
43 pub fn new(n_tasks: uint,
44 init_fn_factory: || -> proc(uint):Send -> T)
45 -> TaskPool<T> {
46 assert!(n_tasks >= 1);
47
48 let channels = Vec::from_fn(n_tasks, |i| {
49 let (tx, rx) = channel::<Msg<T>>();
50 let init_fn = init_fn_factory();
51
52 let task_body = proc() {
53 let local_data = init_fn(i);
54 loop {
55 match rx.recv() {
56 Execute(f) => f(&local_data),
57 Quit => break
58 }
59 }
60 };
61
62 // Run on this scheduler.
63 task::spawn(task_body);
64
65 tx
66 });
67
68 return TaskPool {
69 channels: channels,
70 next_index: 0,
71 };
72 }
73
74 /// Executes the function `f` on a task in the pool. The function
75 /// receives a reference to the local data returned by the `init_fn`.
76 pub fn execute(&mut self, f: proc(&T):Send) {
77 self.channels.get(self.next_index).send(Execute(f));
78 self.next_index += 1;
79 if self.next_index == self.channels.len() { self.next_index = 0; }
80 }
81 }
82
83 #[test]
84 fn test_task_pool() {
85 let f: || -> proc(uint):Send -> uint = || {
86 let g: proc(uint):Send -> uint = proc(i) i;
87 g
88 };
89 let mut pool = TaskPool::new(4, f);
90 for _ in range(0, 8) {
91 pool.execute(proc(i) println!("Hello from thread {}!", *i));
92 }
93 }