use crate::RuntimeError;
use crate::runtime_state::{self as runtime, LifecycleSignals, RuntimeInner, ScopeSlot};
use futures_util::FutureExt;
use std::future::{Future, IntoFuture};
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Poll};
pub async fn on_shutdown() {
match runtime::try_current_runtime() {
Some(runtime) => observe_shutdown(runtime).await,
None => std::future::pending().await,
}
}
async fn observe_shutdown(runtime: Arc<RuntimeInner>) {
runtime
.shutdown_signal()
.wait_observed(|| {
runtime.pause_test_schedule(
crate::runtime_test_support::RuntimeCheckpoint::ShutdownWaitRegistered,
);
std::future::ready(())
})
.await;
}
impl<T> std::fmt::Debug for JoinHandle<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("JoinHandle")
.field("refused", &matches!(self.source, JoinSource::Refused(_)))
.finish()
}
}
enum JoinSource<T> {
Task {
rx: std::sync::mpsc::Receiver<Result<T, RuntimeError>>,
cancel: Arc<AtomicBool>,
cancel_tx: crossbeam_channel::Sender<()>,
},
Refused(RuntimeError),
}
pub struct JoinHandle<T> {
source: JoinSource<T>,
}
fn channel_closed() -> RuntimeError {
RuntimeError::TaskPanicked("task channel closed".into())
}
fn report_dropped_result<E>(sent: Result<(), E>, flavor: &'static str) {
match sent {
Ok(()) => {}
Err(_) => tracing::debug!(flavor, "task result dropped: handle was released"),
}
}
fn recv_task_result<T>(
rx: std::sync::mpsc::Receiver<Result<T, RuntimeError>>,
) -> Result<T, RuntimeError> {
match rx.recv() {
Ok(result) => result,
Err(_) => Err(channel_closed()),
}
}
pub(crate) fn catch_panic<F, T>(f: F) -> Result<T, RuntimeError>
where
F: FnOnce() -> T,
{
std::panic::catch_unwind(AssertUnwindSafe(f)).map_err(panic_to_error)
}
pub(crate) async fn catch_panic_async<F>(f: F) -> Result<F::Output, RuntimeError>
where
F: Future,
{
AssertUnwindSafe(f)
.catch_unwind()
.await
.map_err(panic_to_error)
}
pub(crate) fn panic_message(payload: &(dyn std::any::Any + Send)) -> Option<&str> {
match payload.downcast_ref::<&'static str>() {
Some(text) => Some(text),
None => payload.downcast_ref::<String>().map(String::as_str),
}
}
pub(crate) fn panic_to_error(payload: Box<dyn std::any::Any + Send>) -> RuntimeError {
let message = panic_message(payload.as_ref()).unwrap_or("unknown panic");
RuntimeError::TaskPanicked(message.into())
}
impl<T> JoinHandle<T> {
pub fn cancel(&self) {
match &self.source {
JoinSource::Task {
cancel, cancel_tx, ..
} => {
cancel.store(true, Ordering::Release);
let _ = cancel_tx.try_send(());
}
JoinSource::Refused(_) => {}
}
}
pub fn join(self) -> Result<T, RuntimeError> {
match self.source {
JoinSource::Refused(error) => Err(error),
JoinSource::Task { rx, .. } => recv_task_result(rx),
}
}
fn refused(error: RuntimeError) -> Self {
Self {
source: JoinSource::Refused(error),
}
}
}
pub fn spawn<F, T>(f: F) -> JoinHandle<T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
match admit_blocking_child() {
Ok((rt, slot)) => spawn_admitted(f, rt, slot),
Err(error) => JoinHandle::refused(error),
}
}
fn admit_blocking_child() -> Result<(Arc<RuntimeInner>, ScopeSlot), RuntimeError> {
let rt = runtime::runtime_context()?;
let slot = rt.admit_blocking()?;
Ok((rt, slot))
}
fn spawn_admitted<F, T>(f: F, rt: Arc<RuntimeInner>, slot: ScopeSlot) -> JoinHandle<T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let (tx, rx) = std::sync::mpsc::sync_channel::<Result<T, RuntimeError>>(1);
let cancel = Arc::new(AtomicBool::new(false));
let cancel_child = Arc::clone(&cancel);
let (cancel_tx, cancel_rx) = crossbeam_channel::bounded::<()>(1);
let launched = launch_on_admitting_executor(rt, slot, move || {
deliver_task_result(f, tx, cancel_child, cancel_rx)
});
match launched {
Ok(()) => JoinHandle {
source: JoinSource::Task {
rx,
cancel,
cancel_tx,
},
},
Err(error) => JoinHandle::refused(error),
}
}
fn launch_on_admitting_executor<F>(
rt: Arc<RuntimeInner>,
slot: ScopeSlot,
body: F,
) -> Result<(), RuntimeError>
where
F: FnOnce() + Send + 'static,
{
let executor = rt.executor()?.clone();
drop(executor.spawn_blocking(move || run_in_spawner_context(rt, slot, body)));
Ok(())
}
fn run_in_spawner_context<F>(rt: Arc<RuntimeInner>, slot: ScopeSlot, body: F)
where
F: FnOnce(),
{
let runtime_guard = runtime::install_runtime(rt);
body();
drop(slot);
drop(runtime_guard);
}
fn deliver_task_result<F, T>(
f: F,
tx: std::sync::mpsc::SyncSender<Result<T, RuntimeError>>,
cancel: Arc<AtomicBool>,
cancel_rx: crossbeam_channel::Receiver<()>,
) where
F: FnOnce() -> T,
{
let cancel_guard = runtime::install_cancel_context(cancel, cancel_rx);
let mapped = catch_panic(f);
report_dropped_result(tx.send(mapped), "blocking");
drop(cancel_guard);
}
impl<T> std::fmt::Debug for AsyncJoinHandle<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AsyncJoinHandle")
.field(
"refused",
&matches!(self.source, AsyncJoinSource::Refused(_)),
)
.finish()
}
}
enum AsyncJoinSource<T> {
Task(tokio::sync::oneshot::Receiver<Result<T, RuntimeError>>),
Refused(Option<RuntimeError>),
}
pub struct AsyncJoinHandle<T> {
source: AsyncJoinSource<T>,
cancel: Option<Arc<tokio::sync::Notify>>,
}
impl<T> AsyncJoinHandle<T> {
pub fn cancel(&self) {
if let Some(cancel) = &self.cancel {
cancel.notify_one();
}
}
fn refused(error: RuntimeError) -> Self {
Self {
source: AsyncJoinSource::Refused(Some(error)),
cancel: None,
}
}
}
pub struct AsyncJoinFuture<T> {
source: AsyncJoinSource<T>,
}
impl<T> Future for AsyncJoinFuture<T> {
type Output = Result<T, RuntimeError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let source = &mut self.get_mut().source;
match source {
AsyncJoinSource::Task(rx) => {
let outcome = poll_task_result(Pin::new(rx), cx);
spend_delivered_task(source, outcome)
}
AsyncJoinSource::Refused(error) => take_refusal(error.take()),
}
}
}
fn spend_delivered_task<T>(
source: &mut AsyncJoinSource<T>,
outcome: Poll<Result<T, RuntimeError>>,
) -> Poll<Result<T, RuntimeError>> {
match outcome {
Poll::Ready(result) => {
*source = AsyncJoinSource::Refused(None);
Poll::Ready(result)
}
Poll::Pending => Poll::Pending,
}
}
fn poll_task_result<T>(
rx: Pin<&mut tokio::sync::oneshot::Receiver<Result<T, RuntimeError>>>,
cx: &mut Context<'_>,
) -> Poll<Result<T, RuntimeError>> {
match rx.poll(cx) {
Poll::Ready(Ok(result)) => Poll::Ready(result),
Poll::Ready(Err(_)) => Poll::Ready(Err(channel_closed())),
Poll::Pending => Poll::Pending,
}
}
fn take_refusal<T>(error: Option<RuntimeError>) -> Poll<Result<T, RuntimeError>> {
match error {
Some(error) => Poll::Ready(Err(error)),
None => Poll::Pending,
}
}
impl<T> AsyncJoinFuture<T> {
pub(crate) fn closed() -> Self {
Self {
source: AsyncJoinSource::Refused(Some(channel_closed())),
}
}
}
impl<T> IntoFuture for AsyncJoinHandle<T> {
type Output = Result<T, RuntimeError>;
type IntoFuture = AsyncJoinFuture<T>;
fn into_future(self) -> Self::IntoFuture {
AsyncJoinFuture {
source: self.source,
}
}
}
pub fn spawn_async<F, T>(future: F) -> AsyncJoinHandle<T>
where
F: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
let rt = match runtime::runtime_context() {
Ok(rt) => rt,
Err(error) => return AsyncJoinHandle::refused(error),
};
let (tx, rx) = tokio::sync::oneshot::channel();
let cancel = Arc::new(tokio::sync::Notify::new());
let body = run_async_task(future, tx, Arc::clone(&cancel));
match rt.admit_async(body) {
Ok(()) => AsyncJoinHandle {
source: AsyncJoinSource::Task(rx),
cancel: Some(cancel),
},
Err(error) => AsyncJoinHandle::refused(error),
}
}
pub(crate) fn admit_signalled_on<B, Fut>(
runtime: &Arc<RuntimeInner>,
build: B,
) -> Result<(), RuntimeError>
where
B: FnOnce(LifecycleSignals) -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
let signals = LifecycleSignals::from_runtime(runtime);
runtime.admit_internal_async(build(signals))
}
pub(crate) fn admit_signalled_loop<B, Fut>(build: B) -> Result<(), RuntimeError>
where
B: FnOnce(LifecycleSignals) -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
admit_signalled_on(&runtime::runtime_context()?, build)
}
pub(crate) fn admit_signalled_subsystem_on<B, Fut>(
runtime: &Arc<RuntimeInner>,
subsystem: &str,
build: B,
) -> Result<(), RuntimeError>
where
B: FnOnce(LifecycleSignals) -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
report_subsystem_outcome(subsystem, admit_signalled_on(runtime, build))
}
fn report_subsystem_outcome(
subsystem: &str,
outcome: Result<(), RuntimeError>,
) -> Result<(), RuntimeError> {
if let Err(error) = outcome.as_ref() {
record_subsystem_refusal(subsystem, error);
}
outcome
}
fn record_subsystem_refusal(subsystem: &str, error: &RuntimeError) {
tracing::error!(
subsystem,
%error,
"root scope refused a Camber-owned background subsystem"
);
}
pub(crate) fn spawn_internal_blocking<F>(producer: &'static str, path: &str, f: F)
where
F: FnOnce() + Send + 'static,
{
match admit_blocking_child() {
Ok((rt, slot)) => spawn_admitted_producer(producer, path, rt, slot, f),
Err(RuntimeError::NoRuntime) => detach_producer(producer, path, f),
Err(error) => record_producer_refusal(producer, path, &error),
}
}
fn detach_producer<F>(producer: &'static str, path: &str, f: F)
where
F: FnOnce() + Send + 'static,
{
match tokio::runtime::Handle::try_current() {
Ok(executor) => drop(executor.spawn_blocking(move || run_producer(producer, f))),
Err(_) => record_producer_refusal(producer, path, &RuntimeError::NoRuntime),
}
}
fn spawn_admitted_producer<F>(
producer: &'static str,
path: &str,
rt: Arc<RuntimeInner>,
slot: ScopeSlot,
f: F,
) where
F: FnOnce() + Send + 'static,
{
if let Err(error) = launch_on_admitting_executor(rt, slot, move || run_producer(producer, f)) {
record_producer_refusal(producer, path, &error);
}
}
fn run_producer<F>(producer: &'static str, f: F)
where
F: FnOnce() + Send + 'static,
{
match catch_panic(f) {
Ok(()) => {}
Err(error) => tracing::error!(producer, %error, "per-response producer panicked"),
}
}
fn record_producer_refusal(producer: &'static str, path: &str, error: &RuntimeError) {
tracing::debug!(
producer,
path,
%error,
"root scope refused a per-response producer"
);
}
pub(crate) fn block_in_place<F, T>(f: F) -> T
where
F: FnOnce() -> T,
{
match tokio::runtime::Handle::try_current().map(|handle| handle.runtime_flavor()) {
Ok(tokio::runtime::RuntimeFlavor::MultiThread) => tokio::task::block_in_place(f),
_ => f(),
}
}
pub async fn race<A, B, T>(a: A, b: B) -> T
where
A: Future<Output = T>,
B: Future<Output = T>,
{
tokio::select! {
biased;
result = a => result,
result = b => result,
}
}
pub async fn race_all<F, T>(futures: Vec<F>) -> Result<T, RuntimeError>
where
F: Future<Output = T> + Send,
{
match futures.is_empty() {
true => Err(RuntimeError::InvalidArgument(
"race_all called with empty futures list".into(),
)),
false => {
let pinned: Vec<Pin<Box<F>>> = futures.into_iter().map(Box::pin).collect();
let (result, _, _) = futures_util::future::select_all(pinned).await;
Ok(result)
}
}
}
async fn run_async_task<F, T>(
future: F,
tx: tokio::sync::oneshot::Sender<Result<T, RuntimeError>>,
cancel: Arc<tokio::sync::Notify>,
) where
F: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
let result = tokio::select! {
biased;
() = cancel.notified() => Err(RuntimeError::Cancelled),
outcome = catch_panic_async(future) => outcome,
};
report_dropped_result(tx.send(result), "async");
}