1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450
//! A simple crate for executing work on a thread pool, and getting back a //! future. //! //! This crate provides a simple thread pool abstraction for running work //! externally from the current thread that's running. An instance of `Future` //! is handed back to represent that the work may be done later, and further //! computations can be chained along with it as well. //! //! ```rust //! extern crate futures; //! extern crate futures_cpupool; //! //! use futures::Future; //! use futures_cpupool::CpuPool; //! //! # fn long_running_future(a: u32) -> Box<futures::future::Future<Item = u32, Error = ()> + Send> { //! # Box::new(futures::future::result(Ok(a))) //! # } //! # fn main() { //! //! // Create a worker thread pool with four threads //! let pool = CpuPool::new(4); //! //! // Execute some work on the thread pool, optionally closing over data. //! let a = pool.spawn(long_running_future(2)); //! let b = pool.spawn(long_running_future(100)); //! //! // Express some further computation once the work is completed on the thread //! // pool. //! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap(); //! //! // Print out the result //! println!("{:?}", c); //! # } //! ``` #![deny(missing_docs)] #![deny(missing_debug_implementations)] extern crate futures; extern crate num_cpus; use std::panic::{self, AssertUnwindSafe}; use std::sync::{Arc, Mutex}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::mpsc; use std::thread; use std::fmt; use futures::{IntoFuture, Future, Poll, Async}; use futures::future::{lazy, Executor, ExecuteError}; use futures::sync::oneshot::{channel, Sender, Receiver}; use futures::executor::{self, Run, Executor as OldExecutor}; /// A thread pool intended to run CPU intensive work. /// /// This thread pool will hand out futures representing the completed work /// that happens on the thread pool itself, and the futures can then be later /// composed with other work as part of an overall computation. /// /// The worker threads associated with a thread pool are kept alive so long as /// there is an open handle to the `CpuPool` or there is work running on them. Once /// all work has been drained and all references have gone away the worker /// threads will be shut down. /// /// Currently `CpuPool` implements `Clone` which just clones a new reference to /// the underlying thread pool. /// /// **Note:** if you use CpuPool inside a library it's better accept a /// `Builder` object for thread configuration rather than configuring just /// pool size. This not only future proof for other settings but also allows /// user to attach monitoring tools to lifecycle hooks. pub struct CpuPool { inner: Arc<Inner>, } /// Thread pool configuration object /// /// Builder starts with a number of workers equal to the number /// of CPUs on the host. But you can change it until you call `create()`. pub struct Builder { pool_size: usize, stack_size: usize, name_prefix: Option<String>, after_start: Option<Arc<Fn() + Send + Sync>>, before_stop: Option<Arc<Fn() + Send + Sync>>, } struct MySender<F, T> { fut: F, tx: Option<Sender<T>>, keep_running_flag: Arc<AtomicBool>, } trait AssertSendSync: Send + Sync {} impl AssertSendSync for CpuPool {} struct Inner { tx: Mutex<mpsc::Sender<Message>>, rx: Mutex<mpsc::Receiver<Message>>, cnt: AtomicUsize, size: usize, } impl fmt::Debug for CpuPool { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_struct("CpuPool") .field("size", &self.inner.size) .finish() } } impl fmt::Debug for Builder { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_struct("Builder") .field("pool_size", &self.pool_size) .field("name_prefix", &self.name_prefix) .finish() } } /// The type of future returned from the `CpuPool::spawn` function, which /// proxies the futures running on the thread pool. /// /// This future will resolve in the same way as the underlying future, and it /// will propagate panics. #[must_use] #[derive(Debug)] pub struct CpuFuture<T, E> { inner: Receiver<thread::Result<Result<T, E>>>, keep_running_flag: Arc<AtomicBool>, } enum Message { Run(Run), Close, } impl CpuPool { /// Creates a new thread pool with `size` worker threads associated with it. /// /// The returned handle can use `execute` to run work on this thread pool, /// and clones can be made of it to get multiple references to the same /// thread pool. /// /// This is a shortcut for: /// /// ```rust /// # use futures_cpupool::{Builder, CpuPool}; /// # /// # fn new(size: usize) -> CpuPool { /// Builder::new().pool_size(size).create() /// # } /// ``` /// /// # Panics /// /// Panics if `size == 0`. pub fn new(size: usize) -> CpuPool { Builder::new().pool_size(size).create() } /// Creates a new thread pool with a number of workers equal to the number /// of CPUs on the host. /// /// This is a shortcut for: /// /// ```rust /// # use futures_cpupool::{Builder, CpuPool}; /// # /// # fn new_num_cpus() -> CpuPool { /// Builder::new().create() /// # } /// ``` pub fn new_num_cpus() -> CpuPool { Builder::new().create() } /// Spawns a future to run on this thread pool, returning a future /// representing the produced value. /// /// This function will execute the future `f` on the associated thread /// pool, and return a future representing the finished computation. The /// returned future serves as a proxy to the computation that `F` is /// running. /// /// To simply run an arbitrary closure on a thread pool and extract the /// result, you can use the `future::lazy` combinator to defer work to /// executing on the thread pool itself. /// /// Note that if the future `f` panics it will be caught by default and the /// returned future will propagate the panic. That is, panics will not tear /// down the thread pool and will be propagated to the returned future's /// `poll` method if queried. /// /// If the returned future is dropped then this `CpuPool` will attempt to /// cancel the computation, if possible. That is, if the computation is in /// the middle of working, it will be interrupted when possible. pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error> where F: Future + Send + 'static, F::Item: Send + 'static, F::Error: Send + 'static, { let (tx, rx) = channel(); let keep_running_flag = Arc::new(AtomicBool::new(false)); // AssertUnwindSafe is used here because `Send + 'static` is basically // an alias for an implementation of the `UnwindSafe` trait but we can't // express that in the standard library right now. let sender = MySender { fut: AssertUnwindSafe(f).catch_unwind(), tx: Some(tx), keep_running_flag: keep_running_flag.clone(), }; executor::spawn(sender).execute(self.inner.clone()); CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() } } /// Spawns a closure on this thread pool. /// /// This function is a convenience wrapper around the `spawn` function above /// for running a closure wrapped in `future::lazy`. It will spawn the /// function `f` provided onto the thread pool, and continue to run the /// future returned by `f` on the thread pool as well. /// /// The returned future will be a handle to the result produced by the /// future that `f` returns. pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error> where F: FnOnce() -> R + Send + 'static, R: IntoFuture + 'static, R::Future: Send + 'static, R::Item: Send + 'static, R::Error: Send + 'static, { self.spawn(lazy(f)) } } impl<F> Executor<F> for CpuPool where F: Future<Item = (), Error = ()> + Send + 'static, { fn execute(&self, future: F) -> Result<(), ExecuteError<F>> { executor::spawn(future).execute(self.inner.clone()); Ok(()) } } impl Inner { fn send(&self, msg: Message) { self.tx.lock().unwrap().send(msg).unwrap(); } fn work(&self, after_start: Option<Arc<Fn() + Send + Sync>>, before_stop: Option<Arc<Fn() + Send + Sync>>) { after_start.map(|fun| fun()); loop { let msg = self.rx.lock().unwrap().recv().unwrap(); match msg { Message::Run(r) => r.run(), Message::Close => break, } } before_stop.map(|fun| fun()); } } impl Clone for CpuPool { fn clone(&self) -> CpuPool { self.inner.cnt.fetch_add(1, Ordering::Relaxed); CpuPool { inner: self.inner.clone() } } } impl Drop for CpuPool { fn drop(&mut self) { if self.inner.cnt.fetch_sub(1, Ordering::Relaxed) == 1 { for _ in 0..self.inner.size { self.inner.send(Message::Close); } } } } impl OldExecutor for Inner { fn execute(&self, run: Run) { self.send(Message::Run(run)) } } impl<T, E> CpuFuture<T, E> { /// Drop this future without canceling the underlying future. /// /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying /// future. This function can be used when user wants to drop but keep /// executing the underlying future. pub fn forget(self) { self.keep_running_flag.store(true, Ordering::SeqCst); } } impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> { type Item = T; type Error = E; fn poll(&mut self) -> Poll<T, E> { match self.inner.poll().expect("cannot poll CpuFuture twice") { Async::Ready(Ok(Ok(e))) => Ok(e.into()), Async::Ready(Ok(Err(e))) => Err(e), Async::Ready(Err(e)) => panic::resume_unwind(e), Async::NotReady => Ok(Async::NotReady), } } } impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> { type Item = (); type Error = (); fn poll(&mut self) -> Poll<(), ()> { if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() { if !self.keep_running_flag.load(Ordering::SeqCst) { // Cancelled, bail out return Ok(().into()) } } let res = match self.fut.poll() { Ok(Async::Ready(e)) => Ok(e), Ok(Async::NotReady) => return Ok(Async::NotReady), Err(e) => Err(e), }; // if the receiving end has gone away then that's ok, we just ignore the // send error here. drop(self.tx.take().unwrap().send(res)); Ok(Async::Ready(())) } } impl Builder { /// Create a builder a number of workers equal to the number /// of CPUs on the host. pub fn new() -> Builder { Builder { pool_size: num_cpus::get(), stack_size: 0, name_prefix: None, after_start: None, before_stop: None, } } /// Set size of a future CpuPool /// /// The size of a thread pool is the number of worker threads spawned pub fn pool_size(&mut self, size: usize) -> &mut Self { self.pool_size = size; self } /// Set stack size of threads in the pool. pub fn stack_size(&mut self, stack_size: usize) -> &mut Self { self.stack_size = stack_size; self } /// Set thread name prefix of a future CpuPool /// /// Thread name prefix is used for generating thread names. For example, if prefix is /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc. pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self { self.name_prefix = Some(name_prefix.into()); self } /// Execute function `f` right after each thread is started but before /// running any jobs on it. /// /// This is initially intended for bookkeeping and monitoring uses. /// The `f` will be deconstructed after the `builder` is deconstructed /// and all threads in the pool has executed it. pub fn after_start<F>(&mut self, f: F) -> &mut Self where F: Fn() + Send + Sync + 'static { self.after_start = Some(Arc::new(f)); self } /// Execute function `f` before each worker thread stops. /// /// This is initially intended for bookkeeping and monitoring uses. /// The `f` will be deconstructed after the `builder` is deconstructed /// and all threads in the pool has executed it. pub fn before_stop<F>(&mut self, f: F) -> &mut Self where F: Fn() + Send + Sync + 'static { self.before_stop = Some(Arc::new(f)); self } /// Create CpuPool with configured parameters /// /// # Panics /// /// Panics if `pool_size == 0`. pub fn create(&mut self) -> CpuPool { let (tx, rx) = mpsc::channel(); let pool = CpuPool { inner: Arc::new(Inner { tx: Mutex::new(tx), rx: Mutex::new(rx), cnt: AtomicUsize::new(1), size: self.pool_size, }), }; assert!(self.pool_size > 0); for counter in 0..self.pool_size { let inner = pool.inner.clone(); let after_start = self.after_start.clone(); let before_stop = self.before_stop.clone(); let mut thread_builder = thread::Builder::new(); if let Some(ref name_prefix) = self.name_prefix { thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter)); } if self.stack_size > 0 { thread_builder = thread_builder.stack_size(self.stack_size); } thread_builder.spawn(move || inner.work(after_start, before_stop)).unwrap(); } return pool } } #[cfg(test)] mod tests { use super::*; use std::sync::mpsc; #[test] fn test_drop_after_start() { let (tx, rx) = mpsc::sync_channel(2); let _cpu_pool = Builder::new() .pool_size(2) .after_start(move || tx.send(1).unwrap()).create(); // After Builder is deconstructed, the tx should be droped // so that we can use rx as an iterator. let count = rx.into_iter().count(); assert_eq!(count, 2); } }