mod alef_async_runtime {
use pyo3::prelude::*;
use pyo3_async_runtimes::generic::{ContextExt, Runtime};
use pyo3_async_runtimes::TaskLocals;
use std::future::Future;
use std::pin::Pin;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU32, Ordering};
struct State {
runtime: Option<tokio::runtime::Runtime>,
active: usize,
closed: bool,
}
static STATE: Mutex<State> = Mutex::new(State { runtime: None, active: 0, closed: false });
static OWNER_PID: AtomicU32 = AtomicU32::new(0);
const EXIT_WAIT: std::time::Duration = std::time::Duration::from_secs(5);
tokio::task_local! {
static LOCALS: TaskLocals;
}
fn check_process() -> PyResult<()> {
let owner = OWNER_PID.load(Ordering::Acquire);
if owner != 0 && owner != std::process::id() {
return Err(pyo3::exceptions::PyRuntimeError::new_err(
"async runtime cannot be used after fork once initialized in the parent; use spawn or fork before first async use"
));
}
Ok(())
}
fn ensure_open() -> PyResult<()> {
check_process()?;
OWNER_PID.compare_exchange(0, std::process::id(), Ordering::AcqRel, Ordering::Acquire).ok();
check_process()?;
if STATE.lock().expect("managed runtime lock poisoned").closed {
return Err(pyo3::exceptions::PyRuntimeError::new_err("async runtime is closed at interpreter exit"));
}
Ok(())
}
struct Lease(bool);
impl Lease {
fn acquire() -> Self {
if check_process().is_err() { return Self(false); }
let mut state = STATE.lock().expect("managed runtime lock poisoned");
if state.closed { return Self(false); }
state.active += 1;
Self(true)
}
}
impl Drop for Lease {
fn drop(&mut self) {
if self.0 && check_process().is_ok() {
STATE.lock().expect("managed runtime lock poisoned").active -= 1;
}
}
}
pub(crate) struct RuntimeHandle {
runtime: tokio::runtime::Handle,
_lease: Lease,
}
pub(crate) fn get_runtime_checked() -> PyResult<RuntimeHandle> {
check_process()?;
OWNER_PID.compare_exchange(0, std::process::id(), Ordering::AcqRel, Ordering::Acquire).ok();
check_process()?;
let mut state = STATE.lock().expect("managed runtime lock poisoned");
if state.closed {
return Err(pyo3::exceptions::PyRuntimeError::new_err("async runtime is closed at interpreter exit"));
}
let runtime = state.runtime.get_or_insert_with(|| {
let mut builder = tokio::runtime::Builder::new_multi_thread();
builder.enable_all().thread_stack_size(16 * 1024 * 1024);
builder.build().expect("create managed async runtime")
}).handle().clone();
state.active += 1;
Ok(RuntimeHandle { runtime, _lease: Lease(true) })
}
fn get_runtime() -> RuntimeHandle {
get_runtime_checked().unwrap_or_else(|_| {
// The generic bridge cannot return an error from spawn; a stopped handle cancels racing work. ~keep
let runtime = tokio::runtime::Builder::new_current_thread().build().expect("create stopped handle");
let handle = runtime.handle().clone();
runtime.shutdown_background();
RuntimeHandle { runtime: handle, _lease: Lease(false) }
})
}
pub(crate) fn spawn_blocking<F, T>(function: F) -> tokio::task::JoinHandle<T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
get_runtime().spawn_blocking(function)
}
impl RuntimeHandle {
pub(crate) fn spawn<F>(&self, future: F) -> tokio::task::JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let lease = Lease::acquire();
self.runtime.spawn(async move {
let _lease = lease;
future.await
})
}
pub(crate) fn spawn_blocking<F, T>(&self, function: F) -> tokio::task::JoinHandle<T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let lease = Lease::acquire();
self.runtime.spawn_blocking(move || {
let _lease = lease;
function()
})
}
pub(crate) fn block_on<F: Future>(&self, future: F) -> F::Output {
self.runtime.block_on(future)
}
}
struct ManagedRuntime;
impl Runtime for ManagedRuntime {
type JoinError = tokio::task::JoinError;
type JoinHandle = tokio::task::JoinHandle<()>;
fn spawn<F: Future<Output = ()> + Send + 'static>(future: F) -> Self::JoinHandle {
get_runtime().spawn(future)
}
fn spawn_blocking<F: FnOnce() + Send + 'static>(function: F) -> Self::JoinHandle {
get_runtime().spawn_blocking(function)
}
}
impl ContextExt for ManagedRuntime {
fn scope<F, R>(locals: TaskLocals, future: F) -> Pin<Box<dyn Future<Output = R> + Send>>
where F: Future<Output = R> + Send + 'static {
Box::pin(LOCALS.scope(locals, future))
}
fn get_task_locals() -> Option<TaskLocals> {
LOCALS.try_with(Clone::clone).ok()
}
}
pub(crate) fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where
F: Future<Output = PyResult<T>> + Send + 'static,
T: for<'py> IntoPyObject<'py> + Send + 'static,
{
ensure_open()?;
pyo3_async_runtimes::generic::future_into_py::<ManagedRuntime, F, T>(py, future)
}
pub(crate) fn future_into_py_with_locals<F, T>(py: Python<'_>, locals: TaskLocals, future: F) -> PyResult<Bound<'_, PyAny>>
where
F: Future<Output = PyResult<T>> + Send + 'static,
T: for<'py> IntoPyObject<'py> + Send + 'static,
{
ensure_open()?;
pyo3_async_runtimes::generic::future_into_py_with_locals::<ManagedRuntime, F, T>(py, locals, future)
}
pub(crate) fn get_current_locals(py: Python<'_>) -> PyResult<TaskLocals> {
pyo3_async_runtimes::generic::get_current_locals::<ManagedRuntime>(py)
}
pub(crate) fn get_current_loop(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
pyo3_async_runtimes::generic::get_current_loop::<ManagedRuntime>(py)
}
pub(crate) fn into_future(awaitable: Bound<'_, PyAny>) -> PyResult<impl Future<Output = PyResult<Py<PyAny>>> + Send + use<>> {
pyo3_async_runtimes::generic::into_future::<ManagedRuntime>(awaitable)
}
pub(crate) async fn scope<F, R>(locals: TaskLocals, future: F) -> R
where F: Future<Output = R> + Send + 'static {
ManagedRuntime::scope(locals, future).await
}
pub(crate) fn shutdown(py: Python<'_>) -> PyResult<()> {
check_process()?;
if OWNER_PID.load(Ordering::Acquire) == 0 { return Ok(()); }
py.detach(|| {
let deadline = std::time::Instant::now() + std::time::Duration::from_millis(50);
loop {
let mut state = STATE.lock().expect("managed runtime lock poisoned");
if state.active == 0 {
let runtime = state.runtime.take();
drop(state);
if let Some(runtime) = runtime { runtime.shutdown_timeout(EXIT_WAIT); }
return Ok(());
}
drop(state);
if std::time::Instant::now() >= deadline {
return Err(pyo3::exceptions::PyRuntimeError::new_err("cannot shut down async runtime while work is active"));
}
std::thread::sleep(std::time::Duration::from_millis(1));
}
})
}
#[pyfunction]
fn close_async_runtime_at_exit(py: Python<'_>) {
// A forked child must not lock state inherited from vanished parent threads. ~keep
if check_process().is_err() { return; }
let runtime = {
let mut state = STATE.lock().expect("managed runtime lock poisoned");
state.closed = true;
state.runtime.take()
};
py.detach(|| {
if let Some(runtime) = runtime { runtime.shutdown_timeout(EXIT_WAIT); }
});
}
pub(crate) fn register_exit_hook(py: Python<'_>) -> PyResult<()> {
py.import("atexit")?.call_method1("register", (wrap_pyfunction!(close_async_runtime_at_exit, py)?,))?;
Ok(())
}
}
#[pyfunction]
pub fn shutdown_async_runtime(py: Python<'_>) -> PyResult<()> {
alef_async_runtime::shutdown(py)
}