#!/usr/bin/env python3 """ Async TTS Job Pipeline (Stdlib Only) ------------------------------------ Provides bounded concurrency, backpressure, exponential retry, clean cancellation, and lifecycle callbacks. """ import asyncio import random import uuid from dataclasses import dataclass from typing import Callable, Any, Dict, List # --------------------------------------------------------------------------- # Mock Synthesizer # --------------------------------------------------------------------------- async def mock_synthesize(text: str) -> bytes: """ Simulates TTS synthesis: - Sleeps 50-300ms - Fails ~10% of the time - Returns dummy bytes on success """ await asyncio.sleep(random.uniform(0.05, 0.30)) if random.random() < 0.10: raise RuntimeError("Synthesis failed (mock network/device error)") return b"\x00" * len(text) # --------------------------------------------------------------------------- # Core Pipeline # --------------------------------------------------------------------------- @dataclass class Job: id: str text: str voice: str state: str = "queued" attempts: int = 0 class TTSPipeline: def __init__( self, max_concurrency: int = 4, max_queue_size: int = 100, max_retries: int = 3, ): self._queue: asyncio.Queue = asyncio.Queue() self._semaphore = asyncio.Semaphore(max_concurrency) self._jobs: Dict[str, Job] = {} self._jobs_lock = asyncio.Lock() self._max_queue_size = max_queue_size self._max_retries = max_retries self._callbacks: List[Callable[[str, str], Any]] = [] self._callbacks_lock = asyncio.Lock() # Concurrency tracking self._active_count = 0 self._max_concurrency_seen = 0 self._concurrency_lock = asyncio.Lock() # Worker management self._worker_task: asyncio.Task | None = None self._running = False # Final state tracking (for testing/observability) self._final_states: Dict[str, str] = {} self._final_states_lock = asyncio.Lock() # ----------------------------------------------------------------------- # Lifecycle Management # ----------------------------------------------------------------------- async def start(self) -> None: if self._running: return self._running = True self._worker_task = asyncio.create_task(self._worker_loop()) async def stop(self) -> None: if not self._running: return self._running = False if self._worker_task: self._worker_task.cancel() try: await self._worker_task except asyncio.CancelledError: pass def on_event(self, callback: Callable[[str, str], Any]) -> None: """Register a callback for job lifecycle events.""" async def _safe_callback(job_id: str, event: str) -> None: try: coro = callback(job_id, event) if asyncio.iscoroutine(coro): await coro except Exception: # Fail silently to prevent callback errors from breaking the pipeline pass async with self._callbacks_lock: self._callbacks.append(_safe_callback) async def _notify(self, job_id: str, event: str) -> None: async with self._callbacks_lock: tasks = [cb(job_id, event) for cb in self._callbacks] if tasks: await asyncio.gather(*tasks, return_exceptions=True) # ----------------------------------------------------------------------- # Public API # ----------------------------------------------------------------------- async def submit(self, text: str, voice: str) -> str: """Enqueue a job. Rejects immediately if backpressure threshold is hit.""" if self._queue.qsize() >= self._max_queue_size: raise ValueError( f"Queue full (backpressure): max {self._max_queue_size} jobs allowed" ) job_id = uuid.uuid4().hex job = Job(id=job_id, text=text, voice=voice) async with self._jobs_lock: self._jobs[job_id] = job await self._queue.put(job) await self._notify(job_id, "queued") return job_id async def drain(self) -> None: """Wait until all queued AND in-flight jobs are finished.""" await self._queue.join() async def cancel(self, job_id: str) -> None: """Cancel a job. Safe for queued or in-flight jobs.""" async with self._jobs_lock: job = self._jobs.get(job_id) if job and job.state in ("queued", "started"): job.state = "cancelled" await self._notify(job_id, "cancelled") # ----------------------------------------------------------------------- # Observability (for tests) # ----------------------------------------------------------------------- def get_max_concurrency(self) -> int: return self._max_concurrency_seen async def get_job_states(self) -> Dict[str, str]: async with self._final_states_lock: return dict(self._final_states) # ----------------------------------------------------------------------- # Internal Worker # ----------------------------------------------------------------------- async def _worker_loop(self) -> None: while self._running: try: job = await self._queue.get() except asyncio.CancelledError: break async with self._semaphore: await self._process_job(job) self._queue.task_done() async def _process_job(self, job: Job) -> None: # Track concurrency async with self._concurrency_lock: self._active_count += 1 if self._active_count > self._max_concurrency_seen: self._max_concurrency_seen = self._active_count try: # Skip if cancelled while in queue if job.state == "cancelled": await self._notify(job.id, "cancelled") async with self._final_states_lock: self._final_states[job.id] = "cancelled" return await self._notify(job.id, "started") max_attempts = self._max_retries + 1 # 1 initial + 3 retries for attempt in range(max_attempts): if job.state == "cancelled": break try: await mock_synthesize(job.text) job.state = "completed" await self._notify(job.id, "completed") async with self._final_states_lock: self._final_states[job.id] = "completed" return except Exception: # Failure handling if attempt < max_attempts - 1: backoff = 0.1 * (2 ** attempt) # 0.1, 0.2, 0.4 await asyncio.sleep(backoff) else: job.state = "failed" await self._notify(job.id, "failed") async with self._final_states_lock: self._final_states[job.id] = "failed" return finally: async with self._concurrency_lock: self._active_count -= 1 # --------------------------------------------------------------------------- # Test Suite # --------------------------------------------------------------------------- async def main() -> None: pipeline = TTSPipeline(max_concurrency=4, max_queue_size=100, max_retries=3) await pipeline.start() # ----------------------------------------------------------------------- # a) Bounded Concurrency Test # ----------------------------------------------------------------------- print("Test A: Bounded concurrency (50 jobs)...") ids_a = [await pipeline.submit(f"text_{i}", "voice_1") for i in range(50)] await pipeline.drain() max_c = pipeline.get_max_concurrency() assert max_c <= 4, f"❌ Concurrency exceeded: {max_c}" print(f" ✅ Max concurrency observed: {max_c}") # ----------------------------------------------------------------------- # b) Backpressure Test # ----------------------------------------------------------------------- print("\nTest B: Backpressure cap (100 jobs)...") ids_b = [await pipeline.submit(f"bp_{i}", "voice_1") for i in range(100)] try: await pipeline.submit("bp_overflow", "voice_1") assert False, "❌ Should have raised ValueError" except ValueError as e: print(f" ✅ Backpressure triggered: {e}") # ----------------------------------------------------------------------- # c) Cancellation & Pool Recovery Test # ----------------------------------------------------------------------- print("\nTest C: Cancellation mid-flight & pool recovery...") ids_c = [await pipeline.submit(f"cancel_{i}", "voice_2") for i in range(10)] # Cancel one immediately after submission await pipeline.cancel(ids_c[3]) await pipeline.drain() states = await pipeline.get_job_states() assert len(states) == 10, f"❌ Job count mismatch: {len(states)}" assert len(set(states.keys())) == 10, "❌ Duplicate job IDs found" terminal_states = {"completed", "failed", "cancelled"} for jid, st in states.items(): assert st in terminal_states, f"❌ Non-terminal state for {jid}: {st}" print(f" ✅ All 10 jobs reached terminal state.") print(f" States: {states}") # Verify pool is still usable after cancellation recover_id = await pipeline.submit("recover_test", "voice_3") await pipeline.drain() final_states = await pipeline.get_job_states() assert final_states.get(recover_id) == "completed", "❌ Pool is broken after cancellation" print(" ✅ Pool remains fully functional after cancellation.") await pipeline.stop() print("\n✅ All tests passed.") if __name__ == "__main__": asyncio.run(main())