Improve audio buffer and thread shutdown (#117)

Limit the audio buffer size and harden start/stop behavior for audio input and the WebSocket client. AudioInputStream: set buffer maxsize to 5, clear buffer on start, set running=true, lazily initialize Zenoh session/pub/sub, open PyAudio stream only if not already open, and improve stop() to signal the consumer, join the audio thread with timeout and log termination, and clean up session/pub. Client: track a main client_thread, avoid starting it twice, set running on start, join and log on shutdown, and reset thread references. Tests updated to match the new lifecycle (removed expectations for immediate PyAudio termination/stream close and adjusted assertions).

Author: openminddev

src/om1_speech/audio/audio_input_stream.py

@@ -92,7 +92,7 @@ def __init__(
         self._enable_tts_interrupt = enable_tts_interrupt
 
         # Thread-safe buffer for audio data
-        self._buff: queue.Queue[Optional[bytes]] = queue.Queue()
+        self._buff: queue.Queue[Optional[bytes]] = queue.Queue(maxsize=5)
 
         # audio interface and stream
         self._audio_interface: pyaudio.PyAudio = pyaudio.PyAudio()
@@ -293,36 +293,53 @@ def start(self) -> "AudioInputStream":
         Exception
             If there are errors initializing the audio interface or opening the stream
         """
-        if not self.running:
-            return self
+        self.running = True
+
+        while not self._buff.empty():
+            try:
+                self._buff.get_nowait()
+            except queue.Empty:
+                break
 
         try:
+            if self.session is None:
+                try:
+                    self.session = open_zenoh_session()
+                    self.pub = self.session.declare_publisher(self.topic)
+                    self.session.declare_subscriber(
+                        self.topic, self.zenoh_audio_message
+                    )
+                    logger.info("Zenoh session initialized")
+                except Exception as e:
+                    logger.error(f"Failed to initialize Zenoh: {e}")
+                    self.session = None
+
             if not self.remote_input:
                 logger.info("Remote input is disabled, initializing audio input stream")
                 if self._rate is None:
                     self._rate = 16000
                 if self._chunk is None:
                     self._chunk = int(self._rate * 0.2)
-                self._audio_stream = self._audio_interface.open(
-                    format=pyaudio.paInt16,
-                    input_device_index=(
-                        int(self._device) if self._device is not None else None
-                    ),
-                    channels=1,
-                    rate=self._rate,
-                    input=True,
-                    frames_per_buffer=self._chunk,
-                    stream_callback=self._fill_buffer,  # type: ignore
-                )
 
-            # Start the audio processing thread
+                if self._audio_stream is None:
+                    self._audio_stream = self._audio_interface.open(
+                        format=pyaudio.paInt16,
+                        input_device_index=(
+                            int(self._device) if self._device is not None else None
+                        ),
+                        channels=1,
+                        rate=self._rate,
+                        input=True,
+                        frames_per_buffer=self._chunk,
+                        stream_callback=self._fill_buffer,  # type: ignore
+                    )
+
             self._start_audio_thread()
 
             logger.info(f"Started audio stream with device {self._device}")
 
         except Exception as e:
             logger.error(f"Error opening audio stream: {e}")
-            self._audio_interface.terminate()
             raise
 
         return self
@@ -444,25 +461,23 @@ def on_audio(self):
     def stop(self):
         """
         Stops the audio stream and cleans up resources.
-
-        This method stops the audio stream, terminates the PyAudio interface,
-        and ensures the audio processing thread is properly shut down.
         """
         self.running = False
 
-        if self.session:
-            self.session.close()
+        self._buff.put(None)
 
-        if self._audio_stream:
-            self._audio_stream.stop_stream()
-            self._audio_stream.close()
+        if self._audio_thread and self._audio_thread.is_alive():
+            self._audio_thread.join(timeout=2.0)
+            if self._audio_thread.is_alive():
+                logger.warning("Audio processing thread did not terminate gracefully")
+            else:
+                logger.info("Audio processing thread terminated successfully")
 
-        if self._audio_interface:
-            self._audio_interface.terminate()
+        self._audio_thread = None
 
-        # Clean up the audio processing thread
-        if self._audio_thread and self._audio_thread.is_alive():
-            self._audio_thread.join(timeout=1.0)
+        if self.session:
+            self.session.close()
+            self.session = None
+            self.pub = None
 
-        self._buff.put(None)
         logger.info("Stopped audio stream")

src/om1_utils/ws/client.py

@@ -34,6 +34,7 @@ def __init__(self, url: str = "ws://localhost:6789"):
         self.message_queue: Queue = Queue()
         self.receiver_thread: Optional[threading.Thread] = None
         self.sender_thread: Optional[threading.Thread] = None
+        self.client_thread: Optional[threading.Thread] = None
 
     def _receive_messages(self):
         """
@@ -164,6 +165,11 @@ def start(self):
         Initializes and starts the main client thread that manages the WebSocket
         connection.
         """
+        if self.client_thread and self.client_thread.is_alive():
+            logger.warning("WebSocket client thread is already running")
+            return
+
+        self.running = True
         self.client_thread = threading.Thread(target=self._run_client, daemon=True)
         self.client_thread.start()
         logger.info("WebSocket client thread started")
@@ -232,6 +238,17 @@ def stop(self):
             except Exception as _:
                 pass
 
+        if self.client_thread and self.client_thread.is_alive():
+            self.client_thread.join(timeout=2.0)
+            if self.client_thread.is_alive():
+                logger.warning("Client thread did not terminate gracefully")
+            else:
+                logger.info("Client thread stopped")
+
+        self.client_thread = None
+        self.receiver_thread = None
+        self.sender_thread = None
+
         try:
             while True:
                 self.message_queue.get_nowait()

tests/om1_speech/audio/test_audio_input_stream.py

@@ -23,7 +23,6 @@ def mock_pyaudio():
         mock.return_value.get_default_input_device_info.return_value = default_device
         mock.return_value.get_device_info_by_index.return_value = default_device
 
-        # Mock stream
         mock_stream = MagicMock()
         mock_stream.stop_stream = Mock()
         mock_stream.close = Mock()
@@ -62,13 +61,8 @@ def test_start_with_default_device(audio_stream, mock_pyaudio):
     """Test starting AudioInputStream with default device"""
     audio_stream.start()
 
-    # Verify PyAudio initialization
     mock_pyaudio.assert_called_once()
-
-    # Verify default device was retrieved
     mock_pyaudio.return_value.get_default_input_device_info.assert_called_once()
-
-    # Verify stream was opened with correct parameters
     mock_pyaudio.return_value.open.assert_called_once_with(
         format=pyaudio.paInt16,
         input_device_index=0,
@@ -113,7 +107,6 @@ def test_fill_buffer_with_tts_inactive(audio_stream):
     test_data = b"test_audio_data"
     audio_stream._fill_buffer(test_data, 1024, {}, 0)
 
-    # Verify data was added to buffer
     assert audio_stream._buff.get() == test_data
 
 
@@ -123,24 +116,18 @@ def test_fill_buffer_with_tts_active(audio_stream):
     test_data = b"test_audio_data"
     audio_stream._fill_buffer(test_data, 1024, {}, 0)
 
-    # Verify buffer is empty (data wasn't added)
     with pytest.raises(queue.Empty):
         audio_stream._buff.get_nowait()
 
 
 def test_generator(audio_stream):
     """Test audio data generation"""
-    # Ensure the stream starts in a clean state
     audio_stream.running = True
-
-    # Create test chunks
     test_chunks = [b"chunk1", b"chunk2", b"chunk3"]
 
-    # Add test chunks to buffer
     for chunk in test_chunks:
         audio_stream._buff.put(chunk)
 
-    # Start collecting in a separate thread to avoid blocking
     collected_chunks = []
 
     def collect_data():
@@ -149,19 +136,15 @@ def collect_data():
             if len(collected_chunks) >= len(test_chunks):
                 break
 
-    # Run collection in thread
     collection_thread = threading.Thread(target=collect_data)
     collection_thread.daemon = True
     collection_thread.start()
 
-    # Wait a short time for collection
     collection_thread.join(timeout=1.0)
 
-    # Stop the generator properly
     audio_stream.running = False
     audio_stream._buff.put(None)
 
-    # Verify the results
     assert len(collected_chunks) > 0
     assert all(isinstance(data, dict) for data in collected_chunks)
 
@@ -171,11 +154,9 @@ def test_stop(audio_stream, mock_pyaudio):
     audio_stream.start()
     audio_stream.stop()
 
-    # Verify stream was stopped and closed
+    # Verify stream was stopped
     assert audio_stream.running is False
-    mock_pyaudio.return_value.open.return_value.stop_stream.assert_called_once()
-    mock_pyaudio.return_value.open.return_value.close.assert_called_once()
-    mock_pyaudio.return_value.terminate.assert_called_once()
+    assert audio_stream._audio_thread is None
 
 
 def test_audio_callback(mock_pyaudio):
@@ -220,11 +201,10 @@ def test_error_handling(mock_pyaudio):
     mock_pyaudio.return_value.open.side_effect = Exception("Test error")
 
     stream = AudioInputStream()
-    with pytest.raises(Exception):
+    with pytest.raises(Exception, match="Test error"):
         stream.start()
 
-    # Verify cleanup was performed
-    mock_pyaudio.return_value.terminate.assert_called_once()
+    assert stream._audio_stream is None
 
     stream.stop()