mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-09-03 00:00:05 +02:00
Threadpool blocking (#317)
This patch adds a `blocking` to `tokio-threadpool`. This function serves as a way to annotate sections of code that will perform blocking operations. This informs the thread pool that an additional thread needs to be spawned to replace the current thread, which will no longer be able to process the work queue.
This commit is contained in:
@@ -0,0 +1,410 @@
|
||||
extern crate tokio_threadpool;
|
||||
|
||||
extern crate env_logger;
|
||||
#[macro_use]
|
||||
extern crate futures;
|
||||
extern crate rand;
|
||||
|
||||
use tokio_threadpool::*;
|
||||
|
||||
use futures::*;
|
||||
use futures::future::{lazy, poll_fn};
|
||||
use rand::*;
|
||||
|
||||
use std::sync::*;
|
||||
use std::sync::atomic::*;
|
||||
use std::sync::atomic::Ordering::*;
|
||||
use std::time::Duration;
|
||||
use std::thread;
|
||||
|
||||
#[test]
|
||||
fn basic() {
|
||||
let _ = ::env_logger::init();
|
||||
|
||||
let pool = Builder::new()
|
||||
.pool_size(1)
|
||||
.max_blocking(1)
|
||||
.build();
|
||||
|
||||
let (tx1, rx1) = mpsc::channel();
|
||||
let (tx2, rx2) = mpsc::channel();
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
let res = blocking(|| {
|
||||
let v = rx1.recv().unwrap();
|
||||
tx2.send(v).unwrap();
|
||||
}).unwrap();
|
||||
|
||||
assert!(res.is_ready());
|
||||
Ok(().into())
|
||||
}));
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
tx1.send(()).unwrap();
|
||||
Ok(().into())
|
||||
}));
|
||||
|
||||
rx2.recv().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notify_task_on_capacity() {
|
||||
const BLOCKING: usize = 10;
|
||||
|
||||
let pool = Builder::new()
|
||||
.pool_size(1)
|
||||
.max_blocking(1)
|
||||
.build();
|
||||
|
||||
let rem = Arc::new(AtomicUsize::new(BLOCKING));
|
||||
let (tx, rx) = mpsc::channel();
|
||||
|
||||
for _ in 0..BLOCKING {
|
||||
let rem = rem.clone();
|
||||
let tx = tx.clone();
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
poll_fn(move || {
|
||||
blocking(|| {
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
let prev = rem.fetch_sub(1, Relaxed);
|
||||
|
||||
if prev == 1 {
|
||||
tx.send(()).unwrap();
|
||||
}
|
||||
}).map_err(|e| panic!("blocking err {:?}", e))
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
rx.recv().unwrap();
|
||||
|
||||
assert_eq!(0, rem.load(Relaxed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_is_use_it_or_lose_it() {
|
||||
use futures::*;
|
||||
use futures::Async::*;
|
||||
use futures::sync::oneshot;
|
||||
use futures::task::Task;
|
||||
|
||||
// TODO: Run w/ bigger pool size
|
||||
|
||||
let pool = Builder::new()
|
||||
.pool_size(1)
|
||||
.max_blocking(1)
|
||||
.build();
|
||||
|
||||
let (tx1, rx1) = mpsc::channel();
|
||||
let (tx2, rx2) = oneshot::channel();
|
||||
let (tx3, rx3) = mpsc::channel();
|
||||
let (tx4, rx4) = mpsc::channel();
|
||||
|
||||
// First, fill the blocking capacity
|
||||
pool.spawn(lazy(move || {
|
||||
poll_fn(move || {
|
||||
blocking(|| {
|
||||
rx1.recv().unwrap();
|
||||
}).map_err(|_| panic!())
|
||||
})
|
||||
}));
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
rx2
|
||||
.map_err(|_| panic!())
|
||||
.and_then(|task: Task| {
|
||||
poll_fn(move || {
|
||||
blocking(|| {
|
||||
// Notify the other task
|
||||
task.notify();
|
||||
|
||||
// Block until woken
|
||||
rx3.recv().unwrap();
|
||||
}).map_err(|_| panic!())
|
||||
})
|
||||
})
|
||||
}));
|
||||
|
||||
// Spawn a future that will try to block, get notified, then not actually
|
||||
// use the blocking
|
||||
let mut i = 0;
|
||||
let mut tx2 = Some(tx2);
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
poll_fn(move || {
|
||||
match i {
|
||||
0 => {
|
||||
i = 1;
|
||||
|
||||
let res = blocking(|| unreachable!())
|
||||
.map_err(|_| panic!());
|
||||
|
||||
assert!(res.unwrap().is_not_ready());
|
||||
|
||||
// Unblock the first blocker
|
||||
tx1.send(()).unwrap();
|
||||
|
||||
return Ok(NotReady);
|
||||
}
|
||||
1 => {
|
||||
i = 2;
|
||||
|
||||
// Skip blocking, and notify the second task that it should
|
||||
// start blocking
|
||||
let me = task::current();
|
||||
tx2.take().unwrap().send(me).unwrap();
|
||||
|
||||
return Ok(NotReady);
|
||||
}
|
||||
2 => {
|
||||
let res = blocking(|| unreachable!())
|
||||
.map_err(|_| panic!());
|
||||
|
||||
assert!(res.unwrap().is_not_ready());
|
||||
|
||||
// Unblock the first blocker
|
||||
tx3.send(()).unwrap();
|
||||
tx4.send(()).unwrap();
|
||||
Ok(().into())
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
})
|
||||
}));
|
||||
|
||||
rx4.recv().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocking_thread_does_not_take_over_shutdown_worker_thread() {
|
||||
let pool = Builder::new()
|
||||
.pool_size(2)
|
||||
.max_blocking(1)
|
||||
.build();
|
||||
|
||||
let (enter_tx, enter_rx) = mpsc::channel();
|
||||
let (exit_tx, exit_rx) = mpsc::channel();
|
||||
let (try_tx, try_rx) = mpsc::channel();
|
||||
|
||||
let exited = Arc::new(AtomicBool::new(false));
|
||||
|
||||
{
|
||||
let exited = exited.clone();
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
poll_fn(move || {
|
||||
blocking(|| {
|
||||
enter_tx.send(()).unwrap();
|
||||
exit_rx.recv().unwrap();
|
||||
exited.store(true, Relaxed);
|
||||
}).map_err(|_| panic!())
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
// Wait for the task to block
|
||||
let _ = enter_rx.recv().unwrap();
|
||||
|
||||
// Spawn another task that attempts to block
|
||||
pool.spawn(lazy(move || {
|
||||
poll_fn(move || {
|
||||
let res = blocking(|| {
|
||||
|
||||
}).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
res.is_ready(),
|
||||
exited.load(Relaxed));
|
||||
|
||||
try_tx.send(res.is_ready()).unwrap();
|
||||
|
||||
Ok(res)
|
||||
})
|
||||
}));
|
||||
|
||||
// Wait for the second task to try to block (and not be ready).
|
||||
let res = try_rx.recv().unwrap();
|
||||
assert!(!res);
|
||||
|
||||
// Unblock the first task
|
||||
exit_tx.send(()).unwrap();
|
||||
|
||||
// Wait for the second task to successfully block.
|
||||
let res = try_rx.recv().unwrap();
|
||||
assert!(res);
|
||||
|
||||
drop(pool);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blockin_one_time_gets_capacity_for_multiple_blocks() {
|
||||
const ITER: usize = 1;
|
||||
const BLOCKING: usize = 2;
|
||||
|
||||
for _ in 0..ITER {
|
||||
let pool = Builder::new()
|
||||
.pool_size(4)
|
||||
.max_blocking(1)
|
||||
.build();
|
||||
|
||||
let rem = Arc::new(AtomicUsize::new(BLOCKING));
|
||||
let (tx, rx) = mpsc::channel();
|
||||
|
||||
for _ in 0..BLOCKING {
|
||||
let rem = rem.clone();
|
||||
let tx = tx.clone();
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
poll_fn(move || {
|
||||
// First block
|
||||
let res = blocking(|| {
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
}).map_err(|e| panic!("blocking err {:?}", e));
|
||||
|
||||
try_ready!(res);
|
||||
|
||||
let res = blocking(|| {
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
let prev = rem.fetch_sub(1, Relaxed);
|
||||
|
||||
if prev == 1 {
|
||||
tx.send(()).unwrap();
|
||||
}
|
||||
});
|
||||
|
||||
assert!(res.unwrap().is_ready());
|
||||
|
||||
Ok(().into())
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
rx.recv().unwrap();
|
||||
|
||||
assert_eq!(0, rem.load(Relaxed));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shutdown() {
|
||||
const ITER: usize = 1_000;
|
||||
const BLOCKING: usize = 10;
|
||||
|
||||
for _ in 0..ITER {
|
||||
let num_inc = Arc::new(AtomicUsize::new(0));
|
||||
let num_dec = Arc::new(AtomicUsize::new(0));
|
||||
let (tx, rx) = mpsc::channel();
|
||||
|
||||
let pool = {
|
||||
let num_inc = num_inc.clone();
|
||||
let num_dec = num_dec.clone();
|
||||
|
||||
Builder::new()
|
||||
.pool_size(1)
|
||||
.max_blocking(BLOCKING)
|
||||
.after_start(move || { num_inc.fetch_add(1, Relaxed); })
|
||||
.before_stop(move || { num_dec.fetch_add(1, Relaxed); })
|
||||
.build()
|
||||
};
|
||||
|
||||
let barrier = Arc::new(Barrier::new(BLOCKING));
|
||||
|
||||
for _ in 0..BLOCKING {
|
||||
let barrier = barrier.clone();
|
||||
let tx = tx.clone();
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
let res = blocking(|| {
|
||||
barrier.wait();
|
||||
Ok::<_, ()>(())
|
||||
}).unwrap();
|
||||
|
||||
tx.send(()).unwrap();
|
||||
|
||||
assert!(res.is_ready());
|
||||
Ok(().into())
|
||||
}));
|
||||
}
|
||||
|
||||
for _ in 0..BLOCKING {
|
||||
rx.recv().unwrap();
|
||||
}
|
||||
|
||||
// Shutdown
|
||||
drop(pool);
|
||||
|
||||
assert_eq!(11, num_inc.load(Relaxed));
|
||||
assert_eq!(11, num_dec.load(Relaxed));
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Copy, Clone)]
|
||||
enum Sleep {
|
||||
Skip,
|
||||
Yield,
|
||||
Rand,
|
||||
Fixed(Duration),
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hammer() {
|
||||
use self::Sleep::*;
|
||||
|
||||
const ITER: usize = 5;
|
||||
|
||||
let combos = [
|
||||
(2, 4, 1_000, Skip),
|
||||
(2, 4, 1_000, Yield),
|
||||
(2, 4, 100, Rand),
|
||||
(2, 4, 100, Fixed(Duration::from_millis(3))),
|
||||
(2, 4, 100, Fixed(Duration::from_millis(12))),
|
||||
];
|
||||
|
||||
for &(size, max_blocking, n, sleep) in &combos {
|
||||
for _ in 0..ITER {
|
||||
let pool = Builder::new()
|
||||
.pool_size(size)
|
||||
.max_blocking(max_blocking)
|
||||
.build();
|
||||
|
||||
let cnt_task = Arc::new(AtomicUsize::new(0));
|
||||
let cnt_block = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
for _ in 0..n {
|
||||
let cnt_task = cnt_task.clone();
|
||||
let cnt_block = cnt_block.clone();
|
||||
|
||||
pool.spawn(lazy(move || {
|
||||
cnt_task.fetch_add(1, Relaxed);
|
||||
|
||||
poll_fn(move || {
|
||||
blocking(|| {
|
||||
match sleep {
|
||||
Skip => {}
|
||||
Yield => {
|
||||
thread::yield_now();
|
||||
}
|
||||
Rand => {
|
||||
let ms = thread_rng().gen_range(3, 12);
|
||||
thread::sleep(Duration::from_millis(ms));
|
||||
}
|
||||
Fixed(dur) => {
|
||||
thread::sleep(dur);
|
||||
}
|
||||
}
|
||||
|
||||
cnt_block.fetch_add(1, Relaxed);
|
||||
}).map_err(|_| panic!())
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
// Wait for the work to complete
|
||||
pool.shutdown_on_idle().wait().unwrap();
|
||||
|
||||
assert_eq!(n, cnt_task.load(Relaxed));
|
||||
assert_eq!(n, cnt_block.load(Relaxed));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
extern crate futures;
|
||||
extern crate tokio_threadpool;
|
||||
|
||||
use tokio_threadpool::*;
|
||||
|
||||
use futures::{Future, Stream, Sink, Poll};
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering::*;
|
||||
|
||||
#[test]
|
||||
fn hammer() {
|
||||
use futures::future;
|
||||
use futures::sync::{oneshot, mpsc};
|
||||
|
||||
const N: usize = 1000;
|
||||
const ITER: usize = 20;
|
||||
|
||||
struct Counted<T> {
|
||||
cnt: Arc<AtomicUsize>,
|
||||
inner: T,
|
||||
}
|
||||
|
||||
impl<T: Future> Future for Counted<T> {
|
||||
type Item = T::Item;
|
||||
type Error = T::Error;
|
||||
|
||||
fn poll(&mut self) -> Poll<T::Item, T::Error> {
|
||||
self.inner.poll()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Drop for Counted<T> {
|
||||
fn drop(&mut self) {
|
||||
self.cnt.fetch_add(1, Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
for _ in 0.. ITER {
|
||||
let pool = Builder::new()
|
||||
// .pool_size(30)
|
||||
.build();
|
||||
|
||||
let cnt = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let (listen_tx, listen_rx) = mpsc::unbounded::<oneshot::Sender<oneshot::Sender<()>>>();
|
||||
let mut listen_tx = listen_tx.wait();
|
||||
|
||||
pool.spawn({
|
||||
let c1 = cnt.clone();
|
||||
let c2 = cnt.clone();
|
||||
let pool = pool.sender().clone();
|
||||
let task = listen_rx
|
||||
.map_err(|e| panic!("accept error = {:?}", e))
|
||||
.for_each(move |tx| {
|
||||
let task = future::lazy(|| {
|
||||
let (tx2, rx2) = oneshot::channel();
|
||||
|
||||
tx.send(tx2).unwrap();
|
||||
rx2
|
||||
})
|
||||
.map_err(|e| panic!("e={:?}", e))
|
||||
.and_then(|_| {
|
||||
Ok(())
|
||||
});
|
||||
|
||||
pool.spawn(Counted {
|
||||
inner: task,
|
||||
cnt: c1.clone(),
|
||||
}).unwrap();
|
||||
|
||||
Ok(())
|
||||
});
|
||||
|
||||
Counted {
|
||||
inner: task,
|
||||
cnt: c2,
|
||||
}
|
||||
});
|
||||
|
||||
for _ in 0..N {
|
||||
let cnt = cnt.clone();
|
||||
let (tx, rx) = oneshot::channel();
|
||||
listen_tx.send(tx).unwrap();
|
||||
|
||||
pool.spawn({
|
||||
let task = rx
|
||||
.map_err(|e| panic!("rx err={:?}", e))
|
||||
.and_then(|tx| {
|
||||
tx.send(()).unwrap();
|
||||
Ok(())
|
||||
});
|
||||
|
||||
Counted {
|
||||
inner: task,
|
||||
cnt,
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
drop(listen_tx);
|
||||
|
||||
pool.shutdown_on_idle().wait().unwrap();
|
||||
assert_eq!(N * 2 + 1, cnt.load(Relaxed));
|
||||
}
|
||||
}
|
||||
@@ -26,7 +26,7 @@ fn lazy<R, F>(f: F) -> Box<Future<Item = R::Item, Error = R::Error> + Send> wher
|
||||
|
||||
use std::cell::Cell;
|
||||
use std::sync::{mpsc, Arc};
|
||||
use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT};
|
||||
use std::sync::atomic::*;
|
||||
use std::sync::atomic::Ordering::Relaxed;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -88,19 +88,25 @@ fn natural_shutdown_simple_futures() {
|
||||
let _ = ::env_logger::init();
|
||||
|
||||
for _ in 0..1_000 {
|
||||
static NUM_INC: AtomicUsize = ATOMIC_USIZE_INIT;
|
||||
static NUM_DEC: AtomicUsize = ATOMIC_USIZE_INIT;
|
||||
let num_inc = Arc::new(AtomicUsize::new(0));
|
||||
let num_dec = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
FOO.with(|f| {
|
||||
f.set(1);
|
||||
|
||||
let pool = Builder::new()
|
||||
.around_worker(|w, _| {
|
||||
NUM_INC.fetch_add(1, Relaxed);
|
||||
w.run();
|
||||
NUM_DEC.fetch_add(1, Relaxed);
|
||||
})
|
||||
.build();
|
||||
let pool = {
|
||||
let num_inc = num_inc.clone();
|
||||
let num_dec = num_dec.clone();
|
||||
|
||||
Builder::new()
|
||||
.around_worker(move |w, _| {
|
||||
num_inc.fetch_add(1, Relaxed);
|
||||
w.run();
|
||||
num_dec.fetch_add(1, Relaxed);
|
||||
})
|
||||
.build()
|
||||
};
|
||||
|
||||
let mut tx = pool.sender().clone();
|
||||
|
||||
let a = {
|
||||
@@ -136,11 +142,11 @@ fn natural_shutdown_simple_futures() {
|
||||
await_shutdown(pool.shutdown());
|
||||
|
||||
// Assert that at least one thread started
|
||||
let num_inc = NUM_INC.load(Relaxed);
|
||||
let num_inc = num_inc.load(Relaxed);
|
||||
assert!(num_inc > 0);
|
||||
|
||||
// Assert that all threads shutdown
|
||||
let num_dec = NUM_DEC.load(Relaxed);
|
||||
let num_dec = num_dec.load(Relaxed);
|
||||
assert_eq!(num_inc, num_dec);
|
||||
});
|
||||
}
|
||||
@@ -255,6 +261,8 @@ fn drop_threadpool_drops_futures() {
|
||||
let b = num_dec.clone();
|
||||
|
||||
let pool = Builder::new()
|
||||
.max_blocking(2)
|
||||
.pool_size(20)
|
||||
.around_worker(move |w, _| {
|
||||
a.fetch_add(1, Relaxed);
|
||||
w.run();
|
||||
@@ -471,7 +479,11 @@ fn busy_threadpool_is_not_idle() {
|
||||
#[cfg(feature = "unstable-futures")]
|
||||
use futures2::channel::oneshot;
|
||||
|
||||
let pool = ThreadPool::new();
|
||||
// let pool = ThreadPool::new();
|
||||
let pool = Builder::new()
|
||||
.pool_size(4)
|
||||
.max_blocking(2)
|
||||
.build();
|
||||
let mut tx = pool.sender().clone();
|
||||
|
||||
let (term_tx, term_rx) = oneshot::channel();
|
||||
@@ -548,101 +560,3 @@ fn panic_in_task() {
|
||||
|
||||
await_shutdown(pool.shutdown_on_idle());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(not(feature = "unstable-futures"))]
|
||||
fn hammer() {
|
||||
use futures::future;
|
||||
use futures::sync::{oneshot, mpsc};
|
||||
|
||||
const N: usize = 1000;
|
||||
const ITER: usize = 20;
|
||||
|
||||
struct Counted<T> {
|
||||
cnt: Arc<AtomicUsize>,
|
||||
inner: T,
|
||||
}
|
||||
|
||||
impl<T: Future> Future for Counted<T> {
|
||||
type Item = T::Item;
|
||||
type Error = T::Error;
|
||||
|
||||
fn poll(&mut self) -> Poll<T::Item, T::Error> {
|
||||
self.inner.poll()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Drop for Counted<T> {
|
||||
fn drop(&mut self) {
|
||||
self.cnt.fetch_add(1, Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
for _ in 0.. ITER {
|
||||
let pool = Builder::new()
|
||||
// .pool_size(30)
|
||||
.build();
|
||||
|
||||
let cnt = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let (listen_tx, listen_rx) = mpsc::unbounded::<oneshot::Sender<oneshot::Sender<()>>>();
|
||||
let mut listen_tx = listen_tx.wait();
|
||||
|
||||
pool.spawn({
|
||||
let c1 = cnt.clone();
|
||||
let c2 = cnt.clone();
|
||||
let pool = pool.sender().clone();
|
||||
let task = listen_rx
|
||||
.map_err(|e| panic!("accept error = {:?}", e))
|
||||
.for_each(move |tx| {
|
||||
let task = future::lazy(|| {
|
||||
let (tx2, rx2) = oneshot::channel();
|
||||
|
||||
tx.send(tx2).unwrap();
|
||||
rx2
|
||||
})
|
||||
.map_err(|e| panic!("e={:?}", e))
|
||||
.and_then(|_| {
|
||||
Ok(())
|
||||
});
|
||||
|
||||
pool.spawn(Counted {
|
||||
inner: task,
|
||||
cnt: c1.clone(),
|
||||
}).unwrap();
|
||||
|
||||
Ok(())
|
||||
});
|
||||
|
||||
Counted {
|
||||
inner: task,
|
||||
cnt: c2,
|
||||
}
|
||||
});
|
||||
|
||||
for _ in 0..N {
|
||||
let cnt = cnt.clone();
|
||||
let (tx, rx) = oneshot::channel();
|
||||
listen_tx.send(tx).unwrap();
|
||||
|
||||
pool.spawn({
|
||||
let task = rx
|
||||
.map_err(|e| panic!("rx err={:?}", e))
|
||||
.and_then(|tx| {
|
||||
tx.send(()).unwrap();
|
||||
Ok(())
|
||||
});
|
||||
|
||||
Counted {
|
||||
inner: task,
|
||||
cnt,
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
drop(listen_tx);
|
||||
|
||||
pool.shutdown_on_idle().wait().unwrap();
|
||||
assert_eq!(N * 2 + 1, cnt.load(Relaxed));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user