Simplify shutdown close contracts

Author: kemsta

macloop/__init__.py

@@ -4,6 +4,7 @@
 import os
 import queue
 import uuid
+import warnings
 import weakref
 from dataclasses import dataclass
 from pathlib import Path
@@ -68,23 +69,12 @@ def _raise_on_unexpected_kwargs(name: str, kwargs: dict[str, Any]) -> None:
     raise TypeError(f"{name} got unexpected keyword arguments: {unexpected}")
 
 
-def _close_backend_with_optional_engine(backend: Any, engine_backend: Any | None) -> None:
-    if engine_backend is not None:
-        try:
-            backend.close(engine_backend)
-            return
-        except TypeError as exc:
-            try:
-                backend.close()
-                return
-            except TypeError:
-                raise exc
-
-    close_no_restore = getattr(backend, "close_no_restore", None)
-    if close_no_restore is not None:
-        close_no_restore()
-    else:
-        backend.close()
+def _warn_deferred_native_cleanup(exc: TimeoutError) -> None:
+    warnings.warn(
+        f"AudioEngine.close() deferred native cleanup: {exc}",
+        RuntimeWarning,
+        stacklevel=2,
+    )
 
 
 @dataclass(frozen=True, slots=True)
@@ -370,12 +360,15 @@ def close(self) -> None:
         if self._closed:
             return
 
+        # High-level close is bounded best-effort cleanup. Native runtime timeouts should not
+        # wedge user code; they are surfaced as warnings while the engine still becomes closed.
         self._closed = True
         backend_err: Optional[Exception] = None
         sink_err: Optional[Exception] = None
         try:
             self._backend.close()
-        except TimeoutError:
+        except TimeoutError as exc:
+            _warn_deferred_native_cleanup(exc)
             backend_err = None
         except Exception as exc:
             backend_err = exc
@@ -526,13 +519,14 @@ def close(self) -> None:
         if self._closed:
             return
 
+        # Sink close only detaches the sink and returns its route lease when an engine is alive.
         err: Optional[Exception] = None
         engine = self._engine_ref()
         try:
-            _close_backend_with_optional_engine(
-                self._backend,
-                engine._backend if engine is not None else None,
-            )
+            if engine is not None:
+                self._backend.close_for_engine(engine._backend)
+            else:
+                self._backend.close()
         except Exception as exc:
             err = exc
         finally:
@@ -657,13 +651,14 @@ def close(self) -> None:
         if self._closed:
             return
 
+        # Sink close only detaches the sink and returns its route lease when an engine is alive.
         err: Optional[Exception] = None
         engine = self._engine_ref()
         try:
-            _close_backend_with_optional_engine(
-                self._backend,
-                engine._backend if engine is not None else None,
-            )
+            if engine is not None:
+                self._backend.close_for_engine(engine._backend)
+            else:
+                self._backend.close()
         except Exception as exc:
             err = exc
         finally:

python_ffi/src/lib.rs

@@ -940,30 +940,11 @@ struct PyAsrSinkBackend {
     final_stats: Option<AsrSinkMetricsSnapshot>,
 }
 
-#[pymethods]
 impl PyAsrSinkBackend {
-    fn stats<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
-        let snapshot = match (&self.sink, &self.final_stats) {
-            (Some(sink), _) => sink.stats(),
-            (None, Some(snapshot)) => snapshot.clone(),
-            (None, None) => AsrSinkMetricsSnapshot::default(),
-        };
-
-        let out = PyDict::new(py);
-        for (input_id, stats) in snapshot {
-            out.set_item(
-                input_id,
-                Py::new(py, PyAsrInputStats::from_snapshot(py, stats)?)?,
-            )?;
-        }
-        Ok(out)
-    }
-
-    #[pyo3(signature = (engine=None))]
-    fn close(
+    fn finish_close(
         &mut self,
         py: Python<'_>,
-        mut engine: Option<PyRefMut<'_, PyAudioEngineBackend>>,
+        mut engine: Option<&mut PyAudioEngineBackend>,
     ) -> PyResult<()> {
         let Some(mut sink) = self.sink.take() else {
             return Ok(());
@@ -986,30 +967,49 @@ impl PyAsrSinkBackend {
                         .map(|input| (input.input_id, input.consumer))
                         .collect(),
                 )
-                .map_err(|e| PyRuntimeError::new_err(format!("failed to restore asr sink routes: {e}")))?;
+                .map_err(|e| {
+                    PyRuntimeError::new_err(format!("failed to restore asr sink routes: {e}"))
+                })?;
         }
         Ok(())
     }
+}
 
