Simplify source and route shutdown ownership

Author: kemsta

core_engine/src/outputs/asr_sink.rs

@@ -214,8 +214,10 @@ impl InputState {
         let input_channels = MASTER_FORMAT.channels.max(1) as usize;
         let ready_input_samples = self.drained_master.len() / input_channels * input_channels;
         if ready_input_samples > 0 {
-            self.converter
-                .convert(&self.drained_master[..ready_input_samples], &mut self.converted_output)?;
+            self.converter.convert(
+                &self.drained_master[..ready_input_samples],
+                &mut self.converted_output,
+            )?;
             if !self.converted_output.is_empty() {
                 self.pending_output
                     .extend_from_slice(&self.converted_output);
@@ -452,6 +454,14 @@ impl AsrSink {
             Err(_) => Err(AsrSinkError::ThreadPanic),
         }
     }
+
+    fn abort_on_drop(&mut self) {
+        self.stop.store(true, Ordering::Relaxed);
+        let Some(handle) = self.handle.take() else {
+            return;
+        };
+        let _ = handle.join();
+    }
 }
 
 fn duration_to_u32_us(duration: Duration) -> u32 {
@@ -460,7 +470,7 @@ fn duration_to_u32_us(duration: Duration) -> u32 {
 
 impl Drop for AsrSink {
     fn drop(&mut self) {
-        let _ = self.stop();
+        self.abort_on_drop();
     }
 }
 

core_engine/src/outputs/wav_file.rs

@@ -146,11 +146,7 @@ impl WavFileOutput {
                     }
                 }
 
-                let ready_samples = input_buffers
-                    .iter()
-                    .map(VecDeque::len)
-                    .min()
-                    .unwrap_or(0)
+                let ready_samples = input_buffers.iter().map(VecDeque::len).min().unwrap_or(0)
                     / frame_channels
                     * frame_channels;
 
@@ -312,6 +308,14 @@ impl WavFileOutput {
             Err(_) => Err(WavOutputError::ThreadPanic),
         }
     }
+
+    fn abort_on_drop(&mut self) {
+        self.stop.store(true, Ordering::Relaxed);
+        let Some(handle) = self.handle.take() else {
+            return;
+        };
+        let _ = handle.join();
+    }
 }
 
 fn duration_to_u32_us(duration: Duration) -> u32 {
@@ -320,7 +324,7 @@ fn duration_to_u32_us(duration: Duration) -> u32 {
 
 impl Drop for WavFileOutput {
     fn drop(&mut self) {
-        let _ = self.stop();
+        self.abort_on_drop();
     }
 }
 

macloop/__init__.py

@@ -283,7 +283,6 @@ def __init__(self) -> None:
         self._streams: dict[str, StreamHandle] = {}
         self._routes: dict[str, RouteHandle] = {}
         self._processors: dict[str, ProcessorHandle] = {}
-        self._claimed_routes: set[str] = set()
         self._sink_refs: list[weakref.ReferenceType[_Closable]] = []
         self._closed = False
 
@@ -384,7 +383,6 @@ def close(self) -> None:
                         sink_err = exc
 
             self._sink_refs.clear()
-            self._claimed_routes.clear()
 
         if backend_err is not None:
             raise backend_err
@@ -430,17 +428,12 @@ def _ensure_routes_available(self, routes: Sequence[RouteHandle]) -> None:
         if not routes:
             raise ValueError("routes must not be empty")
 
+        route_ids = []
         for route in routes:
             self._ensure_route_handle(route)
-            if route.id in self._claimed_routes:
-                raise ValueError(f"route '{route.id}' is already in use")
+            route_ids.append(route.id)
 
-    def _claim_routes(self, route_ids: Sequence[str]) -> None:
-        self._claimed_routes.update(route_ids)
-
-    def _release_routes(self, route_ids: Sequence[str]) -> None:
-        for route_id in route_ids:
-            self._claimed_routes.discard(route_id)
+        self._backend.assert_routes_available(route_ids)
 
     def _register_sink(self, sink: _Closable) -> None:
         self._sink_refs.append(weakref.ref(sink))
@@ -472,7 +465,6 @@ def __init__(
 
         self.id = sink_id
         self._engine_ref = weakref.ref(engine)
-        self._route_ids = tuple(route_ids)
         self._queue = out_queue
         self._queue_maxsize = max_queue_size
         self._async_queue: Optional[asyncio.Queue[object]] = None
@@ -496,7 +488,6 @@ def _callback(route_id: str, frames: int, samples: AudioSamples) -> None:
 
         self._backend = backend
 
-        engine._claim_routes(route_ids)
         engine._register_sink(self)
 
     def chunks(self) -> Iterator[AudioChunk]:
@@ -531,8 +522,6 @@ def close(self) -> None:
             err = exc
         finally:
             self._closed = True
-            if engine is not None:
-                engine._release_routes(self._route_ids)
             _drop_oldest_put(self._queue, _STOP)
             async_loop = self._async_loop
             async_queue = self._async_queue
@@ -641,10 +630,8 @@ def __init__(
         self.id = sink_id
         self._backend = backend
         self._engine_ref = weakref.ref(engine)
-        self._route_ids = tuple(route_ids)
         self._closed = False
 
-        engine._claim_routes(route_ids)
         engine._register_sink(self)
 
     def close(self) -> None:
@@ -663,8 +650,6 @@ def close(self) -> None:
             err = exc
         finally:
             self._closed = True
-            if engine is not None:
-                engine._release_routes(self._route_ids)
 
         if err is not None:
             raise err

python_ffi/src/lib.rs

@@ -1336,6 +1336,11 @@ impl PyAudioEngineBackend {
         Ok(out)
     }
 
+    fn assert_routes_available(&self, route_ids: Vec<String>) -> PyResult<()> {
+        self.ensure_open()?;
+        self.ensure_route_consumers_available(&route_ids)
+    }
+
     fn close(&mut self, py: Python<'_>) -> PyResult<()> {
         // `closed` means no new public operations are allowed. If native cleanup timed out earlier,
         // repeated close attempts are still allowed so drop/explicit close can keep retrying the
@@ -1469,6 +1474,44 @@ impl PyAudioEngineBackend {
         Ok(())
     }
 
+    fn start_stream_states_for_routes(
+        &mut self,
+        py: Python<'_>,
+        route_ids: &[String],
+    ) -> PyResult<()> {
+        let stream_ids = self.stream_ids_for_routes(route_ids)?;
+        let stream_states = self.take_stream_states(&stream_ids)?;
+        let started_states = match py.detach(move || {
+            start_stream_states_with_timeout(stream_states, NATIVE_SOURCE_START_TIMEOUT)
+        }) {
+            Ok(states) => states,
+            Err(StartStreamsError {
+                message,
+                states,
+                poison,
+                cleanup,
+            }) => {
+                self.restore_stream_states(states);
+                self.store_pending_cleanups(cleanup);
+                self.poisoned = poison;
+                return Err(PyRuntimeError::new_err(message));
+            }
+        };
+        self.restore_stream_states(started_states);
+        Ok(())
+    }
+
+    fn prepare_route_consumers_for_sink(
+        &mut self,
+        py: Python<'_>,
+        route_ids: &[String],
+    ) -> PyResult<Vec<(String, RouteConsumer)>> {
+        self.ensure_open()?;
+        self.ensure_route_consumers_available(route_ids)?;
+        self.start_stream_states_for_routes(py, route_ids)?;
+        self.take_route_consumers(route_ids)
+    }
+
     fn stream_ids_for_routes(&self, route_ids: &[String]) -> PyResult<Vec<String>> {
         let mut stream_ids = Vec::<String>::new();
 
@@ -1554,9 +1597,6 @@ impl PyAudioEngineBackend {
         chunk_frames: usize,
         callback: Py<PyAny>,
     ) -> PyResult<PyAsrSinkBackend> {
-        self.ensure_open()?;
-        self.ensure_route_consumers_available(&route_ids)?;
-
         let format = StreamFormat::with_sample_format(
             sample_rate,
             channels,
@@ -1568,27 +1608,7 @@ impl PyAudioEngineBackend {
         })
         .map_err(|err| PyValueError::new_err(err.to_string()))?;
 
-        let stream_ids = self.stream_ids_for_routes(&route_ids)?;
-        let stream_states = self.take_stream_states(&stream_ids)?;
-        let started_states = match py.detach(move || {
-            start_stream_states_with_timeout(stream_states, NATIVE_SOURCE_START_TIMEOUT)
-        }) {
-            Ok(states) => states,
-            Err(StartStreamsError {
-                message,
-                states,
-                poison,
-                cleanup,
-            }) => {
-                self.restore_stream_states(states);
-                self.store_pending_cleanups(cleanup);
-                self.poisoned = poison;
-                return Err(PyRuntimeError::new_err(message));
-            }
-        };
-        self.restore_stream_states(started_states);
-
-        let route_consumers = self.take_route_consumers(&route_ids)?;
+        let route_consumers = self.prepare_route_consumers_for_sink(py, &route_ids)?;
         let detached_result: DetachedAsrStartResult = py.detach(move || {
             let inputs = route_consumers
                 .into_iter()
@@ -1640,32 +1660,9 @@ impl PyAudioEngineBackend {
         fd: i32,
         mix_gain: f32,
     ) -> PyResult<PyWavSinkBackend> {
-        self.ensure_open()?;
-        self.ensure_route_consumers_available(&route_ids)?;
-
         let file = duplicate_file_descriptor(fd)
             .map_err(|e| PyOSError::new_err(format!("failed to duplicate file descriptor: {e}")))?;
-        let stream_ids = self.stream_ids_for_routes(&route_ids)?;
-        let stream_states = self.take_stream_states(&stream_ids)?;
-        let started_states = match py.detach(move || {
-            start_stream_states_with_timeout(stream_states, NATIVE_SOURCE_START_TIMEOUT)
-        }) {
-            Ok(states) => states,
-            Err(StartStreamsError {
-                message,
-                states,
-                poison,
-                cleanup,
-            }) => {
-                self.restore_stream_states(states);
-                self.store_pending_cleanups(cleanup);
-                self.poisoned = poison;
-                return Err(PyRuntimeError::new_err(message));
-            }
-        };
-        self.restore_stream_states(started_states);
-
-        let route_consumers = self.take_route_consumers(&route_ids)?;
+        let route_consumers = self.prepare_route_consumers_for_sink(py, &route_ids)?;
         let route_ids_for_sink = route_ids.clone();
         let master_format = self.controller.master_format();
         let detached_result: DetachedWavStartResult = py.detach(move || {

tests/conftest.py

@@ -116,23 +116,26 @@ def __init__(self, write_calls, samples_written, frames_written, write, finalize
 
 
 class _FakeAsrSinkBackend:
-    def __init__(self) -> None:
+    def __init__(self, route_ids: tuple[str, ...]) -> None:
         self.closed = False
         self._stats = {}
+        self._route_ids = route_ids
 
     def stats(self):
         return self._stats
 
     def close(self) -> None:
         self.closed = True
 
-    def close_for_engine(self, _engine) -> None:
+    def close_for_engine(self, engine) -> None:
         self.close()
+        engine.available_routes.update(self._route_ids)
 
 
 class _FakeWavSinkBackend:
-    def __init__(self) -> None:
+    def __init__(self, route_ids: tuple[str, ...]) -> None:
         self.closed = False
+        self._route_ids = route_ids
         write_latency = _FakeLatencyStats(
             last_us=18,
             max_us=30,
@@ -169,14 +172,16 @@ def stats(self):
     def close(self) -> None:
         self.closed = True
 
-    def close_for_engine(self, _engine) -> None:
+    def close_for_engine(self, engine) -> None:
         self.close()
+        engine.available_routes.update(self._route_ids)
 
 
 class _FakeAudioEngineBackend:
     def __init__(self) -> None:
         self.calls: list[tuple[object, ...]] = []
         self.closed = False
+        self.available_routes: set[str] = set()
         pipeline_latency = _FakeLatencyStats(
             last_us=10,
             max_us=20,
@@ -216,10 +221,16 @@ def add_processor(self, stream_id, processor_id, processor_kind, config) -> None
 
     def route(self, route_id, stream_id) -> None:
         self.calls.append(("route", route_id, stream_id))
+        self.available_routes.add(route_id)
 
     def get_stats(self):
         return self.stats
 
+    def assert_routes_available(self, route_ids) -> None:
+        for route_id in route_ids:
+            if route_id not in self.available_routes:
+                raise ValueError(f"route '{route_id}' is not available")
+
     def close(self) -> None:
         self.closed = True
 
@@ -267,7 +278,9 @@ def _fake_create_asr_sink(
         p95_us=64,
         p99_us=64,
     )
-    backend = _FakeAsrSinkBackend()
+    for route_id in route_ids:
+        engine.available_routes.discard(route_id)
+    backend = _FakeAsrSinkBackend(tuple(route_ids))
     backend._stats = {
         route_ids[0]: _FakeAsrInputStats(
             chunks_emitted=2,
@@ -284,7 +297,9 @@ def _fake_create_asr_sink(
 
 def _fake_create_wav_sink(engine, sink_id, route_ids, fd, mix_gain):
     engine.calls.append(("create_wav_sink", sink_id, tuple(route_ids), fd, mix_gain))
-    return _FakeWavSinkBackend()
+    for route_id in route_ids:
+        engine.available_routes.discard(route_id)
+    return _FakeWavSinkBackend(tuple(route_ids))
 
 
 @pytest.fixture