arty 0.6.0

Thread-aware async runtime that keeps each task on one worker thread.
Documentation
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use std::panic::{AssertUnwindSafe, catch_unwind};

use arty_executor::{ShutdownCause, TaskCompletion};

use crate::runtime::ShutdownStart;
use crate::runtime::worker::blocking::PanicPolicy;
use crate::task::handle::{JoinHandle, JoinHandlePool, Sender};
use crate::task::preparation::discard_panicking_drop;

pub(crate) fn prepare_blocking<F, R>(
    result_pool: &JoinHandlePool,
    body: F,
    panic_policy: PanicPolicy,
) -> (impl FnOnce() + Send + 'static, JoinHandle<R>)
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    let (handle, result_tx) = result_pool.handle();
    let task = BlockingTask {
        pending: Some((body, result_tx)),
        panic_policy,
    };
    (move || task.run(), handle)
}

struct BlockingTask<F, R: Send + 'static> {
    pending: Option<(F, Sender<R>)>,
    panic_policy: PanicPolicy,
}

impl<F, R> BlockingTask<F, R>
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    fn run(mut self) {
        let (body, result_tx) = self.pending.take().expect("a blocking task is consumed exactly once");
        match catch_unwind(AssertUnwindSafe(body)) {
            Ok(result) => {
                result_tx.completed(result);
            }
            Err(panic) => {
                if let PanicPolicy::Shutdown(waker) = &self.panic_policy
                    && result_tx.shutdown().begin(ShutdownCause::panic()) == ShutdownStart::Started
                {
                    waker.wake_by_ref();
                }
                result_tx.panicked(panic, None);
            }
        }
    }
}

impl<F, R: Send + 'static> Drop for BlockingTask<F, R> {
    fn drop(&mut self) {
        let Some((body, result_tx)) = self.pending.take() else {
            return;
        };
        discard_panicking_drop(body);
        drop(result_tx);
    }
}