feat: propagate cancellations for async contexts (#1228)

Author: georgeh0

Cargo.lock

@@ -156,3 +156,4 @@ serde_path_to_error = "0.1.17"
 redis = { version = "0.31.0", features = ["tokio-comp", "connection-manager"] }
 expect-test = "1.5.0"
 encoding_rs = "0.8.35"
+tokio-util = { version = "0.7.16", features = ["rt"] }

Cargo.toml

@@ -56,11 +56,13 @@ def run(self, coro: Coroutine[Any, Any, T]) -> T:
                 RuntimeWarning,
                 stacklevel=2,
             )
-            fut = asyncio.run_coroutine_threadsafe(coro, loop)
-            return fut.result()
 
         fut = asyncio.run_coroutine_threadsafe(coro, loop)
-        return fut.result()
+        try:
+            return fut.result()
+        except KeyboardInterrupt:
+            fut.cancel()
+            raise
 
 
 execution_context = _ExecutionContext()

python/cocoindex/runtime.py

@@ -19,7 +19,6 @@
 import asyncio
 import os
 import time
-import atexit
 from .user_app_loader import load_user_app
 from .runtime import execution_context
 import logging
@@ -32,39 +31,14 @@
 # ---------------------------------------------
 _pool_lock = threading.Lock()
 _pool: ProcessPoolExecutor | None = None
-_pool_cleanup_registered = False
 _user_apps: list[str] = []
 _logger = logging.getLogger(__name__)
 
 
-def shutdown_pool_at_exit() -> None:
-    """Best-effort shutdown of the global ProcessPoolExecutor on interpreter exit."""
-    global _pool, _pool_cleanup_registered  # pylint: disable=global-statement
-    with _pool_lock:
-        if _pool is not None:
-            try:
-                _pool.shutdown(wait=True, cancel_futures=True)
-            except Exception as e:
-                _logger.error(
-                    "Error during ProcessPoolExecutor shutdown at exit: %s",
-                    e,
-                    exc_info=True,
-                )
-            finally:
-                _pool = None
-                _pool_cleanup_registered = False
-
-
 def _get_pool() -> ProcessPoolExecutor:
-    global _pool, _pool_cleanup_registered  # pylint: disable=global-statement
+    global _pool  # pylint: disable=global-statement
     with _pool_lock:
         if _pool is None:
-            if not _pool_cleanup_registered:
-                # Register the shutdown at exit at creation time (rather than at import time)
-                # to make sure it's executed earlier in the shutdown sequence.
-                atexit.register(shutdown_pool_at_exit)
-                _pool_cleanup_registered = True
-
             # Single worker process as requested
             _pool = ProcessPoolExecutor(
                 max_workers=1,

python/cocoindex/subprocess_exec.py

@@ -11,6 +11,7 @@ use futures::future::try_join_all;
 use indicatif::ProgressBar;
 use sqlx::PgPool;
 use tokio::{sync::watch, task::JoinSet, time::MissedTickBehavior};
+use tokio_util::task::AbortOnDropHandle;
 
 pub struct FlowLiveUpdaterUpdates {
     pub active_sources: Vec<String>,
@@ -364,7 +365,10 @@ impl SourceUpdateTask {
                     }
                 }
             };
-            (Some(tokio::spawn(report_task)), Some(pb))
+            (
+                Some(AbortOnDropHandle::new(tokio::spawn(report_task))),
+                Some(pb),
+            )
         } else {
             (None, None)
         };

src/execution/live_updater.rs

