import numpy as np import threading from typing import Optional from dataclasses import dataclass from ..config import CHUNK_SAMPLES, BUFFER_SAMPLES, SHIFT_SAMPLES, settings @dataclass class BufferStats: total_written: int = 0 total_read: int = 0 overruns: int = 0 underruns: int = 0 current_fill: float = 0.0 window_size: int = 0 step_size: int = 0 class RingBuffer: """ Thread-safe rolling FIFO ring buffer for 16kHz mono audio. Capacity: BUFFER_SAMPLES (the model's full input window, 24,000 samples/ 1.5s @ 16kHz). Each inference consumes the latest full window and advances the read position by SHIFT_SAMPLES (configurable rolling step). """ def __init__( self, capacity: int = BUFFER_SAMPLES, window: int = BUFFER_SAMPLES, shift: int = SHIFT_SAMPLES, ): self._capacity = capacity self._window = window self._shift = shift self._buffer = np.zeros(capacity, dtype=np.float32) self._write_pos = 0 self._read_pos = 0 self._filled = 0 self._lock = threading.RLock() self._stats = BufferStats() @property def capacity(self) -> int: return self._capacity @property def window_size(self) -> int: return self._window @property def shift(self) -> int: return self._shift @property def filled(self) -> int: with self._lock: return self._filled @property def fill_ratio(self) -> float: with self._lock: return self._filled / self._capacity @property def stats(self) -> BufferStats: with self._lock: self._stats.current_fill = self.fill_ratio return BufferStats( total_written=self._stats.total_written, total_read=self._stats.total_read, overruns=self._stats.overruns, underruns=self._stats.underruns, current_fill=self._stats.current_fill, ) def write(self, data: np.ndarray) -> int: """ Write audio data to buffer. Returns number of samples written. Handles wrap-around and overrun detection. """ if data.ndim != 1: raise ValueError(f"Expected 1D array, got {data.ndim}D") data = data.astype(np.float32, copy=False) n_samples = len(data) with self._lock: if n_samples > self._capacity: data = data[-self._capacity:] n_samples = self._capacity self._stats.overruns += 1 remaining = self._capacity - self._write_pos if n_samples <= remaining: self._buffer[self._write_pos:self._write_pos + n_samples] = data self._write_pos = (self._write_pos + n_samples) % self._capacity else: self._buffer[self._write_pos:] = data[:remaining] self._buffer[:n_samples - remaining] = data[remaining:] self._write_pos = (self._write_pos + n_samples) % self._capacity self._filled = min(self._filled + n_samples, self._capacity) self._stats.total_written += n_samples self._stats.current_fill = self.fill_ratio return n_samples def read_latest(self, n_samples: int) -> np.ndarray: """ Read the most recent n_samples from buffer (zero-copy view when possible). Returns zero-padded array if insufficient data. """ with self._lock: if n_samples > self._capacity: n_samples = self._capacity if self._filled < n_samples: self._stats.underruns += 1 result = np.zeros(n_samples, dtype=np.float32) available = self._filled if available > 0: start = (self._write_pos - available) % self._capacity if start + available <= self._capacity: result[-available:] = self._buffer[start:start + available] else: first_part = self._capacity - start result[-available:-available + first_part] = self._buffer[start:] result[-available + first_part:] = self._buffer[:available - first_part] return result start = (self._write_pos - n_samples) % self._capacity if start + n_samples <= self._capacity: result = self._buffer[start:start + n_samples].copy() else: first_part = self._capacity - start result = np.concatenate([ self._buffer[start:], self._buffer[:n_samples - first_part] ]) self._stats.total_read += n_samples return result def get_processing_window(self) -> Optional[np.ndarray]: """ Get the latest full processing window (model window) for inference. Returns the most recent `window_size` samples (a full 1.5s window under defaults). Returns None if insufficient data accumulated. Advances the read position by `shift` samples after a successful read so windows roll forward by the configured step. """ with self._lock: n = min(self._window, self._capacity) if self._filled < n: return None start = (self._write_pos - n) % self._capacity if start + n <= self._capacity: window = self._buffer[start:start + n].copy() else: first_part = self._capacity - start window = np.concatenate([ self._buffer[start:], self._buffer[:n - first_part] ]) step = min(self._shift, self._capacity) self._read_pos = (self._read_pos + step) % self._capacity self._filled = max(0, self._filled - step) self._stats.total_read += n self._stats.current_fill = self.fill_ratio return window def clear(self) -> None: """Reset buffer to empty state.""" with self._lock: self._buffer.fill(0.0) self._write_pos = 0 self._read_pos = 0 self._filled = 0 self._stats = BufferStats() def is_ready_for_inference(self) -> bool: """Check if buffer has enough data for one inference window.""" with self._lock: return self._filled >= min(self._window, self._capacity) class MultiConnectionBufferManager: """Manages separate ring buffers per WebSocket connection.""" def __init__(self, max_connections: int = settings.max_connections): self._max_connections = max_connections self._buffers: dict[str, RingBuffer] = {} self._lock = threading.RLock() def get_buffer(self, connection_id: str) -> RingBuffer: with self._lock: if connection_id not in self._buffers: if len(self._buffers) >= self._max_connections: raise RuntimeError(f"Max connections ({self._max_connections}) reached") self._buffers[connection_id] = RingBuffer() return self._buffers[connection_id] def remove_buffer(self, connection_id: str) -> None: with self._lock: self._buffers.pop(connection_id, None) def get_all_stats(self) -> dict[str, BufferStats]: with self._lock: return {cid: buf.stats for cid, buf in self._buffers.items()} def cleanup_stale(self, active_ids: set[str]) -> int: with self._lock: stale = set(self._buffers.keys()) - active_ids for cid in stale: del self._buffers[cid] return len(stale)