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::{Arc, Mutex};
struct State {
runtime: Option<Arc<tokio::runtime::Runtime>>,
active: usize,
}
static STATE: Mutex<State> = Mutex::new(State { runtime: None, active: 0 });
tokio::task_local! {
static LOCALS: TaskLocals;
}
struct Lease;
impl Lease {
fn acquire() -> Self {
STATE.lock().expect("managed runtime lock poisoned").active += 1;
Self
}
}
impl Drop for Lease {
fn drop(&mut self) {
STATE.lock().expect("managed runtime lock poisoned").active -= 1;
}
}
pub(crate) struct RuntimeHandle {
runtime: Arc<tokio::runtime::Runtime>,
_lease: Lease,
}
pub(crate) fn get_runtime() -> RuntimeHandle {
let mut state = STATE.lock().expect("managed runtime lock poisoned");
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);
Arc::new(builder.build().expect("create managed async runtime"))
}).clone();
state.active += 1;
RuntimeHandle { runtime, _lease: Lease }
}
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,
{
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,
{
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<()> {
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);
drop(runtime);
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]
pub fn shutdown_async_runtime(py: Python<'_>) -> PyResult<()> {
alef_async_runtime::shutdown(py)
}