tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-19 00:00:09 +02:00
Code Issues Packages Projects Releases Wiki Activity

task: add ways to run a LocalSet from within a rt context (#1971)

Browse Source 
Currently, the only way to run a `tokio::task::LocalSet` is to call its
`block_on` method with a `&mut Runtime`, like

```rust
let mut rt = tokio::runtime::Runtime::new();
let local = tokio::task::LocalSet::new();
local.block_on(&mut rt, async {
  // whatever...
});
```

Unfortunately, this means that `LocalSet` doesn't work with the 
`#[tokio::main]`  and `#[tokio::test]` macros, since the `main` 
function is _already_ inside of a call to `block_on`.

**Solution**

This branch adds a `LocalSet::run` method, which takes a future and
returns a new future that runs that future on the `LocalSet`. This
is analogous to `LocalSet::block_on`, except that it can be called in
an async context.

Additionally, this branch implements `Future` for `LocalSet`. Awaiting
a `LocalSet` will run all spawned local futures until they complete.
This allows code like

```rust
#[tokio::main] 
async fn main() {
    let local = tokio::task::LocalSet::new();

    local.spawn_local(async {
        // ...
    });

    local.spawn_local(async {
        // ...
        tokio::task::spawn_local(...);
        // ...
    });

    local.await;
}
```

The `LocalSet` docs have been updated to show the usage with 
`#[tokio::main]` rather than with manually created runtimes, where
applicable.

Closes #1906 
Closes #1908 
Fixes #2057
This commit is contained in:
Eliza Weisman
2020-01-06 14:44:30 -08:00
committed by Carl Lerche
parent 0193df3a59
commit 798e86821f
2 changed files with 593 additions and 448 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/task/local.rs
+173 -439
View File
@@ -22,16 +22,14 @@ cfg_rt_util! {
/// For example, the following code will not compile:
///
/// ```rust,compile_fail
/// # use tokio::runtime::Runtime;
/// use std::rc::Rc;
///
/// // `Rc` does not implement `Send`, and thus may not be sent between
/// // threads safely.
/// let unsend_data = Rc::new("my unsend data...");
/// #[tokio::main]
/// async fn main() {
/// // `Rc` does not implement `Send`, and thus may not be sent between
/// // threads safely.
/// let unsend_data = Rc::new("my unsend data...");
///
/// let mut rt = Runtime::new().unwrap();
///
/// rt.block_on(async move {
/// let unsend_data = unsend_data.clone();
/// // Because the `async` block here moves `unsend_data`, the future is `!Send`.
/// // Since `tokio::spawn` requires the spawned future to implement `Send`, this
@@ -40,7 +38,7 @@ cfg_rt_util! {
/// println!("{}", unsend_data);
/// // ...
/// }).await.unwrap();
/// });
/// }
/// ```
/// In order to spawn `!Send` futures, we can use a local task set to
/// schedule them on the thread calling [`Runtime::block_on`]. When running
@@ -48,26 +46,60 @@ cfg_rt_util! {
/// spawn `!Send` futures. For example:
///
/// ```rust
/// # use tokio::runtime::Runtime;
/// use std::rc::Rc;
/// use tokio::task;
///
/// let unsend_data = Rc::new("my unsend data...");
/// #[tokio::main]
/// async fn main() {
/// let unsend_data = Rc::new("my unsend data...");
///
/// let mut rt = Runtime::new().unwrap();
/// // Construct a local task set that can run `!Send` futures.
/// let local = task::LocalSet::new();
/// // Construct a local task set that can run `!Send` futures.
/// let local = task::LocalSet::new();
///
/// // Run the local task group.
/// local.block_on(&mut rt, async move {
/// let unsend_data = unsend_data.clone();
/// // `spawn_local` ensures that the future is spawned on the local
/// // task group.
/// task::spawn_local(async move {
/// println!("{}", unsend_data);
/// // Run the local task set.
/// local.run_until(async move {
/// let unsend_data = unsend_data.clone();
/// // `spawn_local` ensures that the future is spawned on the local
/// // task set.
/// task::spawn_local(async move {
/// println!("{}", unsend_data);
/// // ...
/// }).await.unwrap();
/// }).await;
/// }
/// ```
///
/// ## Awaiting a `LocalSet`
///
/// Additionally, a `LocalSet` itself implements `Future`, completing when
/// *all* tasks spawned on the `LocalSet` complete. This can be used to run
/// several futures on a `LocalSet` and drive the whole set until they
/// complete. For example,
///
/// ```rust
/// use tokio::{task, time};
/// use std::rc::Rc;
///
/// #[tokio::main]
/// async fn main() {
/// let unsend_data = Rc::new("world");
/// let local = task::LocalSet::new();
///
/// let unsend_data2 = unsend_data.clone();
/// local.spawn_local(async move {
/// // ...
/// }).await.unwrap();
/// });
/// println!("hello {}", unsend_data2)
/// });
///
/// local.spawn_local(async move {
/// time::delay_for(time::Duration::from_millis(100)).await;
/// println!("goodbye {}", unsend_data)
/// });
///
/// // ...
///
/// local.await;
/// }
/// ```
///
/// [`Send`]: https://doc.rust-lang.org/std/marker/trait.Send.html
@@ -92,6 +124,7 @@ struct Scheduler {
}
pin_project! {
#[derive(Debug)]
struct LocalFuture<F> {
scheduler: Rc<Scheduler>,
#[pin]
@@ -116,23 +149,24 @@ cfg_rt_util! {
/// # Examples
///
/// ```rust
/// # use tokio::runtime::Runtime;
/// use std::rc::Rc;
/// use tokio::task;
///
/// let unsend_data = Rc::new("my unsend data...");
/// #[tokio::main]
/// async fn main() {
/// let unsend_data = Rc::new("my unsend data...");
///
/// let mut rt = Runtime::new().unwrap();
/// let local = task::LocalSet::new();
/// let local = task::LocalSet::new();
///
/// // Run the local task set.
/// local.block_on(&mut rt, async move {
/// let unsend_data = unsend_data.clone();
/// task::spawn_local(async move {
/// println!("{}", unsend_data);
/// // ...
/// }).await.unwrap();
/// });
/// // Run the local task set.
/// local.run_until(async move {
/// let unsend_data = unsend_data.clone();
/// task::spawn_local(async move {
/// println!("{}", unsend_data);
/// // ...
/// }).await.unwrap();
/// }).await;
/// }
/// ```
pub fn spawn_local<F>(future: F) -> JoinHandle<F::Output>
where
@@ -173,34 +207,35 @@ impl LocalSet {
/// This task is guaranteed to be run on the current thread.
///
/// Unlike the free function [`spawn_local`], this method may be used to
/// spawn_local local tasks when the task set is _not_ running. For example:
/// spawn local tasks when the task set is _not_ running. For example:
/// ```rust
/// # use tokio::runtime::Runtime;
/// use tokio::task;
///
/// let mut rt = Runtime::new().unwrap();
/// let local = task::LocalSet::new();
/// #[tokio::main]
/// async fn main() {
/// let local = task::LocalSet::new();
///
/// // Spawn a future on the local set. This future will be run when
/// // we call `block_on` to drive the task set.
/// local.spawn_local(async {
/// // ...
/// });
/// // Spawn a future on the local set. This future will be run when
/// // we call `run_until` to drive the task set.
/// local.spawn_local(async {
/// // ...
/// });
///
/// // Run the local task set.
/// local.block_on(&mut rt, async move {
/// // ...
/// });
/// // Run the local task set.
/// local.run_until(async move {
/// // ...
/// }).await;
///
/// // When `block_on` finishes, we can spawn_local _more_ futures, which will
/// // run in subsequent calls to `block_on`.
/// local.spawn_local(async {
/// // ...
/// });
/// // When `run` finishes, we can spawn _more_ futures, which will
/// // run in subsequent calls to `run_until`.
/// local.spawn_local(async {
/// // ...
/// });
///
/// local.block_on(&mut rt, async move {
/// // ...
/// });
/// local.run_until(async move {
/// // ...
/// }).await;
/// }
/// ```
/// [`spawn_local`]: fn.spawn_local.html
pub fn spawn_local<F>(&self, future: F) -> JoinHandle<F::Output>
@@ -280,12 +315,71 @@ impl LocalSet {
/// [in-place blocking]: ../blocking/fn.in_place.html
/// [`spawn_blocking`]: ../blocking/fn.spawn_blocking.html
pub fn block_on<F>(&self, rt: &mut crate::runtime::Runtime, future: F) -> F::Output
where
F: Future,
{
rt.block_on(self.run_until(future))
}
/// Run a future to completion on the local set, returning its output.
///
/// This returns a future that runs the given future with a local set,
/// allowing it to call [`spawn_local`] to spawn additional `!Send` futures.
/// Any local futures spawned on the local set will be driven in the
/// background until the future passed to `run_until` completes. When the future
/// passed to `run` finishes, any local futures which have not completed
/// will remain on the local set, and will be driven on subsequent calls to
/// `run_until` or when [awaiting the local set] itself.
///
/// # Examples
///
/// ```rust
/// use tokio::task;
///
/// #[tokio::main]
/// async fn main() {
/// task::LocalSet::new().run_until(async {
/// task::spawn_local(async move {
/// // ...
/// }).await.unwrap();
/// // ...
/// }).await;
/// }
/// ```
///
/// [`spawn_local`]: fn.spawn_local.html
/// [awaiting the local set]: #awaiting-a-localset
pub async fn run_until<F>(&self, future: F) -> F::Output
where
F: Future,
{
let scheduler = self.scheduler.clone();
self.scheduler
.with(move || rt.block_on(LocalFuture { scheduler, future }))
let future = LocalFuture { scheduler, future };
future.await
}
}
impl Future for LocalSet {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let scheduler = self.as_ref().scheduler.clone();
scheduler.waker.register_by_ref(cx.waker());
if scheduler.with(|| scheduler.tick()) {
// If `tick` returns true, we need to notify the local future again:
// there are still tasks remaining in the run queue.
cx.waker().wake_by_ref();
Poll::Pending
} else if scheduler.is_empty() {
// If the scheduler has no remaining futures, we're done!
Poll::Ready(())
} else {
// There are still futures in the local set, but we've polled all the
// futures in the run queue. Therefore, we can just return Pending
// since the remaining futures will be woken from somewhere else.
Poll::Pending
}
}
}
@@ -295,6 +389,8 @@ impl Default for LocalSet {
}
}
// === impl LocalFuture ===
impl<F: Future> Future for LocalFuture<F> {
type Output = F::Output;
@@ -303,18 +399,19 @@ impl<F: Future> Future for LocalFuture<F> {
let scheduler = this.scheduler;
let mut future = this.future;
scheduler.waker.register_by_ref(cx.waker());
scheduler.with(|| {
if let Poll::Ready(output) = future.as_mut().poll(cx) {
return Poll::Ready(output);
}
if let Poll::Ready(output) = future.as_mut().poll(cx) {
return Poll::Ready(output);
}
if scheduler.tick() {
// If `tick` returns true, we need to notify the local future again:
// there are still tasks remaining in the run queue.
cx.waker().wake_by_ref();
}
if scheduler.tick() {
// If `tick` returns true, we need to notify the local future again:
// there are still tasks remaining in the run queue.
cx.waker().wake_by_ref();
}
Poll::Pending
Poll::Pending
})
}
}
@@ -424,6 +521,15 @@ impl Scheduler {
true
}
fn is_empty(&self) -> bool {
unsafe {
// safety: this method may not be called from threads other than the
// thread that owns the `Queues`. since `Scheduler` is not `Send` or
// `Sync`, that shouldn't happen.
!self.queues.has_tasks_remaining()
}
}
}
impl Drop for Scheduler {
@@ -450,375 +556,3 @@ impl Drop for Scheduler {
}
}
}
#[cfg(all(test, not(loom)))]
mod tests {
use super::*;
use crate::{
runtime,
sync::{mpsc, oneshot},
task, time,
};
use std::time::Duration;
#[test]
fn local_current_thread() {
let mut rt = runtime::Builder::new().basic_scheduler().build().unwrap();
LocalSet::new().block_on(&mut rt, async {
spawn_local(async {}).await.unwrap();
});
}
#[test]
fn local_threadpool() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Runtime::new().unwrap();
LocalSet::new().block_on(&mut rt, async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
})
.await
.unwrap();
});
}
#[test]
fn local_threadpool_timer() {
// This test ensures that runtime services like the timer are properly
// set for the local task set.
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.enable_all()
.build()
.unwrap();
LocalSet::new().block_on(&mut rt, async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let join = spawn_local(async move {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
crate::time::delay_for(Duration::from_millis(10)).await;
assert!(ON_RT_THREAD.with(|cell| cell.get()));
});
join.await.unwrap();
});
}
#[test]
// This will panic, since the thread that calls `block_on` cannot use
// in-place blocking inside of `block_on`.
#[should_panic]
fn local_threadpool_blocking_in_place() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.enable_all()
.build()
.unwrap();
LocalSet::new().block_on(&mut rt, async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let join = spawn_local(async move {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::block_in_place(|| {});
assert!(ON_RT_THREAD.with(|cell| cell.get()));
});
join.await.unwrap();
});
}
#[test]
fn local_threadpool_blocking_run() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.enable_all()
.build()
.unwrap();
LocalSet::new().block_on(&mut rt, async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let join = spawn_local(async move {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::spawn_blocking(|| {
assert!(
!ON_RT_THREAD.with(|cell| cell.get()),
"blocking must not run on the local task set's thread"
);
})
.await
.unwrap();
assert!(ON_RT_THREAD.with(|cell| cell.get()));
});
join.await.unwrap();
});
}
#[test]
fn all_spawns_are_local() {
use futures::future;
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.build()
.unwrap();
LocalSet::new().block_on(&mut rt, async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let handles = (0..128)
.map(|_| {
spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
})
})
.collect::<Vec<_>>();
for joined in future::join_all(handles).await {
joined.unwrap();
}
})
}
#[test]
fn nested_spawn_is_local() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.build()
.unwrap();
LocalSet::new().block_on(&mut rt, async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
})
.await
.unwrap();
})
.await
.unwrap();
})
.await
.unwrap();
})
.await
.unwrap();
})
}
#[test]
fn join_local_future_elsewhere() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.build()
.unwrap();
let local = LocalSet::new();
local.block_on(&mut rt, async move {
let (tx, rx) = crate::sync::oneshot::channel();
let join = spawn_local(async move {
println!("hello world running...");
assert!(
ON_RT_THREAD.with(|cell| cell.get()),
"local task must run on local thread, no matter where it is awaited"
);
rx.await.unwrap();
println!("hello world task done");
"hello world"
});
let join2 = task::spawn(async move {
assert!(
!ON_RT_THREAD.with(|cell| cell.get()),
"spawned task should be on a worker"
);
tx.send(()).expect("task shouldn't have ended yet");
println!("waking up hello world...");
join.await.expect("task should complete successfully");
println!("hello world task joined");
});
join2.await.unwrap()
});
}
#[test]
fn drop_cancels_tasks() {
// This test reproduces issue #1842
let mut rt = runtime::Builder::new()
.enable_time()
.basic_scheduler()
.build()
.unwrap();
let (started_tx, started_rx) = oneshot::channel();
let local = LocalSet::new();
local.spawn_local(async move {
started_tx.send(()).unwrap();
loop {
time::delay_for(Duration::from_secs(3600)).await;
}
});
local.block_on(&mut rt, async {
started_rx.await.unwrap();
});
drop(local);
drop(rt);
}
#[test]
fn drop_cancels_remote_tasks() {
// This test reproduces issue #1885.
use std::sync::mpsc::RecvTimeoutError;
let (done_tx, done_rx) = std::sync::mpsc::channel();
let thread = std::thread::spawn(move || {
let (tx, mut rx) = crate::sync::mpsc::channel::<()>(1024);
let mut rt = runtime::Builder::new()
.enable_time()
.basic_scheduler()
.build()
.expect("building runtime should succeed");
let local = LocalSet::new();
local.spawn_local(async move { while let Some(_) = rx.recv().await {} });
local.block_on(&mut rt, async {
crate::time::delay_for(Duration::from_millis(1)).await;
});
drop(tx);
// This enters an infinite loop if the remote notified tasks are not
// properly cancelled.
drop(local);
// Send a message on the channel so that the test thread can
// determine if we have entered an infinite loop:
done_tx.send(()).unwrap();
});
// Since the failure mode of this test is an infinite loop, rather than
// something we can easily make assertions about, we'll run it in a
// thread. When the test thread finishes, it will send a message on a
// channel to this thread. We'll wait for that message with a fairly
// generous timeout, and if we don't recieve it, we assume the test
// thread has hung.
//
// Note that it should definitely complete in under a minute, but just
// in case CI is slow, we'll give it a long timeout.
match done_rx.recv_timeout(Duration::from_secs(60)) {
Err(RecvTimeoutError::Timeout) => panic!(
"test did not complete within 60 seconds, \
we have (probably) entered an infinite loop!"
),
// Did the test thread panic? We'll find out for sure when we `join`
// with it.
Err(RecvTimeoutError::Disconnected) => {
println!("done_rx dropped, did the test thread panic?");
}
// Test completed successfully!
Ok(()) => {}
}
thread.join().expect("test thread should not panic!")
}
#[test]
fn local_tasks_are_polled_after_tick() {
// Reproduces issues #1899 and #1900
use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
static RX1: AtomicUsize = AtomicUsize::new(0);
static RX2: AtomicUsize = AtomicUsize::new(0);
static EXPECTED: usize = 500;
let (tx, mut rx) = mpsc::unbounded_channel();
let mut rt = runtime::Builder::new()
.basic_scheduler()
.enable_all()
.build()
.unwrap();
let local = LocalSet::new();
local.block_on(&mut rt, async {
let task2 = task::spawn(async move {
// Wait a bit
time::delay_for(Duration::from_millis(100)).await;
let mut oneshots = Vec::with_capacity(EXPECTED);
// Send values
for _ in 0..EXPECTED {
let (oneshot_tx, oneshot_rx) = oneshot::channel();
oneshots.push(oneshot_tx);
tx.send(oneshot_rx).unwrap();
}
time::delay_for(Duration::from_millis(100)).await;
for tx in oneshots.drain(..) {
tx.send(()).unwrap();
}
time::delay_for(Duration::from_millis(300)).await;
let rx1 = RX1.load(SeqCst);
let rx2 = RX2.load(SeqCst);
println!("EXPECT = {}; RX1 = {}; RX2 = {}", EXPECTED, rx1, rx2);
assert_eq!(EXPECTED, rx1);
assert_eq!(EXPECTED, rx2);
});
while let Some(oneshot) = rx.recv().await {
RX1.fetch_add(1, SeqCst);
task::spawn_local(async move {
oneshot.await.unwrap();
RX2.fetch_add(1, SeqCst);
});
}
task2.await.unwrap();
});
}
}
tokio/tests/task_local_set.rs
+420 -9
View File
@@ -1,18 +1,427 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]
use tokio::runtime::Runtime;
use tokio::sync::oneshot;
use tokio::task::{self, LocalSet};
use std::{
cell::Cell,
sync::atomic::{
AtomicBool, AtomicUsize,
Ordering::{self, SeqCst},
},
time::Duration,
};
use tokio::{
runtime::{self, Runtime},
sync::{mpsc, oneshot},
task::{self, LocalSet},
time,
};
#[tokio::test(basic_scheduler)]
async fn local_basic_scheduler() {
LocalSet::new()
.run_until(async {
task::spawn_local(async {}).await.unwrap();
})
.await;
}
#[tokio::test(threaded_scheduler)]
async fn local_threadpool() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
LocalSet::new()
.run_until(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
})
.await
.unwrap();
})
.await;
}
#[tokio::test(threaded_scheduler)]
async fn localset_future_threadpool() {
thread_local! {
static ON_LOCAL_THREAD: Cell<bool> = Cell::new(false);
}
ON_LOCAL_THREAD.with(|cell| cell.set(true));
let local = LocalSet::new();
local.spawn_local(async move {
assert!(ON_LOCAL_THREAD.with(|cell| cell.get()));
});
local.await;
}
#[tokio::test(threaded_scheduler)]
async fn localset_future_timers() {
static RAN1: AtomicBool = AtomicBool::new(false);
static RAN2: AtomicBool = AtomicBool::new(false);
let local = LocalSet::new();
local.spawn_local(async move {
time::delay_for(Duration::from_millis(10)).await;
RAN1.store(true, Ordering::SeqCst);
});
local.spawn_local(async move {
time::delay_for(Duration::from_millis(20)).await;
RAN2.store(true, Ordering::SeqCst);
});
local.await;
assert!(RAN1.load(Ordering::SeqCst));
assert!(RAN2.load(Ordering::SeqCst));
}
#[tokio::test]
async fn localset_future_drives_all_local_futs() {
static RAN1: AtomicBool = AtomicBool::new(false);
static RAN2: AtomicBool = AtomicBool::new(false);
static RAN3: AtomicBool = AtomicBool::new(false);
let local = LocalSet::new();
local.spawn_local(async move {
task::spawn_local(async {
task::yield_now().await;
RAN3.store(true, Ordering::SeqCst);
});
task::yield_now().await;
RAN1.store(true, Ordering::SeqCst);
});
local.spawn_local(async move {
task::yield_now().await;
RAN2.store(true, Ordering::SeqCst);
});
local.await;
assert!(RAN1.load(Ordering::SeqCst));
assert!(RAN2.load(Ordering::SeqCst));
assert!(RAN3.load(Ordering::SeqCst));
}
#[tokio::test(threaded_scheduler)]
async fn local_threadpool_timer() {
// This test ensures that runtime services like the timer are properly
// set for the local task set.
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
LocalSet::new()
.run_until(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let join = task::spawn_local(async move {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
time::delay_for(Duration::from_millis(10)).await;
assert!(ON_RT_THREAD.with(|cell| cell.get()));
});
join.await.unwrap();
})
.await;
}
#[test]
fn acquire_mutex_in_drop() {
// This will panic, since the thread that calls `block_on` cannot use
// in-place blocking inside of `block_on`.
#[should_panic]
fn local_threadpool_blocking_in_place() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.enable_all()
.build()
.unwrap();
LocalSet::new().block_on(&mut rt, async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let join = task::spawn_local(async move {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::block_in_place(|| {});
assert!(ON_RT_THREAD.with(|cell| cell.get()));
});
join.await.unwrap();
});
}
#[tokio::test(threaded_scheduler)]
async fn local_threadpool_blocking_run() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
LocalSet::new()
.run_until(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let join = task::spawn_local(async move {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::spawn_blocking(|| {
assert!(
!ON_RT_THREAD.with(|cell| cell.get()),
"blocking must not run on the local task set's thread"
);
})
.await
.unwrap();
assert!(ON_RT_THREAD.with(|cell| cell.get()));
});
join.await.unwrap();
})
.await;
}
#[tokio::test(threaded_scheduler)]
async fn all_spawns_are_local() {
use futures::future;
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
LocalSet::new()
.run_until(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
let handles = (0..128)
.map(|_| {
task::spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
})
})
.collect::<Vec<_>>();
for joined in future::join_all(handles).await {
joined.unwrap();
}
})
.await;
}
#[tokio::test(threaded_scheduler)]
async fn nested_spawn_is_local() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
LocalSet::new()
.run_until(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
task::spawn_local(async {
assert!(ON_RT_THREAD.with(|cell| cell.get()));
})
.await
.unwrap();
})
.await
.unwrap();
})
.await
.unwrap();
})
.await
.unwrap();
})
.await;
}
#[test]
fn join_local_future_elsewhere() {
thread_local! {
static ON_RT_THREAD: Cell<bool> = Cell::new(false);
}
ON_RT_THREAD.with(|cell| cell.set(true));
let mut rt = runtime::Builder::new()
.threaded_scheduler()
.build()
.unwrap();
let local = LocalSet::new();
local.block_on(&mut rt, async move {
let (tx, rx) = oneshot::channel();
let join = task::spawn_local(async move {
println!("hello world running...");
assert!(
ON_RT_THREAD.with(|cell| cell.get()),
"local task must run on local thread, no matter where it is awaited"
);
rx.await.unwrap();
println!("hello world task done");
"hello world"
});
let join2 = task::spawn(async move {
assert!(
!ON_RT_THREAD.with(|cell| cell.get()),
"spawned task should be on a worker"
);
tx.send(()).expect("task shouldn't have ended yet");
println!("waking up hello world...");
join.await.expect("task should complete successfully");
println!("hello world task joined");
});
join2.await.unwrap()
});
}
#[test]
fn drop_cancels_tasks() {
// This test reproduces issue #1842
let mut rt = rt();
let (started_tx, started_rx) = oneshot::channel();
let local = LocalSet::new();
local.spawn_local(async move {
started_tx.send(()).unwrap();
loop {
time::delay_for(Duration::from_secs(3600)).await;
}
});
local.block_on(&mut rt, async {
started_rx.await.unwrap();
});
drop(local);
drop(rt);
}
#[test]
fn drop_cancels_remote_tasks() {
// This test reproduces issue #1885.
use std::sync::mpsc::RecvTimeoutError;
let (done_tx, done_rx) = std::sync::mpsc::channel();
let thread = std::thread::spawn(move || {
let (tx, mut rx) = mpsc::channel::<()>(1024);
let mut rt = rt();
let local = LocalSet::new();
local.spawn_local(async move { while let Some(_) = rx.recv().await {} });
local.block_on(&mut rt, async {
time::delay_for(Duration::from_millis(1)).await;
});
drop(tx);
// This enters an infinite loop if the remote notified tasks are not
// properly cancelled.
drop(local);
// Send a message on the channel so that the test thread can
// determine if we have entered an infinite loop:
done_tx.send(()).unwrap();
});
// Since the failure mode of this test is an infinite loop, rather than
// something we can easily make assertions about, we'll run it in a
// thread. When the test thread finishes, it will send a message on a
// channel to this thread. We'll wait for that message with a fairly
// generous timeout, and if we don't recieve it, we assume the test
// thread has hung.
//
// Note that it should definitely complete in under a minute, but just
// in case CI is slow, we'll give it a long timeout.
match done_rx.recv_timeout(Duration::from_secs(60)) {
Err(RecvTimeoutError::Timeout) => panic!(
"test did not complete within 60 seconds, \
we have (probably) entered an infinite loop!"
),
// Did the test thread panic? We'll find out for sure when we `join`
// with it.
Err(RecvTimeoutError::Disconnected) => {
println!("done_rx dropped, did the test thread panic?");
}
// Test completed successfully!
Ok(()) => {}
}
thread.join().expect("test thread should not panic!")
}
#[tokio::test]
async fn local_tasks_are_polled_after_tick() {
// Reproduces issues #1899 and #1900
static RX1: AtomicUsize = AtomicUsize::new(0);
static RX2: AtomicUsize = AtomicUsize::new(0);
static EXPECTED: usize = 500;
let (tx, mut rx) = mpsc::unbounded_channel();
let local = LocalSet::new();
local
.run_until(async {
let task2 = task::spawn(async move {
// Wait a bit
time::delay_for(Duration::from_millis(100)).await;
let mut oneshots = Vec::with_capacity(EXPECTED);
// Send values
for _ in 0..EXPECTED {
let (oneshot_tx, oneshot_rx) = oneshot::channel();
oneshots.push(oneshot_tx);
tx.send(oneshot_rx).unwrap();
}
time::delay_for(Duration::from_millis(100)).await;
for tx in oneshots.drain(..) {
tx.send(()).unwrap();
}
time::delay_for(Duration::from_millis(300)).await;
let rx1 = RX1.load(SeqCst);
let rx2 = RX2.load(SeqCst);
println!("EXPECT = {}; RX1 = {}; RX2 = {}", EXPECTED, rx1, rx2);
assert_eq!(EXPECTED, rx1);
assert_eq!(EXPECTED, rx2);
});
while let Some(oneshot) = rx.recv().await {
RX1.fetch_add(1, SeqCst);
task::spawn_local(async move {
oneshot.await.unwrap();
RX2.fetch_add(1, SeqCst);
});
}
task2.await.unwrap();
})
.await;
}
#[tokio::test]
async fn acquire_mutex_in_drop() {
use futures::future::pending;
let (tx1, rx1) = oneshot::channel();
let (tx2, rx2) = oneshot::channel();
let mut rt = rt();
let local = LocalSet::new();
local.spawn_local(async move {
@@ -33,9 +442,11 @@ fn acquire_mutex_in_drop() {
});
// Tick the loop
local.block_on(&mut rt, async {
task::yield_now().await;
});
local
.run_until(async {
task::yield_now().await;
})
.await;
// Drop the LocalSet
drop(local);