use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use pin_project_lite::pin_project; use crate::io; use crate::task::{JoinHandle, Task, TaskLocalsWrapper}; /// Task builder that configures the settings of a new task. #[derive(Debug, Default)] pub struct Builder { pub(crate) name: Option, } impl Builder { /// Creates a new builder. #[inline] pub fn new() -> Builder { Builder { name: None } } /// Configures the name of the task. #[inline] pub fn name(mut self, name: String) -> Builder { self.name = Some(name); self } fn build(self, future: F) -> SupportTaskLocals where F: Future, { let name = self.name.map(Arc::new); // Create a new task handle. let task = Task::new(name); #[cfg(not(target_os = "unknown"))] once_cell::sync::Lazy::force(&crate::rt::RUNTIME); let tag = TaskLocalsWrapper::new(task); SupportTaskLocals { tag, future } } /// Spawns a task with the configured settings. #[cfg(not(target_os = "unknown"))] pub fn spawn(self, future: F) -> io::Result> where F: Future + Send + 'static, T: Send + 'static, { let wrapped = self.build(future); kv_log_macro::trace!("spawn", { task_id: wrapped.tag.id().0, parent_task_id: TaskLocalsWrapper::get_current(|t| t.id().0).unwrap_or(0), }); let task = wrapped.tag.task().clone(); let handle = crate::task::executor::spawn(wrapped); Ok(JoinHandle::new(handle, task)) } /// Spawns a task locally with the configured settings. #[cfg(all(not(target_os = "unknown"), feature = "unstable"))] pub fn local(self, future: F) -> io::Result> where F: Future + 'static, T: 'static, { let wrapped = self.build(future); kv_log_macro::trace!("spawn_local", { task_id: wrapped.tag.id().0, parent_task_id: TaskLocalsWrapper::get_current(|t| t.id().0).unwrap_or(0), }); let task = wrapped.tag.task().clone(); let handle = crate::task::executor::local(wrapped); Ok(JoinHandle::new(handle, task)) } /// Spawns a task locally with the configured settings. #[cfg(all(target_arch = "wasm32", feature = "unstable"))] pub fn local(self, future: F) -> io::Result> where F: Future + 'static, T: 'static, { use futures_channel::oneshot::channel; let (sender, receiver) = channel(); let wrapped = self.build(async move { let res = future.await; let _ = sender.send(res); }); kv_log_macro::trace!("spawn_local", { task_id: wrapped.tag.id().0, parent_task_id: TaskLocalsWrapper::get_current(|t| t.id().0).unwrap_or(0), }); let task = wrapped.tag.task().clone(); wasm_bindgen_futures::spawn_local(wrapped); Ok(JoinHandle::new(receiver, task)) } /// Spawns a task locally with the configured settings. #[cfg(all(target_arch = "wasm32", not(feature = "unstable")))] pub(crate) fn local(self, future: F) -> io::Result> where F: Future + 'static, T: 'static, { use futures_channel::oneshot::channel; let (sender, receiver) = channel(); let wrapped = self.build(async move { let res = future.await; let _ = sender.send(res); }); kv_log_macro::trace!("spawn_local", { task_id: wrapped.tag.id().0, parent_task_id: TaskLocalsWrapper::get_current(|t| t.id().0).unwrap_or(0), }); let task = wrapped.tag.task().clone(); wasm_bindgen_futures::spawn_local(wrapped); Ok(JoinHandle::new(receiver, task)) } /// Spawns a task with the configured settings, blocking on its execution. #[cfg(not(target_os = "unknown"))] pub fn blocking(self, future: F) -> T where F: Future, { use std::cell::Cell; let wrapped = self.build(future); // Log this `block_on` operation. kv_log_macro::trace!("block_on", { task_id: wrapped.tag.id().0, parent_task_id: TaskLocalsWrapper::get_current(|t| t.id().0).unwrap_or(0), }); thread_local! { /// Tracks the number of nested block_on calls. static NUM_NESTED_BLOCKING: Cell = Cell::new(0); } // Run the future as a task. NUM_NESTED_BLOCKING.with(|num_nested_blocking| { let count = num_nested_blocking.get(); let should_run = count == 0; // increase the count num_nested_blocking.replace(count + 1); unsafe { TaskLocalsWrapper::set_current(&wrapped.tag, || { let res = if should_run { // The first call should run the executor crate::task::executor::run(wrapped) } else { futures_lite::future::block_on(wrapped) }; num_nested_blocking.replace(num_nested_blocking.get() - 1); res }) } }) } } pin_project! { /// Wrapper to add support for task locals. struct SupportTaskLocals { tag: TaskLocalsWrapper, #[pin] future: F, } } impl Future for SupportTaskLocals { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { unsafe { TaskLocalsWrapper::set_current(&self.tag, || { let this = self.project(); this.future.poll(cx) }) } } }