Merge #228: Release 34: Fix deadlocks

Approved-by: Qqwy
Priority: Normal
Auto-deploy: false

Author: OpsBotPrime

Cargo.lock

@@ -47,7 +47,7 @@ test-integration *TEST_ARGS: build-bin build-python
   cd libs/opsqueue_python
   source "./.setup_local_venv.sh"
 
-  pytest --color=yes {{TEST_ARGS}}
+  timeout 30 pytest --color=yes {{TEST_ARGS}}
 
 # Python integration test suite, using artefacts built through Nix. Args are forwarded to pytest
 [group('nix')]
@@ -61,7 +61,7 @@ nix-test-integration *TEST_ARGS:
   export OPSQUEUE_VIA_NIX=true
   export RUST_LOG="opsqueue=debug"
 
-  pytest --color=yes {{TEST_ARGS}}
+  timeout 30 pytest --color=yes {{TEST_ARGS}}
 
 # Run all linters, fast and slow
 [group('lint')]

justfile

@@ -32,6 +32,7 @@ thiserror = "1.0.65"
 # Python FFI:
 pyo3 = { version = "0.23.4", features = ["chrono"] }
 pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime", "unstable-streams"] }
+once_cell = "1.21.3" # Only currently used for `unsync::OnceCell` as part of PyO3 async helpers.
 
 # Logging/tracing:
 pyo3-log = "0.12.1"

libs/opsqueue_python/Cargo.toml

@@ -52,10 +52,15 @@ test = [
     "pytest==8.3.3",
     "pytest-random-order==1.1.1",
     "pytest-parallel==0.1.1",
+    "pytest-timeout==2.4.0",
     "py==1.11.0", # Needs to be manually specified because of this issue: https://github.com/kevlened/pytest-parallel/issues/118
 ]
 
 [tool.pytest.ini_options]
 # We ensure tests never rely on global state,
 # by running them in a random order, and in parallel:
 addopts = "--random-order --workers=auto"
+# Individual tests should be very fast. They should never take multiple seconds
+# If after 20sec (accomodating for a toaster-like PC) there is no progress,
+# assume a deadlock
+timeout=20

libs/opsqueue_python/pyproject.toml

