You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
290 lines
7.5 KiB
Rust
290 lines
7.5 KiB
Rust
5 years ago
|
use std::cell::{Cell, UnsafeCell};
|
||
|
use std::fmt::Arguments;
|
||
|
use std::future::Future;
|
||
|
use std::io;
|
||
|
use std::mem;
|
||
|
use std::panic::{self, AssertUnwindSafe};
|
||
|
use std::pin::Pin;
|
||
|
use std::ptr;
|
||
|
use std::thread;
|
||
|
|
||
|
use crossbeam::channel::{unbounded, Sender};
|
||
|
use futures::prelude::*;
|
||
|
use lazy_static::lazy_static;
|
||
|
|
||
|
use super::task;
|
||
|
use super::{JoinHandle, Task};
|
||
|
|
||
|
/// Returns a handle to the current task.
|
||
|
///
|
||
|
/// # Panics
|
||
|
///
|
||
|
/// This function will panic if not called within the context of a task created by [`block_on`],
|
||
|
/// [`spawn`], or [`Builder::spawn`].
|
||
|
///
|
||
|
/// [`block_on`]: fn.block_on.html
|
||
|
/// [`spawn`]: fn.spawn.html
|
||
|
/// [`Builder::spawn`]: struct.Builder.html#method.spawn
|
||
|
///
|
||
|
/// # Examples
|
||
|
///
|
||
|
/// ```
|
||
|
/// # #![feature(async_await)]
|
||
|
/// use async_std::task::current;
|
||
|
///
|
||
|
/// # async_std::task::block_on(async {
|
||
|
/// println!("The name of this task is {:?}", current().name());
|
||
|
/// # });
|
||
|
pub fn current() -> Task {
|
||
|
get_task(|task| task.clone()).expect("`task::current()` called outside the context of a task")
|
||
|
}
|
||
|
|
||
|
/// Spawns a task.
|
||
|
///
|
||
|
/// This function is similar to [`std::thread::spawn`], except it spawns an asynchronous task.
|
||
|
///
|
||
|
/// [`std::thread`]: https://doc.rust-lang.org/std/thread/fn.spawn.html
|
||
|
///
|
||
|
/// # Examples
|
||
|
///
|
||
|
/// ```
|
||
|
/// # #![feature(async_await)]
|
||
|
/// use async_std::task;
|
||
|
///
|
||
|
/// # async_std::task::block_on(async {
|
||
|
/// let handle = task::spawn(async {
|
||
|
/// 1 + 2
|
||
|
/// });
|
||
|
///
|
||
|
/// assert_eq!(handle.await, 3);
|
||
|
/// # });
|
||
|
/// ```
|
||
|
pub fn spawn<F, T>(future: F) -> JoinHandle<T>
|
||
|
where
|
||
|
F: Future<Output = T> + Send + 'static,
|
||
|
T: Send + 'static,
|
||
|
{
|
||
|
spawn_with_builder(Builder::new(), future, "spawn")
|
||
|
}
|
||
|
|
||
|
/// Spawns a task and blocks the current thread on its result.
|
||
|
///
|
||
|
/// Calling this function is similar to [spawning] a thread and immediately [joining] it, except an
|
||
|
/// asynchronous task will be spawned.
|
||
|
///
|
||
|
/// [spawning]: https://doc.rust-lang.org/std/thread/fn.spawn.html
|
||
|
/// [joining]: https://doc.rust-lang.org/std/thread/struct.JoinHandle.html#method.join
|
||
|
///
|
||
|
/// # Examples
|
||
|
///
|
||
|
/// ```no_run
|
||
|
/// # #![feature(async_await)]
|
||
|
/// use async_std::task;
|
||
|
///
|
||
|
/// fn main() {
|
||
|
/// task::block_on(async {
|
||
|
/// println!("Hello, world!");
|
||
|
/// })
|
||
|
/// }
|
||
|
/// ```
|
||
|
pub fn block_on<F, T>(future: F) -> T
|
||
|
where
|
||
|
F: Future<Output = T> + Send,
|
||
|
T: Send,
|
||
|
{
|
||
|
unsafe {
|
||
|
// A place on the stack where the result will be stored.
|
||
|
let out = &mut UnsafeCell::new(None);
|
||
|
|
||
|
// Wrap the future into one that stores the result into `out`.
|
||
|
let future = {
|
||
|
let out = out.get();
|
||
|
async move {
|
||
|
let v = AssertUnwindSafe(future).catch_unwind().await;
|
||
|
*out = Some(v);
|
||
|
}
|
||
|
};
|
||
|
|
||
|
// Pin the future onto the stack.
|
||
|
futures::pin_mut!(future);
|
||
|
|
||
|
// Transmute the future into one that is static and sendable.
|
||
|
let future = mem::transmute::<
|
||
|
Pin<&mut dyn Future<Output = ()>>,
|
||
|
Pin<&'static mut (dyn Future<Output = ()> + Send)>,
|
||
|
>(future);
|
||
|
|
||
|
// Spawn the future and wait for it to complete.
|
||
|
futures::executor::block_on(spawn_with_builder(Builder::new(), future, "block_on"));
|
||
|
|
||
|
// Take out the result.
|
||
|
match (*out.get()).take().unwrap() {
|
||
|
Ok(v) => v,
|
||
|
Err(err) => panic::resume_unwind(err),
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
/// Task builder that configures the settings of a new task.
|
||
|
#[derive(Debug)]
|
||
|
pub struct Builder {
|
||
|
pub(crate) name: Option<String>,
|
||
|
}
|
||
|
|
||
|
impl Builder {
|
||
|
/// Creates a new builder.
|
||
|
pub fn new() -> Builder {
|
||
|
Builder { name: None }
|
||
|
}
|
||
|
|
||
|
/// Configures the name of the task.
|
||
|
pub fn name(mut self, name: String) -> Builder {
|
||
|
self.name = Some(name);
|
||
|
self
|
||
|
}
|
||
|
|
||
|
/// Spawns a task with the configured settings.
|
||
|
pub fn spawn<F, T>(self, future: F) -> io::Result<JoinHandle<T>>
|
||
|
where
|
||
|
F: Future<Output = T> + Send + 'static,
|
||
|
T: Send + 'static,
|
||
|
{
|
||
|
Ok(spawn_with_builder(self, future, "spawn"))
|
||
|
}
|
||
|
}
|
||
|
|
||
|
pub(crate) fn spawn_with_builder<F, T>(
|
||
|
builder: Builder,
|
||
|
future: F,
|
||
|
fn_name: &'static str,
|
||
|
) -> JoinHandle<T>
|
||
|
where
|
||
|
F: Future<Output = T> + Send + 'static,
|
||
|
T: Send + 'static,
|
||
|
{
|
||
|
let Builder { name } = builder;
|
||
|
|
||
|
type Job = async_task::Task<task::Tag>;
|
||
|
|
||
|
lazy_static! {
|
||
|
static ref QUEUE: Sender<Job> = {
|
||
|
let (sender, receiver) = unbounded::<Job>();
|
||
|
|
||
|
for _ in 0..num_cpus::get().max(1) {
|
||
|
let receiver = receiver.clone();
|
||
|
thread::Builder::new()
|
||
|
.name("async-task-driver".to_string())
|
||
|
.spawn(|| {
|
||
|
TAG.with(|tag| {
|
||
|
for job in receiver {
|
||
|
tag.set(job.tag());
|
||
|
abort_on_panic(|| job.run());
|
||
|
tag.set(ptr::null());
|
||
|
}
|
||
|
});
|
||
|
})
|
||
|
.expect("cannot start a thread driving tasks");
|
||
|
}
|
||
|
|
||
|
sender
|
||
|
};
|
||
|
}
|
||
|
|
||
|
let tag = task::Tag::new(name);
|
||
|
let schedule = |job| QUEUE.send(job).unwrap();
|
||
|
|
||
|
let child_id = tag.task_id().as_u64();
|
||
|
let parent_id = get_task(|t| t.id().as_u64()).unwrap_or(0);
|
||
|
print(
|
||
|
format_args!("{}", fn_name),
|
||
|
LogData {
|
||
|
parent_id,
|
||
|
child_id,
|
||
|
},
|
||
|
);
|
||
|
|
||
|
// Wrap the future into one that drops task-local variables on exit.
|
||
|
let future = async move {
|
||
|
let res = future.await;
|
||
|
|
||
|
// Abort on panic because thread-local variables behave the same way.
|
||
|
abort_on_panic(|| get_task(|task| task.metadata().local_map.clear()));
|
||
|
|
||
|
print(
|
||
|
format_args!("{} completed", fn_name),
|
||
|
LogData {
|
||
|
parent_id,
|
||
|
child_id,
|
||
|
},
|
||
|
);
|
||
|
res
|
||
|
};
|
||
|
|
||
|
let (task, handle) = async_task::spawn(future, schedule, tag);
|
||
|
task.schedule();
|
||
|
JoinHandle::new(handle)
|
||
|
}
|
||
|
|
||
|
thread_local! {
|
||
|
static TAG: Cell<*const task::Tag> = Cell::new(ptr::null_mut());
|
||
|
}
|
||
|
|
||
|
pub(crate) fn get_task<F: FnOnce(&Task) -> R, R>(f: F) -> Option<R> {
|
||
|
let res = TAG.try_with(|tag| unsafe { tag.get().as_ref().map(task::Tag::task).map(f) });
|
||
|
|
||
|
match res {
|
||
|
Ok(Some(val)) => Some(val),
|
||
|
Ok(None) | Err(_) => None,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
/// Calls a function and aborts if it panics.
|
||
|
///
|
||
|
/// This is useful in unsafe code where we can't recover from panics.
|
||
|
#[inline]
|
||
|
fn abort_on_panic<T>(f: impl FnOnce() -> T) -> T {
|
||
|
struct Bomb;
|
||
|
|
||
|
impl Drop for Bomb {
|
||
|
fn drop(&mut self) {
|
||
|
std::process::abort();
|
||
|
}
|
||
|
}
|
||
|
|
||
|
let bomb = Bomb;
|
||
|
let t = f();
|
||
|
mem::forget(bomb);
|
||
|
t
|
||
|
}
|
||
|
|
||
|
/// This struct only exists because kv logging isn't supported from the macros right now.
|
||
|
struct LogData {
|
||
|
parent_id: u64,
|
||
|
child_id: u64,
|
||
|
}
|
||
|
|
||
|
impl<'a> log::kv::Source for LogData {
|
||
|
fn visit<'kvs>(
|
||
|
&'kvs self,
|
||
|
visitor: &mut dyn log::kv::Visitor<'kvs>,
|
||
|
) -> Result<(), log::kv::Error> {
|
||
|
visitor.visit_pair("parent_id".into(), self.parent_id.into())?;
|
||
|
visitor.visit_pair("child_id".into(), self.child_id.into())?;
|
||
|
Ok(())
|
||
|
}
|
||
|
}
|
||
|
|
||
|
fn print(msg: Arguments<'_>, key_values: impl log::kv::Source) {
|
||
|
log::logger().log(
|
||
|
&log::Record::builder()
|
||
|
.args(msg)
|
||
|
.key_values(&key_values)
|
||
|
.level(log::Level::Trace)
|
||
|
.target(module_path!())
|
||
|
.module_path(Some(module_path!()))
|
||
|
.file(Some(file!()))
|
||
|
.line(Some(line!()))
|
||
|
.build(),
|
||
|
);
|
||
|
}
|