mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-27 00:00:12 +02:00
threadpool: move threadpool into tokio-executor (#1452)
The threadpool is behind a feature flag. Refs: #1264
This commit is contained in:
Carl Lerche
@@ -0,0 +1,526 @@
|
||||
#![warn(rust_2018_idioms)]
|
||||
#![feature(async_await)]
|
||||
|
||||
use tokio_executor::park::{Park, Unpark};
|
||||
use tokio_executor::threadpool;
|
||||
use tokio_executor::threadpool::park::{DefaultPark, DefaultUnpark};
|
||||
use tokio_executor::threadpool::*;
|
||||
use tokio_test::assert_pending;
|
||||
|
||||
use std::cell::Cell;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::Ordering::Relaxed;
|
||||
use std::sync::atomic::*;
|
||||
use std::sync::{mpsc, Arc};
|
||||
use std::task::{Context, Poll, Waker};
|
||||
use std::time::Duration;
|
||||
|
||||
thread_local!(static FOO: Cell<u32> = Cell::new(0));
|
||||
|
||||
#[test]
|
||||
fn natural_shutdown_simple_futures() {
|
||||
for _ in 0..1_000 {
|
||||
let num_inc = Arc::new(AtomicUsize::new(0));
|
||||
let num_dec = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
FOO.with(|f| {
|
||||
f.set(1);
|
||||
|
||||
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 tx = pool.sender().clone();
|
||||
|
||||
let a = {
|
||||
let (t, rx) = mpsc::channel();
|
||||
tx.spawn(async move {
|
||||
// Makes sure this runs on a worker thread
|
||||
FOO.with(|f| assert_eq!(f.get(), 0));
|
||||
|
||||
t.send("one").unwrap();
|
||||
})
|
||||
.unwrap();
|
||||
rx
|
||||
};
|
||||
|
||||
let b = {
|
||||
let (t, rx) = mpsc::channel();
|
||||
tx.spawn(async move {
|
||||
// Makes sure this runs on a worker thread
|
||||
FOO.with(|f| assert_eq!(f.get(), 0));
|
||||
|
||||
t.send("two").unwrap();
|
||||
})
|
||||
.unwrap();
|
||||
rx
|
||||
};
|
||||
|
||||
drop(tx);
|
||||
|
||||
assert_eq!("one", a.recv().unwrap());
|
||||
assert_eq!("two", b.recv().unwrap());
|
||||
|
||||
// Wait for the pool to shutdown
|
||||
pool.shutdown().wait();
|
||||
|
||||
// Assert that at least one thread started
|
||||
let num_inc = num_inc.load(Relaxed);
|
||||
assert!(num_inc > 0);
|
||||
|
||||
// Assert that all threads shutdown
|
||||
let num_dec = num_dec.load(Relaxed);
|
||||
assert_eq!(num_inc, num_dec);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn force_shutdown_drops_futures() {
|
||||
for _ in 0..1_000 {
|
||||
let num_inc = Arc::new(AtomicUsize::new(0));
|
||||
let num_dec = Arc::new(AtomicUsize::new(0));
|
||||
let num_drop = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
struct Never(Arc<AtomicUsize>);
|
||||
|
||||
impl Future for Never {
|
||||
type Output = ();
|
||||
|
||||
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Never {
|
||||
fn drop(&mut self) {
|
||||
self.0.fetch_add(1, Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
let a = num_inc.clone();
|
||||
let b = num_dec.clone();
|
||||
|
||||
let pool = Builder::new()
|
||||
.around_worker(move |w| {
|
||||
a.fetch_add(1, Relaxed);
|
||||
w.run();
|
||||
b.fetch_add(1, Relaxed);
|
||||
})
|
||||
.build();
|
||||
let tx = pool.sender().clone();
|
||||
|
||||
tx.spawn(Never(num_drop.clone())).unwrap();
|
||||
|
||||
// Wait for the pool to shutdown
|
||||
pool.shutdown_now().wait();
|
||||
|
||||
// Assert that only a single thread was spawned.
|
||||
let a = num_inc.load(Relaxed);
|
||||
assert!(a >= 1);
|
||||
|
||||
// Assert that all threads shutdown
|
||||
let b = num_dec.load(Relaxed);
|
||||
assert_eq!(a, b);
|
||||
|
||||
// Assert that the future was dropped
|
||||
let c = num_drop.load(Relaxed);
|
||||
assert_eq!(c, 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_threadpool_drops_futures() {
|
||||
for _ in 0..1_000 {
|
||||
let num_inc = Arc::new(AtomicUsize::new(0));
|
||||
let num_dec = Arc::new(AtomicUsize::new(0));
|
||||
let num_drop = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
struct Never(Arc<AtomicUsize>);
|
||||
|
||||
impl Future for Never {
|
||||
type Output = ();
|
||||
|
||||
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Never {
|
||||
fn drop(&mut self) {
|
||||
self.0.fetch_add(1, Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
let a = num_inc.clone();
|
||||
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();
|
||||
b.fetch_add(1, Relaxed);
|
||||
})
|
||||
.build();
|
||||
let tx = pool.sender().clone();
|
||||
|
||||
tx.spawn(Never(num_drop.clone())).unwrap();
|
||||
|
||||
// Wait for the pool to shutdown
|
||||
drop(pool);
|
||||
|
||||
// Assert that only a single thread was spawned.
|
||||
let a = num_inc.load(Relaxed);
|
||||
assert!(a >= 1);
|
||||
|
||||
// Assert that all threads shutdown
|
||||
let b = num_dec.load(Relaxed);
|
||||
assert_eq!(a, b);
|
||||
|
||||
// Assert that the future was dropped
|
||||
let c = num_drop.load(Relaxed);
|
||||
assert_eq!(c, 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn many_oneshot_futures() {
|
||||
const NUM: usize = 10_000;
|
||||
|
||||
for _ in 0..50 {
|
||||
let pool = ThreadPool::new();
|
||||
let tx = pool.sender().clone();
|
||||
let cnt = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
for _ in 0..NUM {
|
||||
let cnt = cnt.clone();
|
||||
tx.spawn(async move {
|
||||
cnt.fetch_add(1, Relaxed);
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Wait for the pool to shutdown
|
||||
pool.shutdown().wait();
|
||||
|
||||
let num = cnt.load(Relaxed);
|
||||
assert_eq!(num, NUM);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn many_multishot_futures() {
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
const CHAIN: usize = 200;
|
||||
const CYCLES: usize = 5;
|
||||
const TRACKS: usize = 50;
|
||||
|
||||
for _ in 0..50 {
|
||||
let pool = ThreadPool::new();
|
||||
let pool_tx = pool.sender().clone();
|
||||
|
||||
let mut start_txs = Vec::with_capacity(TRACKS);
|
||||
let mut final_rxs = Vec::with_capacity(TRACKS);
|
||||
|
||||
for _ in 0..TRACKS {
|
||||
let (start_tx, mut chain_rx) = mpsc::channel(10);
|
||||
|
||||
for _ in 0..CHAIN {
|
||||
let (mut next_tx, next_rx) = mpsc::channel(10);
|
||||
|
||||
// Forward all the messages
|
||||
pool_tx
|
||||
.spawn(async move {
|
||||
while let Some(v) = chain_rx.recv().await {
|
||||
next_tx.send(v).await.unwrap();
|
||||
}
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
chain_rx = next_rx;
|
||||
}
|
||||
|
||||
// This final task cycles if needed
|
||||
let (mut final_tx, final_rx) = mpsc::channel(10);
|
||||
let mut cycle_tx = start_tx.clone();
|
||||
let mut rem = CYCLES;
|
||||
|
||||
pool_tx
|
||||
.spawn(async move {
|
||||
for _ in 0..CYCLES {
|
||||
let msg = chain_rx.recv().await.unwrap();
|
||||
|
||||
rem -= 1;
|
||||
|
||||
if rem == 0 {
|
||||
final_tx.send(msg).await.unwrap();
|
||||
} else {
|
||||
cycle_tx.send(msg).await.unwrap();
|
||||
}
|
||||
}
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
start_txs.push(start_tx);
|
||||
final_rxs.push(final_rx);
|
||||
}
|
||||
|
||||
{
|
||||
let mut e = tokio_executor::enter().unwrap();
|
||||
|
||||
e.block_on(async move {
|
||||
for mut start_tx in start_txs {
|
||||
start_tx.send("ping").await.unwrap();
|
||||
}
|
||||
|
||||
for mut final_rx in final_rxs {
|
||||
final_rx.recv().await.unwrap();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Shutdown the pool
|
||||
pool.shutdown().wait();
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn global_executor_is_configured() {
|
||||
let pool = ThreadPool::new();
|
||||
let tx = pool.sender().clone();
|
||||
|
||||
let (signal_tx, signal_rx) = mpsc::channel();
|
||||
|
||||
tx.spawn(async move {
|
||||
tokio_executor::spawn(async move {
|
||||
signal_tx.send(()).unwrap();
|
||||
});
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
signal_rx.recv().unwrap();
|
||||
|
||||
pool.shutdown().wait();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_threadpool_is_idle() {
|
||||
let pool = ThreadPool::new();
|
||||
pool.shutdown_on_idle().wait();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn busy_threadpool_is_not_idle() {
|
||||
use tokio_sync::oneshot;
|
||||
|
||||
// let pool = ThreadPool::new();
|
||||
let pool = Builder::new().pool_size(4).max_blocking(2).build();
|
||||
let tx = pool.sender().clone();
|
||||
|
||||
let (term_tx, term_rx) = oneshot::channel();
|
||||
|
||||
tx.spawn(async move {
|
||||
term_rx.await.unwrap();
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let mut idle = pool.shutdown_on_idle();
|
||||
|
||||
struct IdleFut<'a>(&'a mut Shutdown);
|
||||
|
||||
impl Future for IdleFut<'_> {
|
||||
type Output = ();
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
|
||||
assert_pending!(Pin::new(&mut self.as_mut().0).poll(cx));
|
||||
Poll::Ready(())
|
||||
}
|
||||
}
|
||||
|
||||
let idle_fut = IdleFut(&mut idle);
|
||||
tokio_executor::enter().unwrap().block_on(idle_fut);
|
||||
|
||||
term_tx.send(()).unwrap();
|
||||
|
||||
let idle_fut = IdleFut(&mut idle);
|
||||
tokio_executor::enter().unwrap().block_on(idle_fut);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn panic_in_task() {
|
||||
let pool = ThreadPool::new();
|
||||
let tx = pool.sender().clone();
|
||||
|
||||
struct Boom;
|
||||
|
||||
impl Future for Boom {
|
||||
type Output = ();
|
||||
|
||||
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
|
||||
panic!();
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Boom {
|
||||
fn drop(&mut self) {
|
||||
assert!(::std::thread::panicking());
|
||||
}
|
||||
}
|
||||
|
||||
tx.spawn(Boom).unwrap();
|
||||
|
||||
pool.shutdown_on_idle().wait();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn count_panics() {
|
||||
let counter = Arc::new(AtomicUsize::new(0));
|
||||
let counter_ = counter.clone();
|
||||
let pool = threadpool::Builder::new()
|
||||
.panic_handler(move |_err| {
|
||||
// We caught a panic.
|
||||
counter_.fetch_add(1, Relaxed);
|
||||
})
|
||||
.build();
|
||||
// Spawn a future that will panic.
|
||||
pool.spawn(async { panic!() });
|
||||
pool.shutdown_on_idle().wait();
|
||||
let counter = counter.load(Relaxed);
|
||||
assert_eq!(counter, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_threadpool() {
|
||||
use tokio_sync::oneshot;
|
||||
|
||||
let pool1 = ThreadPool::new();
|
||||
let pool2 = ThreadPool::new();
|
||||
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let (done_tx, done_rx) = mpsc::channel();
|
||||
|
||||
pool2.spawn(async move {
|
||||
rx.await.unwrap();
|
||||
done_tx.send(()).unwrap();
|
||||
});
|
||||
|
||||
pool1.spawn(async move {
|
||||
tx.send(()).unwrap();
|
||||
});
|
||||
|
||||
done_rx.recv().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eagerly_drops_futures() {
|
||||
use std::sync::mpsc;
|
||||
|
||||
struct MyPark {
|
||||
inner: DefaultPark,
|
||||
#[allow(dead_code)]
|
||||
park_tx: mpsc::SyncSender<()>,
|
||||
unpark_tx: mpsc::SyncSender<()>,
|
||||
}
|
||||
|
||||
impl Park for MyPark {
|
||||
type Unpark = MyUnpark;
|
||||
type Error = <DefaultPark as Park>::Error;
|
||||
|
||||
fn unpark(&self) -> Self::Unpark {
|
||||
MyUnpark {
|
||||
inner: self.inner.unpark(),
|
||||
unpark_tx: self.unpark_tx.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn park(&mut self) -> Result<(), Self::Error> {
|
||||
self.inner.park()
|
||||
}
|
||||
|
||||
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
|
||||
self.inner.park_timeout(duration)
|
||||
}
|
||||
}
|
||||
|
||||
struct MyUnpark {
|
||||
inner: DefaultUnpark,
|
||||
#[allow(dead_code)]
|
||||
unpark_tx: mpsc::SyncSender<()>,
|
||||
}
|
||||
|
||||
impl Unpark for MyUnpark {
|
||||
fn unpark(&self) {
|
||||
self.inner.unpark()
|
||||
}
|
||||
}
|
||||
|
||||
let (task_tx, task_rx) = mpsc::channel();
|
||||
let (drop_tx, drop_rx) = mpsc::channel();
|
||||
let (park_tx, park_rx) = mpsc::sync_channel(0);
|
||||
let (unpark_tx, unpark_rx) = mpsc::sync_channel(0);
|
||||
|
||||
let pool = threadpool::Builder::new()
|
||||
.custom_park(move |_| MyPark {
|
||||
inner: DefaultPark::new(),
|
||||
park_tx: park_tx.clone(),
|
||||
unpark_tx: unpark_tx.clone(),
|
||||
})
|
||||
.build();
|
||||
|
||||
struct MyTask {
|
||||
task_tx: Option<mpsc::Sender<Waker>>,
|
||||
drop_tx: mpsc::Sender<()>,
|
||||
}
|
||||
|
||||
impl Future for MyTask {
|
||||
type Output = ();
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
|
||||
if let Some(tx) = self.get_mut().task_tx.take() {
|
||||
tx.send(cx.waker().clone()).unwrap();
|
||||
}
|
||||
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MyTask {
|
||||
fn drop(&mut self) {
|
||||
self.drop_tx.send(()).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
pool.spawn(MyTask {
|
||||
task_tx: Some(task_tx),
|
||||
drop_tx,
|
||||
});
|
||||
|
||||
// Wait until we get the task handle.
|
||||
let task = task_rx.recv().unwrap();
|
||||
|
||||
// Drop the pool, this should result in futures being forcefully dropped.
|
||||
drop(pool);
|
||||
|
||||
// Make sure `MyPark` and `MyUnpark` were dropped during shutdown.
|
||||
assert_eq!(park_rx.try_recv(), Err(mpsc::TryRecvError::Disconnected));
|
||||
assert_eq!(unpark_rx.try_recv(), Err(mpsc::TryRecvError::Disconnected));
|
||||
|
||||
// If the future is forcefully dropped, then we will get a signal here.
|
||||
drop_rx.recv().unwrap();
|
||||
|
||||
// Ensure `task` lives until after the test completes.
|
||||
drop(task);
|
||||
}
|
||||
@@ -0,0 +1,413 @@
|
||||
#![warn(rust_2018_idioms)]
|
||||
#![feature(async_await)]
|
||||
|
||||
use tokio_executor::threadpool::*;
|
||||
use tokio_test::*;
|
||||
|
||||
use futures_core::ready;
|
||||
use futures_util::future::poll_fn;
|
||||
use rand::*;
|
||||
use std::sync::atomic::Ordering::*;
|
||||
use std::sync::atomic::*;
|
||||
use std::sync::*;
|
||||
use std::task::{Poll, Waker};
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
#[test]
|
||||
fn basic() {
|
||||
let pool = Builder::new().pool_size(1).max_blocking(1).build();
|
||||
|
||||
let (tx1, rx1) = mpsc::channel();
|
||||
let (tx2, rx2) = mpsc::channel();
|
||||
|
||||
pool.spawn(async move {
|
||||
let res = blocking(|| {
|
||||
let v = rx1.recv().unwrap();
|
||||
tx2.send(v).unwrap();
|
||||
});
|
||||
|
||||
assert_ready!(res).unwrap();
|
||||
});
|
||||
|
||||
pool.spawn(async move {
|
||||
tx1.send(()).unwrap();
|
||||
});
|
||||
|
||||
rx2.recv().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn other_executors_can_run_inside_blocking() {
|
||||
let pool = Builder::new().pool_size(1).max_blocking(1).build();
|
||||
|
||||
let (tx, rx) = mpsc::channel();
|
||||
|
||||
pool.spawn(async move {
|
||||
let res = blocking(|| {
|
||||
let _e = tokio_executor::enter().expect("nested blocking enter");
|
||||
tx.send(()).unwrap();
|
||||
});
|
||||
|
||||
assert_ready!(res).unwrap();
|
||||
});
|
||||
|
||||
rx.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(async move {
|
||||
poll_fn(move |_| {
|
||||
blocking(|| {
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
let prev = rem.fetch_sub(1, SeqCst);
|
||||
|
||||
if prev == 1 {
|
||||
tx.send(()).unwrap();
|
||||
}
|
||||
})
|
||||
.map_err(|e| panic!("blocking err {:?}", e))
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
rx.recv().unwrap();
|
||||
|
||||
assert_eq!(0, rem.load(SeqCst));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_is_use_it_or_lose_it() {
|
||||
use tokio_sync::oneshot;
|
||||
|
||||
// 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(async move {
|
||||
poll_fn(move |_| {
|
||||
blocking(|| {
|
||||
rx1.recv().unwrap();
|
||||
})
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
pool.spawn(async move {
|
||||
let task: Waker = rx2.await.unwrap();
|
||||
|
||||
poll_fn(move |_| {
|
||||
blocking(|| {
|
||||
// Notify the other task
|
||||
task.wake_by_ref();
|
||||
|
||||
// Block until woken
|
||||
rx3.recv().unwrap();
|
||||
})
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
// 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(async move {
|
||||
poll_fn(move |cx| {
|
||||
match i {
|
||||
0 => {
|
||||
i = 1;
|
||||
|
||||
let res = blocking(|| unreachable!()).map_err(|_| panic!());
|
||||
|
||||
assert_pending!(res);
|
||||
|
||||
// Unblock the first blocker
|
||||
tx1.send(()).unwrap();
|
||||
|
||||
return Poll::Pending;
|
||||
}
|
||||
1 => {
|
||||
i = 2;
|
||||
|
||||
// Skip blocking, and notify the second task that it should
|
||||
// start blocking
|
||||
let me = cx.waker().clone();
|
||||
tx2.take().unwrap().send(me).unwrap();
|
||||
|
||||
return Poll::Pending;
|
||||
}
|
||||
2 => {
|
||||
let res = blocking(|| unreachable!()).map_err(|_| panic!());
|
||||
|
||||
assert_pending!(res);
|
||||
|
||||
// Unblock the first blocker
|
||||
tx3.send(()).unwrap();
|
||||
tx4.send(()).unwrap();
|
||||
Poll::Ready(())
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
})
|
||||
.await
|
||||
});
|
||||
|
||||
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(async move {
|
||||
poll_fn(move |_| {
|
||||
blocking(|| {
|
||||
enter_tx.send(()).unwrap();
|
||||
exit_rx.recv().unwrap();
|
||||
exited.store(true, SeqCst);
|
||||
})
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
// Wait for the task to block
|
||||
let _ = enter_rx.recv().unwrap();
|
||||
|
||||
// Spawn another task that attempts to block
|
||||
pool.spawn(async move {
|
||||
poll_fn(move |_| {
|
||||
let res = blocking(|| {});
|
||||
|
||||
assert_eq!(res.is_ready(), exited.load(SeqCst));
|
||||
|
||||
try_tx.send(res.is_ready()).unwrap();
|
||||
|
||||
res.map(|_| ())
|
||||
})
|
||||
.await
|
||||
});
|
||||
|
||||
// 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 blocking_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(async move {
|
||||
poll_fn(move |_| {
|
||||
// First block
|
||||
let res = blocking(|| {
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
});
|
||||
|
||||
ready!(res).unwrap();
|
||||
|
||||
let res = blocking(|| {
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
let prev = rem.fetch_sub(1, SeqCst);
|
||||
|
||||
if prev == 1 {
|
||||
tx.send(()).unwrap();
|
||||
}
|
||||
});
|
||||
|
||||
assert!(res.is_ready());
|
||||
|
||||
Poll::Ready(())
|
||||
})
|
||||
.await
|
||||
});
|
||||
}
|
||||
|
||||
rx.recv().unwrap();
|
||||
|
||||
assert_eq!(0, rem.load(SeqCst));
|
||||
}
|
||||
}
|
||||
|
||||
#[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, SeqCst);
|
||||
})
|
||||
.before_stop(move || {
|
||||
num_dec.fetch_add(1, SeqCst);
|
||||
})
|
||||
.build()
|
||||
};
|
||||
|
||||
let barrier = Arc::new(Barrier::new(BLOCKING));
|
||||
|
||||
for _ in 0..BLOCKING {
|
||||
let barrier = barrier.clone();
|
||||
let tx = tx.clone();
|
||||
|
||||
pool.spawn(async move {
|
||||
let res = blocking(|| {
|
||||
barrier.wait();
|
||||
Ok::<_, ()>(())
|
||||
});
|
||||
|
||||
tx.send(()).unwrap();
|
||||
|
||||
assert!(res.is_ready());
|
||||
});
|
||||
}
|
||||
|
||||
for _ in 0..BLOCKING {
|
||||
rx.recv().unwrap();
|
||||
}
|
||||
|
||||
// Shutdown
|
||||
drop(pool);
|
||||
|
||||
assert_eq!(11, num_inc.load(SeqCst));
|
||||
assert_eq!(11, num_dec.load(SeqCst));
|
||||
}
|
||||
}
|
||||
|
||||
#[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(async move {
|
||||
cnt_task.fetch_add(1, SeqCst);
|
||||
|
||||
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, SeqCst);
|
||||
})
|
||||
.map_err(|_| panic!())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
// Wait for the work to complete
|
||||
pool.shutdown_on_idle().wait();
|
||||
|
||||
assert_eq!(n, cnt_task.load(SeqCst));
|
||||
assert_eq!(n, cnt_block.load(SeqCst));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
#![warn(rust_2018_idioms)]
|
||||
#![feature(async_await)]
|
||||
|
||||
use tokio_executor::threadpool::*;
|
||||
use tokio_sync::{mpsc, oneshot};
|
||||
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering::*;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
#[test]
|
||||
fn hammer() {
|
||||
const N: usize = 1000;
|
||||
const ITER: usize = 20;
|
||||
|
||||
struct Counted<T> {
|
||||
cnt: Arc<AtomicUsize>,
|
||||
inner: T,
|
||||
}
|
||||
|
||||
impl<T: Future> Future for Counted<T> {
|
||||
type Output = T::Output;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T::Output> {
|
||||
unsafe {
|
||||
let inner = &mut self.get_unchecked_mut().inner;
|
||||
Pin::new_unchecked(inner).poll(cx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 (mut listen_tx, mut listen_rx) =
|
||||
mpsc::unbounded_channel::<oneshot::Sender<oneshot::Sender<()>>>();
|
||||
|
||||
pool.spawn({
|
||||
let c1 = cnt.clone();
|
||||
let c2 = cnt.clone();
|
||||
let pool = pool.sender().clone();
|
||||
let task = async move {
|
||||
while let Some(tx) = listen_rx.recv().await {
|
||||
let task = async {
|
||||
let (tx2, rx2) = oneshot::channel();
|
||||
tx.send(tx2).unwrap();
|
||||
rx2.await.unwrap()
|
||||
};
|
||||
|
||||
pool.spawn(Counted {
|
||||
inner: task,
|
||||
cnt: c1.clone(),
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
};
|
||||
|
||||
/*
|
||||
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.try_send(tx).unwrap();
|
||||
|
||||
pool.spawn(async {
|
||||
let task = async {
|
||||
let tx = rx.await.unwrap();
|
||||
tx.send(()).unwrap();
|
||||
};
|
||||
|
||||
/*
|
||||
let task = rx.map_err(|e| panic!("rx err={:?}", e)).and_then(|tx| {
|
||||
tx.send(()).unwrap();
|
||||
Ok(())
|
||||
});
|
||||
*/
|
||||
|
||||
Counted { inner: task, cnt }.await
|
||||
});
|
||||
}
|
||||
|
||||
drop(listen_tx);
|
||||
|
||||
pool.shutdown_on_idle().wait();
|
||||
assert_eq!(N * 2 + 1, cnt.load(Relaxed));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user