@@ -0,0 +1,98 @@
+//! Helpers to bridge async Rust and PyO3's flavour of async Python
+//!
+use pyo3::{Bound, IntoPyObject, PyAny, PyResult, Python};
+use pyo3_async_runtimes::TaskLocals;
+use std::{
+    future::Future,
+    pin::{pin, Pin},
+    task::{Context, Poll},
+};
+
+/// Unlock the GIL across `.await` points
+///
+/// Essentially `py.allow_threads` but for async code.
+///
+/// Based on https://pyo3.rs/v0.25.1/async-await.html#release-the-gil-across-await
+pub struct AsyncAllowThreads<F>(F);
+
+pub fn async_allow_threads<F>(fut: F) -> AsyncAllowThreads<F>
+where
+    F: Future,
+{
+    AsyncAllowThreads(fut)
+}
+
+impl<F> Future for AsyncAllowThreads<F>
+where
+    F: Future + Unpin + Send,
+    F::Output: Send,
+{
+    type Output = F::Output;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let waker = cx.waker();
+        Python::with_gil(|py| {
+            py.allow_threads(|| pin!(&mut self.0).poll(&mut Context::from_waker(waker)))
+        })
+    }
+}
+
+/// Version of `future_into_py` that uses the `TokioRuntimeThatIsInScope`
+pub fn future_into_py<T, F>(py: Python<'_>, fut: F) -> PyResult<Bound<'_, PyAny>>
+where
+    F: Future<Output = PyResult<T>> + Send + 'static,
+    T: for<'py> IntoPyObject<'py>,
+{
+    pyo3_async_runtimes::generic::future_into_py::<TokioRuntimeThatIsInScope, F, T>(py, fut)
+}
+
+/// An alternative runtime for `pyo3_async_runtimes`
+/// since `pyo3_async_runtimes::tokio` runs its _own_ Tokio runtime
+/// rather than using whatever is in scope or allowing the user to pass the specific runtime.
+///
+/// How this runtime works, is to use whatever Tokio runtime was entered using `runtime.enter()` beforehand
+struct TokioRuntimeThatIsInScope();
+
+use once_cell::unsync::OnceCell as UnsyncOnceCell;
+
+tokio::task_local! {
+    static TASK_LOCALS: UnsyncOnceCell<TaskLocals>;
+}
+
+impl pyo3_async_runtimes::generic::Runtime for TokioRuntimeThatIsInScope {
+    type JoinError = tokio::task::JoinError;
+    type JoinHandle = tokio::task::JoinHandle<()>;
+
+    fn spawn<F>(fut: F) -> Self::JoinHandle
+    where
+        F: std::future::Future<Output = ()> + Send + 'static,
+    {
+        tokio::task::spawn(async move {
+            fut.await;
+        })
+    }
+}
+
+impl pyo3_async_runtimes::generic::ContextExt for TokioRuntimeThatIsInScope {
+    fn scope<F, R>(
+        locals: TaskLocals,
+        fut: F,
+    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = R> + Send>>
+    where
+        F: std::future::Future<Output = R> + Send + 'static,
+    {
+        let cell = UnsyncOnceCell::new();
+        cell.set(locals).unwrap();
+
+        Box::pin(TASK_LOCALS.scope(cell, fut))
+    }
+
+    fn get_task_locals() -> Option<TaskLocals> {
+        TASK_LOCALS
+            .try_with(|c| {
+                c.get()
+                    .map(|locals| Python::with_gil(|py| locals.clone_ref(py)))
+            })
+            .unwrap_or_default()
+    }
+}

libs/opsqueue_python/src/async_util.rs

@@ -443,10 +443,24 @@ pub async fn check_signals_in_background() -> FatalPythonException {
 
 /// Sets up a Tokio runtime to use for a client.
 ///
-/// Rather than the current-thread scheduler,
-/// we use a (single extra!) background thread,
-/// allowing us to keep (GIL-less) tasks alive in the background
-/// even when returning back to Python
+/// Note that we very intentionally use the multi-threaded scheduler
+/// but with a single thread.
+///
+/// We **cannot** use the current_thread scheduler,
+/// since that would result in the Python task scheduler (e.g. `asyncio`)
+/// and the Rust task scheduler (`Tokio`) to run on the same thread.
+/// This seems to work fine, until you end up with a Python future
+/// that depends on a Rust future completing or vice-versa:
+/// Since the task schedulers each have their own task queues,
+/// work co-operatively, and know nothing of each-other,
+/// they will not (nor can they) yield to the other.
+/// The result: deadlocks!
+///
+/// Therefore, we run the Tokio scheduler on a separate thread.
+/// Since switching between the 'Python scheduler thread' and the
+/// 'Tokio scheduler thread' is preemptive,
+/// the same problem now no longer occurs:
+/// Both schedulers are able to make forward progress (even on a 1-CPU machine).
 pub fn start_runtime() -> Arc<tokio::runtime::Runtime> {
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .worker_threads(1)

libs/opsqueue_python/src/common.rs

@@ -1,3 +1,4 @@
+mod async_util;
 pub mod common;
 pub mod consumer;
 pub mod errors;

libs/opsqueue_python/src/lib.rs

@@ -23,6 +23,7 @@ use opsqueue::{
 use ux_serde::u63;
 
 use crate::{
+    async_util,
     common::{run_unless_interrupted, start_runtime, SubmissionId, SubmissionStatus},
     errors::{self, CError, CPyResult, FatalPythonException},
 };
@@ -217,17 +218,19 @@ impl ProducerClient {
         py.allow_threads(|| {
             let prefix = uuid::Uuid::now_v7().to_string();
             tracing::debug!("Uploading submission chunks to object store subfolder {prefix}...");
-            let chunk_count = Python::with_gil(|py| {
-                self.block_unless_interrupted(async {
-                    let chunk_contents = chunk_contents.bind(py);
-                    let stream = futures::stream::iter(chunk_contents)
-                        .map(|item| item.and_then(|item| item.extract()).map_err(Into::into));
+            let chunk_count = self.block_unless_interrupted(async {
+                    let chunk_contents = std::iter::from_fn(move || {
+                        Python::with_gil(|py|
+                            chunk_contents.bind(py).clone().next()
+                                .map(|item| item.and_then(
+                                    |item| item.extract()).map_err(Into::into)))
+                    });
+                    let stream = futures::stream::iter(chunk_contents);
                     self.object_store_client
                         .store_chunks(&prefix, ChunkType::Input, stream)
                         .await
                         .map_err(|e| CError(R(L(e))))
-                })
-            })?;
+                })?;
             let chunk_count = chunk::ChunkIndex::from(chunk_count);
             tracing::debug!("Finished uploading to object store. {prefix} contains {chunk_count} chunks");
 
@@ -360,15 +363,18 @@ impl ProducerClient {
     ) -> PyResult<Bound<'p, PyAny>> {
         let me = self.clone();
         let _tokio_active_runtime_guard = me.runtime.enter();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            match me.stream_completed_submission_chunks(submission_id).await {
-                Ok(iter) => {
-                    let async_iter = PyChunksAsyncIter::from(iter);
-                    Ok(async_iter)
+        async_util::future_into_py(
+            py,
+            async_util::async_allow_threads(Box::pin(async move {
+                match me.stream_completed_submission_chunks(submission_id).await {
+                    Ok(iter) => {
+                        let async_iter = PyChunksAsyncIter::from(iter);
+                        Ok(async_iter)
+                    }
+                    Err(e) => PyResult::Err(e.into()),
                 }
-                Err(e) => PyResult::Err(e.into()),
-            }
-        })
+            })),
+        )
     }
 }
 
@@ -462,7 +468,7 @@ pub type ChunksStream = BoxStream<'static, CPyResult<Vec<u8>, ChunkRetrievalErro
 
 #[pyclass]
 pub struct PyChunksIter {
-    stream: tokio::sync::Mutex<ChunksStream>,
+    stream: Arc<tokio::sync::Mutex<ChunksStream>>,
     runtime: Arc<tokio::runtime::Runtime>,
 }
 
@@ -475,7 +481,7 @@ impl PyChunksIter {
             .map_err(CError)
             .boxed();
         Self {
-            stream: tokio::sync::Mutex::new(stream),
+            stream: Arc::new(tokio::sync::Mutex::new(stream)),
             runtime: client.runtime.clone(),
         }
     }
@@ -487,11 +493,21 @@ impl PyChunksIter {
         slf
     }
 
-    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<CPyResult<Vec<u8>, ChunkRetrievalError>> {
-        let me = &mut *slf;
-        let runtime = &mut me.runtime;
-        let stream = &mut me.stream;
-        runtime.block_on(async { stream.lock().await.next().await })
+    fn __next__(&self, py: Python<'_>) -> Option<CPyResult<Vec<u8>, ChunkRetrievalError>> {
+        // The only time we need the GIL is when turning the result back.
+        // By unlocking here, we reduce the chance of deadlocks.
+        py.allow_threads(move || {
+            let runtime = self.runtime.clone();
+            let stream = self.stream.clone();
+            runtime.block_on(async {
+                // We lock the stream in a separate Tokio task
+                // that explicitly runs on the runtime thread rather than on the main Python thread.
+                // This reduces the possibility for deadlocks even further.
+                tokio::task::spawn(async move { stream.lock().await.next().await })
+                    .await
+                    .expect("Top-level spawn to succeed")
+            })
+        })
     }
 
     fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
@@ -508,7 +524,7 @@ pub struct PyChunksAsyncIter {
 impl From<PyChunksIter> for PyChunksAsyncIter {
     fn from(iter: PyChunksIter) -> Self {
         Self {
-            stream: Arc::new(iter.stream),
+            stream: iter.stream,
             runtime: iter.runtime,
         }
     }
@@ -520,16 +536,27 @@ impl PyChunksAsyncIter {
         slf
     }
 
-    fn __anext__(slf: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
-        let _tokio_active_runtime_guard = slf.runtime.enter();
-        let stream = slf.stream.clone();
-        pyo3_async_runtimes::tokio::future_into_py(slf.py(), async move {
-            let res = stream.lock().await.next().await;
-            match res {
-                None => Err(PyStopAsyncIteration::new_err(())),
-                Some(Ok(val)) => Ok(Some(val)),
-                Some(Err(e)) => Err(e.into()),
-            }
-        })
+    fn __anext__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let stream = self.stream.clone();
+        let _tokio_active_runtime_guard = self.runtime.enter();
+
+        async_util::future_into_py(
+            py,
+            // The only time we need the GIL is when turning the result into Python datatypes.
+            // By unlocking here, we reduce the chance of deadlocks.
+            async_util::async_allow_threads(Box::pin(async move {
+                // We lock the stream in a separate Tokio task
+                // that explicitly runs on the runtime thread rather than on the main Python thread.
+                // This reduces the possibility for deadlocks even further.
+                let res = tokio::task::spawn(async move { stream.lock().await.next().await })
+                    .await
+                    .expect("Top-level spawn to succeed");
+                match res {
+                    None => Err(PyStopAsyncIteration::new_err(())),
+                    Some(Ok(val)) => Ok(Some(val)),
+                    Some(Err(e)) => Err(e.into()),
+                }
+            })),
+        )
     }
 }

libs/opsqueue_python/src/producer.rs

@@ -19,6 +19,8 @@
 #     print("A")
 #     multiprocessing.set_start_method('forkserver')
 
+PROJECT_ROOT = Path(__file__).parents[3]
+
 
 @dataclass
 class OpsqueueProcess:
@@ -34,8 +36,12 @@ def opsqueue_bin_location() -> Path:
         )
         return Path(deriv_path) / "bin" / "opsqueue"
     else:
-        subprocess.run(["cargo", "build", "--quiet", "--bin", "opsqueue"])
-        return Path(".", "target", "debug", "opsqueue")
+        subprocess.run(
+            ["cargo", "build", "--quiet", "--bin", "opsqueue"],
+            cwd=PROJECT_ROOT,
+            check=True,
+        )
+        return PROJECT_ROOT / Path("target", "debug", "opsqueue")
 
 
 @pytest.fixture
@@ -62,12 +68,11 @@ def opsqueue_service(
         "--database-filename",
         temp_dbname,
     ]
-    cwd = "../../"
     env = os.environ.copy()  # We copy the env so e.g. RUST_LOG and other env vars are propagated from outside of the invocation of pytest
     if env.get("RUST_LOG") is None:
         env["RUST_LOG"] = "off"
 
-    with subprocess.Popen(command, cwd=cwd, env=env) as process:
+    with subprocess.Popen(command, cwd=PROJECT_ROOT, env=env) as process:
         try:
             wrapper = OpsqueueProcess(port=port, process=process)
             yield wrapper