@@ -1,4 +1,4 @@
-use crate::prelude::*;
+use crate::{prelude::*, py::future::from_py_future};
 
 use pyo3::{
     Bound, IntoPyObjectExt, Py, PyAny, Python, pyclass, pymethods,
@@ -93,10 +93,7 @@ impl interface::SimpleFunctionExecutor for Arc<PyFunctionExecutor> {
             let result_coro = self.call_py_fn(py, input)?;
             let task_locals =
                 pyo3_async_runtimes::TaskLocals::new(self.py_exec_ctx.event_loop.bind(py).clone());
-            Ok(pyo3_async_runtimes::into_future_with_locals(
-                &task_locals,
-                result_coro,
-            )?)
+            Ok(from_py_future(py, &task_locals, result_coro)?)
         })?;
         let result = result_fut.await;
         Python::with_gil(|py| -> Result<_> {
@@ -188,7 +185,8 @@ impl interface::SimpleFunctionFactory for PyFunctionFactory {
                         let prepare_coro = executor
                             .call_method(py, "prepare", (), None)
                             .to_result_with_py_trace(py)?;
-                        let prepare_fut = pyo3_async_runtimes::into_future_with_locals(
+                        let prepare_fut = from_py_future(
+                            py,
                             &pyo3_async_runtimes::TaskLocals::new(
                                 py_exec_ctx.event_loop.bind(py).clone(),
                             ),
@@ -294,7 +292,8 @@ impl interface::SourceExecutor for PySourceExecutor {
                 .to_result_with_py_trace(py)?;
             let task_locals =
                 pyo3_async_runtimes::TaskLocals::new(py_exec_ctx.event_loop.bind(py).clone());
-            Ok(pyo3_async_runtimes::into_future_with_locals(
+            Ok(from_py_future(
+                py,
                 &task_locals,
                 result_coro.into_bound(py),
             )?)
@@ -333,10 +332,7 @@ impl PySourceExecutor {
                 .to_result_with_py_trace(py)?;
             let task_locals =
                 pyo3_async_runtimes::TaskLocals::new(py_exec_ctx.event_loop.bind(py).clone());
-            Ok(pyo3_async_runtimes::into_future_with_locals(
-                &task_locals,
-                coro.into_bound(py),
-            )?)
+            Ok(from_py_future(py, &task_locals, coro.into_bound(py))?)
         })?;
 
         // Await the future to get the next item
@@ -523,10 +519,7 @@ impl interface::SourceFactory for PySourceConnectorFactory {
                     .to_result_with_py_trace(py)?;
                 let task_locals =
                     pyo3_async_runtimes::TaskLocals::new(py_exec_ctx.event_loop.bind(py).clone());
-                let create_future = pyo3_async_runtimes::into_future_with_locals(
-                    &task_locals,
-                    create_coro.into_bound(py),
-                )?;
+                let create_future = from_py_future(py, &task_locals, create_coro.into_bound(py))?;
                 Ok(create_future)
             })?;
 
@@ -671,7 +664,8 @@ impl interface::TargetFactory for PyExportTargetFactory {
                         let task_locals = pyo3_async_runtimes::TaskLocals::new(
                             py_exec_ctx.event_loop.bind(py).clone(),
                         );
-                        anyhow::Ok(pyo3_async_runtimes::into_future_with_locals(
+                        anyhow::Ok(from_py_future(
+                            py,
                             &task_locals,
                             prepare_coro.into_bound(py),
                         )?)
@@ -807,7 +801,8 @@ impl interface::TargetFactory for PyExportTargetFactory {
                 .to_result_with_py_trace(py)?;
             let task_locals =
                 pyo3_async_runtimes::TaskLocals::new(py_exec_ctx.event_loop.bind(py).clone());
-            Ok(pyo3_async_runtimes::into_future_with_locals(
+            Ok(from_py_future(
+                py,
                 &task_locals,
                 result_coro.into_bound(py),
             )?)
@@ -867,7 +862,8 @@ impl interface::TargetFactory for PyExportTargetFactory {
                 .to_result_with_py_trace(py)?;
             let task_locals =
                 pyo3_async_runtimes::TaskLocals::new(py_exec_ctx.event_loop.bind(py).clone());
-            Ok(pyo3_async_runtimes::into_future_with_locals(
+            Ok(from_py_future(
+                py,
                 &task_locals,
                 result_coro.into_bound(py),
             )?)

src/ops/py_factory.rs

@@ -0,0 +1,85 @@
+use crate::prelude::*;
+
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use pyo3_async_runtimes::TaskLocals;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+struct CancelOnDropPy {
+    inner: BoxFuture<'static, pyo3::PyResult<pyo3::PyObject>>,
+    task: Py<PyAny>,
+    event_loop: Py<PyAny>,
+    ctx: Py<PyAny>,
+    done: AtomicBool,
+}
+
+impl Future for CancelOnDropPy {
+    type Output = pyo3::PyResult<pyo3::PyObject>;
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match Pin::new(&mut self.inner).poll(cx) {
+            Poll::Ready(out) => {
+                self.done.store(true, Ordering::SeqCst);
+                Poll::Ready(out)
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+impl Drop for CancelOnDropPy {
+    fn drop(&mut self) {
+        if self.done.load(Ordering::SeqCst) {
+            return;
+        }
+        Python::with_gil(|py| {
+            let kwargs = PyDict::new(py);
+            let result = || -> PyResult<()> {
+                // pass context so cancellation runs under the right contextvars
+                kwargs.set_item("context", self.ctx.bind(py))?;
+                self.event_loop.bind(py).call_method(
+                    "call_soon_threadsafe",
+                    (self.task.bind(py).getattr("cancel")?,),
+                    Some(&kwargs),
+                )?;
+                // self.task.bind(py).call_method0("cancel")?;
+                Ok(())
+            }();
+            if let Err(e) = result {
+                error!("Error cancelling task: {e:?}");
+            }
+        });
+    }
+}
+
+pub fn from_py_future<'py, 'fut>(
+    py: Python<'py>,
+    locals: &TaskLocals,
+    awaitable: Bound<'py, PyAny>,
+) -> pyo3::PyResult<impl Future<Output = pyo3::PyResult<pyo3::PyObject>> + Send + use<'fut>> {
+    // 1) Capture loop + context from TaskLocals for thread-safe cancellation
+    let event_loop: Bound<'py, PyAny> = locals.event_loop(py).into();
+    let ctx: Bound<'py, PyAny> = locals.context(py);
+
+    // 2) Create a Task so we own a handle we can cancel later
+    let kwarg = PyDict::new(py);
+    kwarg.set_item("context", &ctx)?;
+    let task: Bound<'py, PyAny> = event_loop
+        .call_method("create_task", (awaitable,), Some(&kwarg))?
+        .into();
+
+    // 3) Bridge it to a Rust Future as usual
+    let fut = pyo3_async_runtimes::into_future_with_locals(locals, task.clone())?.boxed();
+
+    Ok(CancelOnDropPy {
+        inner: fut,
+        task: task.unbind(),
+        event_loop: event_loop.unbind(),
+        ctx: ctx.unbind(),
+        done: AtomicBool::new(false),
+    })
+}

src/py/future.rs

@@ -21,6 +21,7 @@ use std::sync::Arc;
 
 mod convert;
 pub(crate) use convert::*;
+pub mod future;
 
 pub struct PythonExecutionContext {
     pub event_loop: Py<PyAny>,
@@ -478,7 +479,8 @@ impl Flow {
                     let task_locals = pyo3_async_runtimes::TaskLocals::new(
                         py_exec_ctx.event_loop.bind(py).clone(),
                     );
-                    Ok(pyo3_async_runtimes::into_future_with_locals(
+                    Ok(future::from_py_future(
+                        py,
                         &task_locals,
                         result_coro.into_bound(py),
                     )?)