src/libsync/task_pool.rs

   1 // Copyright 2012 The Rust Project Developers. See the COPYRIGHT
   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:Send(&T)),
  20     Quit
  21 }
  22
  23 pub struct TaskPool<T> {
  24     priv channels: Vec<Sender<Msg<T>>>,
  25     priv 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:Send(uint) -> 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:Send(&T)) {
  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:Send(uint) -> uint = || {
  86         let g: proc:Send(uint) -> 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 }