-    fn close_no_restore(&mut self, py: Python<'_>) -> PyResult<()> {
-        let Some(mut sink) = self.sink.take() else {
-            return Ok(());
+#[pymethods]
+impl PyAsrSinkBackend {
+    fn stats<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
+        let snapshot = match (&self.sink, &self.final_stats) {
+            (Some(sink), _) => sink.stats(),
+            (None, Some(snapshot)) => snapshot.clone(),
+            (None, None) => AsrSinkMetricsSnapshot::default(),
         };
 
-        let (stop_result, final_stats) = py.detach(move || {
-            let stop_result = sink.stop().map_err(|e| e.to_string());
-            let final_stats = sink.stats();
-            (stop_result, final_stats)
-        });
-        self.final_stats = Some(final_stats);
-        let _ = stop_result;
-        Ok(())
+        let out = PyDict::new(py);
+        for (input_id, stats) in snapshot {
+            out.set_item(
+                input_id,
+                Py::new(py, PyAsrInputStats::from_snapshot(py, stats)?)?,
+            )?;
+        }
+        Ok(out)
+    }
+
+    fn close(&mut self, py: Python<'_>) -> PyResult<()> {
+        self.finish_close(py, None)
+    }
+
+    fn close_for_engine(
+        &mut self,
+        py: Python<'_>,
+        mut engine: PyRefMut<'_, PyAudioEngineBackend>,
+    ) -> PyResult<()> {
+        self.finish_close(py, Some(&mut *engine))
     }
 }
 
 impl Drop for PyAsrSinkBackend {
     fn drop(&mut self) {
-        let _ = Python::try_attach(|py| self.close_no_restore(py));
+        let _ = Python::try_attach(|py| self.finish_close(py, None));
     }
 }
 
@@ -1020,22 +1020,11 @@ struct PyWavSinkBackend {
     route_ids: Vec<String>,
 }
 
-#[pymethods]
 impl PyWavSinkBackend {
-    fn stats(&self, py: Python<'_>) -> PyResult<Py<PyWavSinkStats>> {
-        let snapshot = match (&self.sink, &self.final_stats) {
-            (Some(sink), _) => sink.stats(),
-            (None, Some(snapshot)) => snapshot.clone(),
-            (None, None) => WavSinkMetricsSnapshot::default(),
-        };
-        Py::new(py, PyWavSinkStats::from_snapshot(py, snapshot)?)
-    }
-
-    #[pyo3(signature = (engine=None))]
-    fn close(
+    fn finish_close(
         &mut self,
         py: Python<'_>,
-        mut engine: Option<PyRefMut<'_, PyAudioEngineBackend>>,
+        mut engine: Option<&mut PyAudioEngineBackend>,
     ) -> PyResult<()> {
         let Some(mut sink) = self.sink.take() else {
             return Ok(());
@@ -1054,30 +1043,41 @@ impl PyWavSinkBackend {
         if let Some(engine) = engine.as_mut() {
             engine
                 .restore_route_consumers(route_ids.into_iter().zip(consumers).collect())
-                .map_err(|e| PyRuntimeError::new_err(format!("failed to restore wav sink routes: {e}")))?;
+                .map_err(|e| {
+                    PyRuntimeError::new_err(format!("failed to restore wav sink routes: {e}"))
+                })?;
         }
         Ok(())
     }
+}
 
-    fn close_no_restore(&mut self, py: Python<'_>) -> PyResult<()> {
-        let Some(mut sink) = self.sink.take() else {
-            return Ok(());
+#[pymethods]
+impl PyWavSinkBackend {
+    fn stats(&self, py: Python<'_>) -> PyResult<Py<PyWavSinkStats>> {
+        let snapshot = match (&self.sink, &self.final_stats) {
+            (Some(sink), _) => sink.stats(),
+            (None, Some(snapshot)) => snapshot.clone(),
+            (None, None) => WavSinkMetricsSnapshot::default(),
         };
+        Py::new(py, PyWavSinkStats::from_snapshot(py, snapshot)?)
+    }
 
-        let (stop_result, final_stats) = py.detach(move || {
-            let stop_result = sink.stop().map_err(|e| e.to_string());
-            let final_stats = sink.stats();
-            (stop_result, final_stats)
-        });
-        self.final_stats = Some(final_stats);
-        let _ = stop_result;
-        Ok(())
+    fn close(&mut self, py: Python<'_>) -> PyResult<()> {
+        self.finish_close(py, None)
+    }
+
+    fn close_for_engine(
+        &mut self,
+        py: Python<'_>,
+        mut engine: PyRefMut<'_, PyAudioEngineBackend>,
+    ) -> PyResult<()> {
+        self.finish_close(py, Some(&mut *engine))
     }
 }
 
 impl Drop for PyWavSinkBackend {
     fn drop(&mut self) {
-        let _ = Python::try_attach(|py| self.close_no_restore(py));
+        let _ = Python::try_attach(|py| self.finish_close(py, None));
     }
 }
 

ref.md

@@ -0,0 +1,83 @@
+# Shutdown Refactor Plan
+
+## Context
+- The recent fixes made shutdown reliable again: synthetic hangs are gone, native source shutdown is bounded, and the real-source rerun repro now completes.
+- The remaining concern is architectural: stop/close semantics currently work, but they are spread across Python high-level API, Python FFI, and Rust sink/source implementations in a way that feels patchy.
+- In particular, the desired product semantics are now clearer than the code structure: **shutdown should be bounded best-effort cleanup, not guaranteed full pipeline delivery or full drain**.
+- The current code still contains compatibility-oriented and cross-layer shutdown logic that obscures that model.
+- User preference: preserve only the high-level Python API contract; low-level/internal `_macloop` shutdown compatibility may be simplified or broken if that leads to a cleaner design.
+- User preference for high-level close behavior: if a native source does not stop within the bounded timeout, `AudioEngine.close()` should still complete without hanging, but it should surface a warning/log signal rather than silently succeeding or raising by default.
+
+## Approach
+Refactor shutdown around a simpler lifecycle contract:
+1. **Engine owns source lifecycle** — native sources start/stop on the engine lifecycle, not on sink lifecycle.
+2. **Sink close only detaches the sink** — bounded worker stop + route lease/consumer return, but no source orchestration.
+3. **Bounded best-effort close is explicit** — native stop timeouts are part of the normal close model, not an exceptional hidden path.
+4. **High-level close warns on deferred native cleanup** — `AudioEngine.close()` remains non-hanging and non-fatal by default, but emits a warning/log when native cleanup is deferred.
+5. **Cross-layer close API becomes explicit** — remove signature probing / fallback behavior and replace it with one clear sink-close contract.
+6. **Allow small semantic cleanup** — internal/FFI shutdown contracts may be cleaned up rather than preserved verbatim if that materially simplifies the design, as long as the high-level Python API remains stable.
+
+This should be done incrementally, preserving user-visible reliability while shrinking the shutdown state machine.
+4. **Cross-layer close API becomes explicit** — remove signature probing / fallback behavior and replace it with one clear sink-close contract.
+5. **Allow small semantic cleanup** — internal/FFI shutdown contracts may be cleaned up rather than preserved verbatim if that materially simplifies the design, as long as the high-level Python API remains stable.
+
+This should be done incrementally, preserving user-visible reliability while shrinking the shutdown state machine.
+
+## Files to modify
+- `macloop/__init__.py`
+- `python_ffi/src/lib.rs`
+- `core_engine/src/outputs/asr_sink.rs`
+- `core_engine/src/outputs/wav_file.rs`
+- `core_engine/src/engine.rs`
+- `tests/test_runtime_helpers.py`
+- `tests/test_e2e_synthetic.py`
+- `tests/test_medium_e2e_real_capture.py`
+- possibly `tests/test_ffi_backend.py`
+
+## Reuse
+- Bounded native lifecycle helpers already exist in `python_ffi/src/lib.rs`:
+  - `start_native_source_with_deadline(...)`
+  - `stop_native_source_with_deadline(...)`
+  - `cleanup_pending_cleanups_with_deadline(...)`
+- Route consumer handoff/restore already exists and should be reused rather than re-invented:
+  - `AudioEngineController::take_output_consumer(...)` in `core_engine/src/engine.rs`
+  - `AudioEngineController::restore_output_consumer(...)` in `core_engine/src/engine.rs`
+  - `restore_route_consumers(...)` in `python_ffi/src/lib.rs`
+- Bounded worker shutdown patterns already exist in Rust sinks and should stay as the basis:
+  - bounded polling in `core_engine/src/outputs/asr_sink.rs`
+  - bounded polling in `core_engine/src/outputs/wav_file.rs`
+- Existing regression coverage to preserve during refactor:
+  - `tests/test_e2e_synthetic.py::test_hot_synthetic_engine_close_completes_with_wav_and_asr_sinks`
+  - `tests/test_e2e_synthetic.py::test_hot_synthetic_restart_on_single_engine_does_not_hang`
+  - `tests/test_medium_e2e_real_capture.py::test_medium_real_source_asr_rerun_engine_close_remains_stable`
+
+## Steps
+- [ ] Step 1: Document and codify the target lifecycle contract in code comments/tests:
+  - engine close = bounded best-effort source shutdown
+  - sink close = bounded detach + route return
+  - close does not guarantee full drain/delivery
+- [ ] Step 2: Simplify Python high-level close flow in `macloop/__init__.py` so it no longer performs feature detection via exception shape/signature fallback, and define one clear warning/log path for deferred native cleanup.
+- [ ] Step 3: Simplify FFI sink close API in `python_ffi/src/lib.rs` into one explicit contract for route restoration instead of optional engine-aware probing, even if that changes low-level/internal `_macloop` shutdown semantics.
+- [ ] Step 4: Keep source lifecycle orchestration centralized in `PyAudioEngineBackend`, and remove any remaining sink/source lifecycle coupling that is only there for historical shutdown bugs.
+- [ ] Step 5: Separate explicit sink shutdown semantics from fallback/drop cleanup semantics in Rust sink implementations, so `Drop` remains best-effort cleanup and explicit close remains the ownership-return path.
+- [ ] Step 6: Revisit route ownership bookkeeping and reduce duplicate state where possible (Python `_claimed_routes` vs backend consumer availability), without regressing same-engine restart behavior.
+- [ ] Step 7: Tighten and align tests around the simplified model, especially:
+  - backend timeout is non-hanging and best-effort
+  - high-level `engine.close()` emits the intended warning/log when native cleanup is deferred
+  - same-engine restart still works
+  - real-source rerun stays stable
+  - low-level tests are updated to validate the new internal contract only where still intentionally supported
+
+## Verification
+- Rust tests:
+  - `cargo test -p python_ffi --quiet`
+  - `cargo test -p core_engine --quiet`
+- Python tests:
+  - `.venv/bin/python -m pytest -q tests`
+- Focused regressions:
+  - `.venv/bin/python -m pytest -q tests/test_e2e_synthetic.py::test_hot_synthetic_restart_on_single_engine_does_not_hang`
+  - `.venv/bin/python -m pytest -q tests/test_medium_e2e_real_capture.py::test_medium_real_source_asr_rerun_engine_close_remains_stable --run-medium`
+- Manual reasoning check after refactor:
+  - sink close no longer starts/stops sources
+  - engine close remains bounded even if native source stop times out
+  - route reuse still works on the same engine after explicit sink close

tests/conftest.py

@@ -126,6 +126,9 @@ def stats(self):
     def close(self) -> None:
         self.closed = True
 
+    def close_for_engine(self, _engine) -> None:
+        self.close()
+
 
 class _FakeWavSinkBackend:
     def __init__(self) -> None:
@@ -166,6 +169,9 @@ def stats(self):
     def close(self) -> None:
         self.closed = True
 
+    def close_for_engine(self, _engine) -> None:
+        self.close()
+
 
 class _FakeAudioEngineBackend:
     def __init__(self) -> None: