# tts_pipeline.py import asyncio import random import uuid from collections import deque # ---------- mock synthesizer ---------- _global_active = 0 _global_max = 0 _global_active_lock = asyncio.Lock() async def mock_synthesize(text: str) -> bytes: global _global_active, _global_max async with _global_active_lock: _global_active += 1 if _global_active > _global_max: _global_max = _global_active try: await asyncio.sleep(random.uniform(0.05, 0.30)) if random.random() < 0.10: raise RuntimeError("synthesis failed") return b"\x00" * len(text) finally: async with _global_active_lock: _global_active -= 1 # ---------- pipeline ---------- class TTSPipeline: def __init__(self): self._queue = deque() self._queue_lock = asyncio.Lock() self._queue_not_empty = asyncio.Condition(self._queue_lock) self._jobs_lock = asyncio.Lock() self._jobs = {} # id -> meta self._active_jobs = set() self._callbacks = [] self._workers = [] def register_callback(self, cb): self._callbacks.append(cb) async def _emit(self, event, job_id, **kw): for cb in self._callbacks: try: cb(event, job_id, **kw) except Exception: pass async def start(self): for _ in range(4): self._workers.append(asyncio.create_task(self._worker())) async def _worker(self): while True: async with self._queue_lock: while not self._queue: await self._queue_not_empty.wait() job = self._queue.popleft() job_id = job["id"] async with self._jobs_lock: meta = self._jobs.get(job_id) if meta and meta.get("cancelled"): await self._emit("cancelled", job_id) continue self._active_jobs.add(job_id) try: await self._process_job(job) finally: async with self._jobs_lock: self._active_jobs.discard(job_id) async def _process_job(self, job): job_id = job["id"] async with self._jobs_lock: if job_id in self._jobs: self._jobs[job_id]["status"] = "started" await self._emit("started", job_id) backoff = 0.1 for attempt in range(1, 4): async with self._jobs_lock: meta = self._jobs.get(job_id) if not meta or meta.get("cancelled"): await self._emit("cancelled", job_id) return try: await mock_synthesize(job["text"]) await self._emit("completed", job_id) async with self._jobs_lock: if job_id in self._jobs: self._jobs[job_id]["status"] = "completed" return except asyncio.CancelledError: await self._emit("cancelled", job_id) raise except Exception as e: async with self._jobs_lock: if job_id in self._jobs: self._jobs[job_id]["attempts"] = attempt if attempt >= 3: await self._emit("failed", job_id, error=e) async with self._jobs_lock: if job_id in self._jobs: self._jobs[job_id]["status"] = "failed" return await asyncio.sleep(backoff) backoff *= 2 async def submit(self, text: str, voice: str) -> str: job_id = uuid.uuid4().hex async with self._jobs_lock: self._jobs[job_id] = { "text": text, "voice": voice, "cancelled": False, "status": "queued", "attempts": 0 } async with self._queue_lock: if len(self._queue) >= 100: async with self._jobs_lock: self._jobs.pop(job_id, None) raise RuntimeError("Backpressure: queue full") self._queue.append({"id": job_id, "text": text, "voice": voice}) self._queue_not_empty.notify() await self._emit("queued", job_id) return job_id async def cancel(self, job_id: str) -> bool: async with self._jobs_lock: meta = self._jobs.get(job_id) if not meta: return False if meta["status"] in ("completed", "failed", "cancelled"): return False meta["cancelled"] = True removed = False async with self._queue_lock: for i, j in enumerate(self._queue): if j["id"] == job_id: del self._queue[i] removed = True break if removed: await self._emit("cancelled", job_id) return True async def drain(self): while True: async with self._queue_lock: q_empty = len(self._queue) == 0 async with self._jobs_lock: active_empty = len(self._active_jobs) == 0 if q_empty and active_empty: break await asyncio.sleep(0.01) # ---------- tests ---------- async def main(): global _global_active, _global_max _global_active = 0 _global_max = 0 # a) concurrency limit pipeline = TTSPipeline() await pipeline.start() ids = [await pipeline.submit(f"text {i}", "v1") for i in range(50)] await pipeline.drain() print("max concurrency", _global_max) assert _global_max <= 4, f"max concurrency {_global_max} > 4" # b) backpressure _global_active = 0; _global_max = 0 pipeline2 = TTSPipeline() await pipeline2.start() for i in range(100): await pipeline2.submit(f"t{i}", "v") try: await pipeline2.submit("overflow", "v") assert False, "should have raised" except RuntimeError as e: print("backpressure works:", e) await pipeline2.drain() # c) cancel mid-flight and reuse _global_active = 0; _global_max = 0 pipeline3 = TTSPipeline() await pipeline3.start() events3 = [] pipeline3.register_callback(lambda e, jid, **kw: events3.append((e, jid))) ids3 = [await pipeline3.submit(f"t{i}", "v") for i in range(20)] await asyncio.sleep(0.05) cancel_id = ids3[5] await pipeline3.cancel(cancel_id) await pipeline3.drain() queued = [e for e in events3 if e[0] == "queued"] completed = [e for e in events3 if e[0] == "completed"] failed = [e for e in events3 if e[0] == "failed"] cancelled = [e for e in events3 if e[0] == "cancelled"] print("queued", len(queued), "finished", len(completed)+len(failed)+len(cancelled)) assert len(queued) == 20 assert len(completed) + len(failed) + len(cancelled) == 20 assert any(e[1] == cancel_id and e[0] == "cancelled" for e in events3) # pool still usable new_id = await pipeline3.submit("after cancel", "v") await pipeline3.drain() print("pool still usable") print("All tests passed") if __name__ == "__main__": asyncio.run(main())