From fece2bf9e80d137adc18176524971bed1af3d1ed Mon Sep 17 00:00:00 2001 From: David Leong <116610336+leongdl@users.noreply.github.com> Date: Wed, 17 Jun 2026 11:45:09 -0700 Subject: [PATCH] feat: unwrap hardening from code review Signed-off-by: David Leong <116610336+leongdl@users.noreply.github.com> --- rust-bindings/src/expr/errors.rs | 21 ++ rust-bindings/src/expr/symbol_table.rs | 16 +- .../src/model/step_dependency_graph.rs | 70 +++-- rust-bindings/src/model/step_param_space.rs | 27 +- rust-bindings/src/sessions/session.rs | 289 ++++++++++++------ 5 files changed, 294 insertions(+), 129 deletions(-) diff --git a/rust-bindings/src/expr/errors.rs b/rust-bindings/src/expr/errors.rs index e6125b7b..f37654ca 100644 --- a/rust-bindings/src/expr/errors.rs +++ b/rust-bindings/src/expr/errors.rs @@ -70,6 +70,27 @@ pyo3::create_exception!( // caret column shifted accordingly. Falls back to // `str(self)` for context-free or multi-line errors. // +// Rendering shape of `message_with_expr_prefix` (the "printing" +// path). When `expr` is attached and single-line, it produces a +// three-line, caret-annotated message: +// +// +// +// ^ +// +// * line 1 is `self._base_message` (the original message text); +// * line 2 is the expression source, two-space-indented, with +// the caller-supplied `prefix` prepended +// (`" " + prefix + expr`); +// * line 3 is emitted only when `col_offset is not None`: a +// caret under the offending column. The caret is shifted right +// by `len(prefix)` (`" " * (col_offset + len(prefix))`) so it +// stays aligned with the source character after the prefix +// pushes the expression text rightward. +// +// If `expr is None` or the expression is multi-line (`"\n" in +// expr`), it returns plain `str(self)` with no caret annotation. +// // Implementing these as `#[pyfunction]`s installed via `setattr` // would *almost* work, except that PyO3's `#[pyfunction]` builds // a `PyCFunction` (builtin function) which is not bound to its diff --git a/rust-bindings/src/expr/symbol_table.rs b/rust-bindings/src/expr/symbol_table.rs index d481a1c1..13762f44 100644 --- a/rust-bindings/src/expr/symbol_table.rs +++ b/rust-bindings/src/expr/symbol_table.rs @@ -73,10 +73,10 @@ impl PySymbolTable { use pyo3::IntoPyObjectExt; match self.inner.get(key) { Some(openjd_expr::symbol_table::SymbolTableEntry::Value(v)) => { - Ok(PyExprValue { inner: v.clone() }.into_py_any(py).unwrap()) + PyExprValue { inner: v.clone() }.into_py_any(py) } Some(openjd_expr::symbol_table::SymbolTableEntry::Table(t)) => { - Ok(PySymbolTable { inner: t.clone() }.into_py_any(py).unwrap()) + PySymbolTable { inner: t.clone() }.into_py_any(py) } None => Err(pyo3::exceptions::PyKeyError::new_err(key.to_string())), } @@ -85,12 +85,12 @@ impl PySymbolTable { fn get(&self, py: Python<'_>, name: &str) -> PyResult>> { use pyo3::IntoPyObjectExt; match self.inner.get(name) { - Some(openjd_expr::symbol_table::SymbolTableEntry::Value(v)) => Ok(Some( - PyExprValue { inner: v.clone() }.into_py_any(py).unwrap(), - )), - Some(openjd_expr::symbol_table::SymbolTableEntry::Table(t)) => Ok(Some( - PySymbolTable { inner: t.clone() }.into_py_any(py).unwrap(), - )), + Some(openjd_expr::symbol_table::SymbolTableEntry::Value(v)) => { + Ok(Some(PyExprValue { inner: v.clone() }.into_py_any(py)?)) + } + Some(openjd_expr::symbol_table::SymbolTableEntry::Table(t)) => { + Ok(Some(PySymbolTable { inner: t.clone() }.into_py_any(py)?)) + } None => Ok(None), } } diff --git a/rust-bindings/src/model/step_dependency_graph.rs b/rust-bindings/src/model/step_dependency_graph.rs index 0715751f..27e55c1b 100644 --- a/rust-bindings/src/model/step_dependency_graph.rs +++ b/rust-bindings/src/model/step_dependency_graph.rs @@ -18,37 +18,45 @@ pub(crate) struct PyStepDependencyGraph { } impl PyStepDependencyGraph { - fn make_node(&self, node_index: usize) -> PyStepDependencyNode { - let node = self.inner.node(node_index).unwrap(); + fn make_node(&self, node_index: usize) -> PyResult { + let node = self.inner.node(node_index).ok_or_else(|| { + pyo3::exceptions::PyIndexError::new_err(format!( + "step dependency graph has no node at index {node_index}" + )) + })?; + let step_for = |idx: usize| -> PyResult { + self.job_steps.get(idx).cloned().ok_or_else(|| { + pyo3::exceptions::PyIndexError::new_err(format!( + "step dependency edge references out-of-range step index {idx}" + )) + }) + }; + let edge_for = |edge_idx| -> PyResult> { + match self.inner.edge(edge_idx) { + Some(edge) => Ok(Some(PyStepDependencyEdge { + origin_step: step_for(edge.origin)?, + dependent_step: step_for(edge.dependent)?, + })), + None => Ok(None), + } + }; let in_edges: Vec = node .in_edges .iter() - .filter_map(|&edge_idx| { - let edge = self.inner.edge(edge_idx)?; - Some(PyStepDependencyEdge { - origin_step: self.job_steps[edge.origin].clone(), - dependent_step: self.job_steps[edge.dependent].clone(), - }) - }) - .collect(); + .filter_map(|&edge_idx| edge_for(edge_idx).transpose()) + .collect::>()?; let out_edges: Vec = node .out_edges .iter() - .filter_map(|&edge_idx| { - let edge = self.inner.edge(edge_idx)?; - Some(PyStepDependencyEdge { - origin_step: self.job_steps[edge.origin].clone(), - dependent_step: self.job_steps[edge.dependent].clone(), - }) - }) - .collect(); - PyStepDependencyNode { + .filter_map(|&edge_idx| edge_for(edge_idx).transpose()) + .collect::>()?; + Ok(PyStepDependencyNode { step: PyStep { - inner: self.job_steps[node.step_index].clone(), + inner: step_for(node.step_index)?, }, in_edges, out_edges, - } + }) } } @@ -66,7 +74,7 @@ impl PyStepDependencyGraph { } #[getter] - fn _nodes(&self) -> Vec { + fn _nodes(&self) -> PyResult> { (0..self.inner.node_count()) .map(|i| self.make_node(i)) .collect() @@ -76,17 +84,25 @@ impl PyStepDependencyGraph { let node = self.inner.step_node(stepname).ok_or_else(|| { pyo3::exceptions::PyKeyError::new_err(format!("No step named '{stepname}'")) })?; - Ok(self.make_node(node.step_index)) + self.make_node(node.step_index) } fn topo_sorted(&self) -> PyResult> { let indices = self.inner.topo_sorted().map_err(model_err_to_py)?; - Ok(indices + indices .into_iter() - .map(|i| PyStep { - inner: self.job_steps[i].clone(), + .map(|i| { + self.job_steps + .get(i) + .cloned() + .map(|inner| PyStep { inner }) + .ok_or_else(|| { + pyo3::exceptions::PyIndexError::new_err(format!( + "topo sort produced out-of-range step index {i}" + )) + }) }) - .collect()) + .collect() } fn step_names(&self) -> PyResult> { diff --git a/rust-bindings/src/model/step_param_space.rs b/rust-bindings/src/model/step_param_space.rs index 78782da8..c17cc468 100644 --- a/rust-bindings/src/model/step_param_space.rs +++ b/rust-bindings/src/model/step_param_space.rs @@ -91,6 +91,15 @@ fn param_type_to_expr_type(pt: TaskParameterType) -> openjd_expr::ExprType { } } +/// Lock the iterator mutex, recovering the guard if a previous holder +/// panicked. A panic inside `next()` / `contains()` / `validate_containment()` +/// would otherwise poison the mutex, and every subsequent `.lock().unwrap()` +/// would raise an uncatchable `PanicException` on the Python side — wedging +/// the iterator object permanently. Recovering keeps it usable. +fn lock_recover(m: &Mutex) -> std::sync::MutexGuard<'_, T> { + m.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) +} + #[cfg_attr(feature = "stub-gen", gen_stub_pyclass(module = "openjd._openjd_rs"))] #[pyclass(module = "openjd.model._v1.job", name = "StepParameterSpaceIterator")] pub(crate) struct PyStepParameterSpaceIterator { @@ -154,7 +163,7 @@ impl PyStepParameterSpaceIterator { // Match the pure-Python reference: adaptive-chunked spaces // cannot answer `len()` because the count depends on the // dynamic chunk size that may change during execution. - let iter = self.iter.lock().unwrap(); + let iter = lock_recover(&self.iter); if iter.chunks_adaptive() { return Err(pyo3::exceptions::PyValueError::new_err( "Length is not available because the parameter space uses adaptive chunking.", @@ -195,7 +204,7 @@ impl PyStepParameterSpaceIterator { } fn __next__(&self, py: Python<'_>) -> PyResult>> { - let mut iter = self.iter.lock().unwrap(); + let mut iter = lock_recover(&self.iter); match iter.next() { Some(params) => Ok(Some(task_param_set_to_py(py, ¶ms)?)), None => Ok(None), @@ -204,7 +213,7 @@ impl PyStepParameterSpaceIterator { fn __contains__(&self, item: &Bound<'_, PyDict>) -> PyResult { let params = extract_task_parameter_set(item)?; - let iter = self.iter.lock().unwrap(); + let iter = lock_recover(&self.iter); Ok(iter.contains(¶ms)) } @@ -217,13 +226,13 @@ impl PyStepParameterSpaceIterator { /// crate's ``StepParameterSpaceIterator::validate_containment``. fn validate_containment(&self, params: &Bound<'_, PyDict>) -> PyResult<()> { let params = extract_task_parameter_set(params)?; - let iter = self.iter.lock().unwrap(); + let iter = lock_recover(&self.iter); iter.validate_containment(¶ms) .map_err(pyo3::exceptions::PyValueError::new_err) } fn reset_iter(&self) { - let mut iter = self.iter.lock().unwrap(); + let mut iter = lock_recover(&self.iter); iter.reset(); } @@ -234,19 +243,19 @@ impl PyStepParameterSpaceIterator { #[getter] fn chunks_adaptive(&self) -> bool { - let iter = self.iter.lock().unwrap(); + let iter = lock_recover(&self.iter); iter.chunks_adaptive() } #[getter] fn chunks_parameter_name(&self) -> Option { - let iter = self.iter.lock().unwrap(); + let iter = lock_recover(&self.iter); iter.chunks_parameter_name().map(|s| s.to_string()) } #[getter] fn chunks_default_task_count(&self) -> Option { - let iter = self.iter.lock().unwrap(); + let iter = lock_recover(&self.iter); iter.chunks_default_task_count() } @@ -257,7 +266,7 @@ impl PyStepParameterSpaceIterator { "chunks_default_task_count must be a positive integer.", )); } - let mut iter = self.iter.lock().unwrap(); + let mut iter = lock_recover(&self.iter); if !iter.chunks_adaptive() { return Err(pyo3::exceptions::PyValueError::new_err( "The parameter space does not use adaptive chunking, so cannot modify chunks_default_task_count.", diff --git a/rust-bindings/src/sessions/session.rs b/rust-bindings/src/sessions/session.rs index 017da034..6bcdda70 100644 --- a/rust-bindings/src/sessions/session.rs +++ b/rust-bindings/src/sessions/session.rs @@ -117,6 +117,147 @@ struct StateSnapshot { environments_entered: Vec, } +/// Lock a mutex, recovering the guard even if a previous holder panicked. +/// +/// Action work runs on detached `std::thread::spawn` threads. A panic on +/// one of those threads (e.g. inside async session code) would poison +/// `session` / `snapshot`, after which every later `.lock().unwrap()` — +/// including the read-only property getters invoked from the Python +/// thread — would itself panic and surface as an uncatchable +/// `PanicException`, permanently wedging the session object. Recovering +/// the poisoned guard keeps the session readable so the failure can be +/// reported through the normal `ActionStatus` channel instead of +/// cascading. +/// +/// Caveat: `into_inner()` recovers the lock but not the invariants — a +/// thread that panicked mid-mutation may have left the data readable but +/// not necessarily consistent. For the read-only snapshot getters that's +/// harmless (worst case: stale fields). For the `session` slot, the +/// panic-handling path in `run_action` forces a terminal FAILED +/// `ActionStatus`, which closes the practical risk of re-dispatching an +/// action onto a half-mutated session. +fn lock_recover(m: &Mutex) -> std::sync::MutexGuard<'_, T> { + m.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +/// Build a terminal FAILED `ActionStatus` carrying `message`. +/// +/// Used when an action thread can't even start (e.g. the Tokio runtime +/// fails to build) so the Python poller observes a terminal state and +/// stops waiting on RUNNING rather than spinning forever. +fn failed_action_status(message: &str) -> ActionStatus { + ActionStatus { + state: openjd_sessions::action::ActionState::Failed, + progress: None, + status_message: None, + fail_message: Some(message.to_string()), + exit_code: None, + started_at: None, + ended_at: None, + } +} + +/// Body of a spawned action thread. +/// +/// Builds a lightweight current-thread Tokio runtime, runs the action +/// closure to completion while catching panics, then **always** returns +/// the `Session` to the shared slot and refreshes the snapshot — even if +/// the runtime failed to build or the future panicked. This guarantees: +/// +/// * the session is never left permanently "taken" (which would wedge +/// every later call into `"An action is already running"`), and +/// * a panic in async session code can't escape the detached thread to +/// poison shared state silently and tear down the interpreter's view +/// of the session. +/// +/// A runtime-build failure is surfaced as a terminal FAILED +/// `ActionStatus` so the Python-side poller terminates cleanly. A panic +/// inside the action is likewise converted into a terminal FAILED +/// `ActionStatus` (rather than copying the possibly-still-Running real +/// session state), so the poller never hangs on a panicked action. +fn run_action( + mut session: Session, + session_arc: Arc>>, + snapshot: Arc>, + action: F, +) where + F: FnOnce(&tokio::runtime::Runtime, &mut Session), +{ + let rt = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(rt) => rt, + Err(e) => { + // Couldn't build a runtime (e.g. fd/thread exhaustion). Mark + // the action FAILED and drop out of RUNNING so the poller + // delivers a terminal callback instead of hanging. + { + let mut snap = lock_recover(&snapshot); + snap.state = SessionState::Ready; + snap.action_status = Some(failed_action_status(&format!( + "failed to start action runtime: {e}" + ))); + } + *lock_recover(&session_arc) = Some(session); + return; + } + }; + + // Catch panics so a panic inside async session code still restores + // the session below rather than leaking out of this detached thread. + let action_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + action(&rt, &mut session); + })); + + // Restore the session BEFORE refreshing the snapshot, holding the guard + // across both so there's no window where a caller finds the slot empty. + // Python callers watch `state` via the snapshot and dispatch the next + // action the moment it transitions out of Running; if the snapshot + // updated first they could find the mutex still empty ("An action is + // already running"). + let guard = { + let mut slot = lock_recover(&session_arc); + *slot = Some(session); + slot + }; + + match action_result { + Ok(()) => { + if let Some(s) = guard.as_ref() { + PySession::update_snapshot(s, &snapshot); + } + } + Err(payload) => { + // The action panicked mid-run. `catch_unwind` kept the session + // usable for future calls, but the real `session.state()` may + // still read Running with no terminal `ActionStatus` — which + // would leave the Python poller waiting forever. Force a + // terminal FAILED status and drop out of RUNNING, mirroring the + // runtime-build-failure branch above. + let mut snap = lock_recover(&snapshot); + snap.state = SessionState::Ready; + snap.action_status = Some(failed_action_status(&format!( + "action panicked during execution: {}", + panic_detail(payload.as_ref()) + ))); + } + } +} + +/// Best-effort extraction of a human-readable message from a caught panic +/// payload (`std::panic::catch_unwind`'s `Err` value), which is usually the +/// `&str` or `String` passed to `panic!`. +fn panic_detail(payload: &(dyn std::any::Any + Send)) -> String { + if let Some(s) = payload.downcast_ref::<&'static str>() { + (*s).to_string() + } else if let Some(s) = payload.downcast_ref::() { + s.clone() + } else { + "unknown panic".to_string() + } +} + #[cfg_attr(feature = "stub-gen", gen_stub_pyclass(module = "openjd._openjd_rs"))] #[pyclass(module = "openjd.sessions._v1", name = "Session")] pub(crate) struct PySession { @@ -128,7 +269,7 @@ pub(crate) struct PySession { impl PySession { fn update_snapshot(session: &Session, snapshot: &Arc>) { - let mut snap = snapshot.lock().unwrap(); + let mut snap = lock_recover(snapshot); snap.state = session.state(); snap.action_status = session.action_status(); snap.environments_entered = session.environments_entered().to_vec(); @@ -213,7 +354,7 @@ impl PySession { }; let session = Session::with_config(config).map_err(session_err_to_py)?; { - let mut snap = snapshot.lock().unwrap(); + let mut snap = lock_recover(&snapshot); snap.state = session.state(); snap.working_directory = session.working_directory().to_string_lossy().to_string(); snap.files_directory = session.files_directory().to_string_lossy().to_string(); @@ -229,29 +370,27 @@ impl PySession { #[getter] fn session_id(&self) -> String { - self.snapshot.lock().unwrap().session_id.clone() + lock_recover(&self.snapshot).session_id.clone() } #[getter] fn state(&self) -> PySessionState { - self.snapshot.lock().unwrap().state.into() + lock_recover(&self.snapshot).state.into() } #[getter] fn working_directory(&self) -> String { - self.snapshot.lock().unwrap().working_directory.clone() + lock_recover(&self.snapshot).working_directory.clone() } #[getter] fn files_directory(&self) -> String { - self.snapshot.lock().unwrap().files_directory.clone() + lock_recover(&self.snapshot).files_directory.clone() } #[getter] fn action_status(&self) -> Option { - self.snapshot - .lock() - .unwrap() + lock_recover(&self.snapshot) .action_status .clone() .map(PyActionStatus::from) @@ -259,7 +398,7 @@ impl PySession { #[getter] fn environments_entered(&self) -> Vec { - self.snapshot.lock().unwrap().environments_entered.clone() + lock_recover(&self.snapshot).environments_entered.clone() } /// Extend the session's path mapping rules with additional rules. @@ -278,7 +417,7 @@ impl PySession { &self, additional: Vec, ) -> PyResult<()> { - let mut guard = self.session.lock().unwrap(); + let mut guard = lock_recover(&self.session); let session = guard.as_mut().ok_or_else(|| { pyo3::exceptions::PyRuntimeError::new_err( "cannot extend path mapping rules while an action is running", @@ -306,7 +445,7 @@ impl PySession { let env_id = match &identifier { Some(id) => id.clone(), None => { - let snap = self.snapshot.lock().unwrap(); + let snap = lock_recover(&self.snapshot); format!("{}:{}", snap.session_id, uuid::Uuid::new_v4().simple()) } }; @@ -321,8 +460,8 @@ impl PySession { let resolved = resolved_symtab.map(|st| st.inner.clone()); // Take the session out of the mutex so the background thread owns it - let mut guard = self.session.lock().unwrap(); - let mut session = guard.take().ok_or_else(|| { + let mut guard = lock_recover(&self.session); + let session = guard.take().ok_or_else(|| { pyo3::exceptions::PyRuntimeError::new_err("An action is already running") })?; drop(guard); @@ -334,7 +473,7 @@ impl PySession { // for the new action — causing the worker to think envEnter / run_task // finished in ~1ms with the previous action's SUCCEEDED result. { - let mut snap = self.snapshot.lock().unwrap(); + let mut snap = lock_recover(&self.snapshot); snap.state = SessionState::Running; snap.action_status = None; } @@ -343,24 +482,15 @@ impl PySession { let snapshot = self.snapshot.clone(); std::thread::spawn(move || { - let rt = tokio::runtime::Runtime::new().unwrap(); - let env_ref = os_env_vars.as_ref(); - let _ = rt.block_on(session.enter_environment( - &env, - resolved.as_ref(), - Some(&env_id), - env_ref, - )); - // Put session back BEFORE updating the snapshot. Python callers - // watch `state` via the snapshot and dispatch the next action the - // moment it transitions out of Running. If we updated the snapshot - // first, a Python caller could call run_task / enter_environment - // and find the mutex still empty — "An action is already running". - *session_arc.lock().unwrap() = Some(session); - let guard = session_arc.lock().unwrap(); - if let Some(s) = guard.as_ref() { - Self::update_snapshot(s, &snapshot); - } + run_action(session, session_arc, snapshot, move |rt, session| { + let env_ref = os_env_vars.as_ref(); + let _ = rt.block_on(session.enter_environment( + &env, + resolved.as_ref(), + Some(&env_id), + env_ref, + )); + }); }); Ok(return_id) @@ -378,14 +508,14 @@ impl PySession { ) -> PyResult<()> { let resolved = resolved_symtab.map(|st| st.inner.clone()); - let mut guard = self.session.lock().unwrap(); - let mut session = guard.take().ok_or_else(|| { + let mut guard = lock_recover(&self.session); + let session = guard.take().ok_or_else(|| { pyo3::exceptions::PyRuntimeError::new_err("An action is already running") })?; drop(guard); { - let mut snap = self.snapshot.lock().unwrap(); + let mut snap = lock_recover(&self.snapshot); snap.state = SessionState::Running; snap.action_status = None; } @@ -394,19 +524,15 @@ impl PySession { let snapshot = self.snapshot.clone(); std::thread::spawn(move || { - let rt = tokio::runtime::Runtime::new().unwrap(); - let env_ref = os_env_vars.as_ref(); - let _ = rt.block_on(session.exit_environment( - &identifier, - resolved.as_ref(), - keep_session_running, - env_ref, - )); - *session_arc.lock().unwrap() = Some(session); - let guard = session_arc.lock().unwrap(); - if let Some(s) = guard.as_ref() { - Self::update_snapshot(s, &snapshot); - } + run_action(session, session_arc, snapshot, move |rt, session| { + let env_ref = os_env_vars.as_ref(); + let _ = rt.block_on(session.exit_environment( + &identifier, + resolved.as_ref(), + keep_session_running, + env_ref, + )); + }); }); Ok(()) @@ -428,14 +554,14 @@ impl PySession { }; let resolved = resolved_symtab.map(|st| st.inner.clone()); - let mut guard = self.session.lock().unwrap(); - let mut session = guard.take().ok_or_else(|| { + let mut guard = lock_recover(&self.session); + let session = guard.take().ok_or_else(|| { pyo3::exceptions::PyRuntimeError::new_err("An action is already running") })?; drop(guard); { - let mut snap = self.snapshot.lock().unwrap(); + let mut snap = lock_recover(&self.snapshot); snap.state = SessionState::Running; snap.action_status = None; } @@ -444,15 +570,12 @@ impl PySession { let snapshot = self.snapshot.clone(); std::thread::spawn(move || { - let rt = tokio::runtime::Runtime::new().unwrap(); - let env_ref = os_env_vars.as_ref(); - let task_ref = task_params.as_ref(); - let _ = rt.block_on(session.run_task(&script, task_ref, resolved.as_ref(), env_ref)); - *session_arc.lock().unwrap() = Some(session); - let guard = session_arc.lock().unwrap(); - if let Some(s) = guard.as_ref() { - Self::update_snapshot(s, &snapshot); - } + run_action(session, session_arc, snapshot, move |rt, session| { + let env_ref = os_env_vars.as_ref(); + let task_ref = task_params.as_ref(); + let _ = + rt.block_on(session.run_task(&script, task_ref, resolved.as_ref(), env_ref)); + }); }); Ok(()) @@ -469,14 +592,14 @@ impl PySession { use_session_env_vars: bool, log_banner_message: Option, ) -> PyResult<()> { - let mut guard = self.session.lock().unwrap(); - let mut session = guard.take().ok_or_else(|| { + let mut guard = lock_recover(&self.session); + let session = guard.take().ok_or_else(|| { pyo3::exceptions::PyRuntimeError::new_err("An action is already running") })?; drop(guard); { - let mut snap = self.snapshot.lock().unwrap(); + let mut snap = lock_recover(&self.snapshot); snap.state = SessionState::Running; snap.action_status = None; } @@ -485,24 +608,20 @@ impl PySession { let snapshot = self.snapshot.clone(); std::thread::spawn(move || { - let rt = tokio::runtime::Runtime::new().unwrap(); - let duration = timeout.map(std::time::Duration::from_secs_f64); - let env_ref = os_env_vars.as_ref(); - let args_ref = args.as_deref(); - let banner_ref = log_banner_message.as_deref(); - let _ = rt.block_on(session.run_subprocess( - &command, - args_ref, - duration, - env_ref, - use_session_env_vars, - banner_ref, - )); - *session_arc.lock().unwrap() = Some(session); - let guard = session_arc.lock().unwrap(); - if let Some(s) = guard.as_ref() { - Self::update_snapshot(s, &snapshot); - } + run_action(session, session_arc, snapshot, move |rt, session| { + let duration = timeout.map(std::time::Duration::from_secs_f64); + let env_ref = os_env_vars.as_ref(); + let args_ref = args.as_deref(); + let banner_ref = log_banner_message.as_deref(); + let _ = rt.block_on(session.run_subprocess( + &command, + args_ref, + duration, + env_ref, + use_session_env_vars, + banner_ref, + )); + }); }); Ok(()) @@ -518,7 +637,7 @@ impl PySession { // subprocess runner. But we don't have access to &mut Session here. // For now, this is a limitation — cancel requires the session to not be taken. let duration = time_limit.map(std::time::Duration::from_secs_f64); - let mut guard = self.session.lock().unwrap(); + let mut guard = lock_recover(&self.session); match guard.as_mut() { Some(session) => session .cancel_action(duration, mark_action_failed.unwrap_or(false)) @@ -530,7 +649,7 @@ impl PySession { } fn cleanup(&self) { - let mut guard = self.session.lock().unwrap(); + let mut guard = lock_recover(&self.session); if let Some(session) = guard.as_mut() { session.cleanup(); Self::update_snapshot(session, &self.snapshot); @@ -538,7 +657,7 @@ impl PySession { } fn __repr__(&self) -> String { - let snap = self.snapshot.lock().unwrap(); + let snap = lock_recover(&self.snapshot); format!("Session(id={:?}, state={:?})", snap.session_id, snap.state) } }