diff --git a/.github/workflows/test-realtime-ws.yml b/.github/workflows/test-realtime-ws.yml new file mode 100644 index 000000000..74aa23b3a --- /dev/null +++ b/.github/workflows/test-realtime-ws.yml @@ -0,0 +1,64 @@ +name: Validate realtime WebSocket service + +on: + pull_request: + paths: + - "funasr/bin/realtime_ws.py" + - "funasr/utils/postprocess_hotwords.py" + - "examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py" + - "tests/test_realtime*.py" + - "examples/industrial_data_pretraining/fun_asr_nano/realtime_ws_benchmark.py" + - ".github/workflows/test-realtime-ws.yml" + push: + branches: [main] + paths: + - "funasr/bin/realtime_ws.py" + - "funasr/utils/postprocess_hotwords.py" + - "examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py" + - "tests/test_realtime*.py" + - "examples/industrial_data_pretraining/fun_asr_nano/realtime_ws_benchmark.py" + - ".github/workflows/test-realtime-ws.yml" + +permissions: + contents: read + +jobs: + transport: + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + include: + - python: "3.10" + websockets: "10.4" + - python: "3.12" + websockets: "17.1" + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + CUDA_VISIBLE_DEVICES: "" + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + cache: pip + - name: Install CPU test dependencies + run: | + python -m pip install torch==2.10.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/cpu + python -m pip install -e . pytest einops pytorch-wpe "websockets==${{ matrix.websockets }}" + - name: Check sessions, bounded ingress and real loopback transport + run: | + python -m pytest -q \ + tests/test_realtime_ws_service.py \ + tests/test_realtime_ws_benchmark.py \ + tests/test_realtime_ws_receive_pressure.py \ + tests/test_realtime_receive_configuration.py \ + --junitxml=realtime-ws-results.xml + - uses: actions/upload-artifact@v4 + if: always() + with: + name: realtime-ws-${{ matrix.python }}-${{ matrix.websockets }} + path: realtime-ws-results.xml + if-no-files-found: error diff --git a/docs/benchmark/realtime_ws_benchmark.md b/docs/benchmark/realtime_ws_benchmark.md index 6cf8f57f7..9abc0cf70 100644 --- a/docs/benchmark/realtime_ws_benchmark.md +++ b/docs/benchmark/realtime_ws_benchmark.md @@ -117,6 +117,29 @@ limit. In particular, long speech segments create synchronized, expensive final decodes that are not representative of every meeting or voice-agent workload. +## Receive Pressure Acceptance + +The bounded receive FIFO and deferred provisional decodes must be evaluated +separately from acoustic throughput. Record both receive limits, actual +`websockets` version, frame size and send interval, and whether the load is paced +or a burst. With `--log-decode-profile`, retain the receive high-water observations +(`peak_messages`, `peak_bytes`, `skipped_partials`) and engine profile together. + +Acceptance covers: Ping/Pong during a finite decode stall; explicit close1013 +when either receive limit is exceeded; no silent PCM loss/reordering; ordered +START, hotword/language changes, COMMIT and STOP; full completed-segment inputs; +and reader/worker cleanup after disconnect or cancellation. A disconnected or +overloaded session is a failed session, not a completed final. The queue limits +do not bound the whole process: protocol and session buffers remain additional. + +Fewer partials mean different work and potentially different context/fallback +observations. Compare first text, preview freshness, complete transcripts and +final latency alongside request counts and encoded seconds; a lower wall time +alone is not a win. Keep low-concurrency interactive cases in the matrix, then +run the reporter's L20 paced47-second/16-client workload and another GPU before +claiming a performance fix. Synthetic blocked-decoder tests establish transport +behavior only; they do not demonstrate GPU capacity or resolve issue#3528. + ## Report Template When publishing a realtime WebSocket benchmark or issue report, include: diff --git a/docs/vllm_guide.md b/docs/vllm_guide.md index 7ea949b14..c14bcb845 100644 --- a/docs/vllm_guide.md +++ b/docs/vllm_guide.md @@ -713,11 +713,34 @@ CUDA_VISIBLE_DEVICES=0 python examples/industrial_data_pretraining/fun_asr_nano/ Set a positive `--ws-ping-timeout` only after measuring the worst-case decode and queue delay for the production traffic shape; keep it above that delay and coordinate it with the gateway idle-timeout policy. The `websockets` library's -`max_queue` setting bounds receive buffering for incoming messages; it doesn't -change ping/pong timeout semantics, so increasing it doesn't fix keepalive -timeouts. Set `--ws-ping-interval 0` only when an external gateway already owns +receive high-water mark (`max_queue`) can pause socket reads when its queue fills; +this also delays processing Ping frames even if inference runs off the event loop. +Increasing that mark only moves the overload threshold. It isn't a throughput +fix. Set `--ws-ping-interval 0` only when an external gateway already owns keepalive/reconnect policy. +The source server receives messages independently of session inference, with a +bounded application FIFO: `--ws-receive-max-messages 128` and +`--ws-receive-max-bytes 16777216`. Both limits must be positive. They count queued +messages and payload bytes (UTF-8 bytes for text commands), not total process +memory: protocol buffering, an in-flight message and session audio add to it. +An overflow closes the connection with **1013**, not a successful final result. +Do not automatically replay a partially processed session without an +application-level recovery policy. + +Audio and commands remain ordered. When the next queued message is more audio, +an otherwise-due provisional decode is deferred until that backlog is consumed; +audio itself isn't dropped or combined. VAD completed-segment decoding and +`COMMIT`/`STOP` final decoding still process their full input. Preview cadence and +context/fallback observations may change under load, so identical transcript +text or improved hardware throughput isn't guaranteed. Use +`--log-decode-profile` to record `peak_messages`, `peak_bytes` and +`skipped_partials` alongside the existing engine profile. + +These source changes are not in the published `funasr==1.4.15` package. Transport +tests with synthetic decoding do not replace L20 or production-traffic +acceptance; see [the benchmark contract](benchmark/realtime_ws_benchmark.md). + For long-session debugging, especially with `--enable-spk`, enable periodic session-state logs: diff --git a/docs/vllm_guide_zh.md b/docs/vllm_guide_zh.md index 94faa1ca6..7446c4730 100644 --- a/docs/vllm_guide_zh.md +++ b/docs/vllm_guide_zh.md @@ -683,11 +683,28 @@ CUDA_VISIBLE_DEVICES=0 python examples/industrial_data_pretraining/fun_asr_nano/ 只有在按生产流量测得最坏推理和排队延迟后，才设置正数 `--ws-ping-timeout`；该值应高于实测延迟，并与网关 idle timeout 策略配合。 -`websockets` 库的 `max_queue` 设置只限制入站消息的接收缓冲，不会改变 -ping/pong 的超时语义，因此增大它不能解决 keepalive timeout。只有外部网关 +`websockets` 库的接收高水位 `max_queue` 在队列满时可能暂停 socket 读取； +即使推理已经移出事件循环，后续 Ping 帧也会因此延迟处理。 +增大高水位只是推后过载触发点，不是吞吐修复。只有外部网关 已经统一负责 keepalive / reconnect 策略时，才设置 `--ws-ping-interval 0` 关闭服务端 ping。 +源码服务将消息接收与会话推理解耦，应用层 FIFO 默认限制为 +`--ws-receive-max-messages 128` 和 `--ws-receive-max-bytes 16777216`， +两项都必须为正数。限制分别统计排队消息数和载荷字节数（文本命令按 UTF-8 +字节计），并非进程总内存上限；协议缓冲、正在处理的消息与会话音频还会占用内存。 +超限时连接以 **1013** 关闭，不会作为成功的 final 返回。没有应用层恢复策略时， +不要自动重放可能已经部分处理的会话。 + +音频与命令保持原顺序。当队列中的下一条消息仍是音频时，会推迟一次本已到期的 +临时预览解码，先消费积压音频；不会丢弃或合并音频帧。VAD 完整段解码和 +`COMMIT`/`STOP` 最终解码仍处理完整输入。负载下的预览频率、上下文与回退观测可能 +变化，因此不保证转写文本逐字一致或硬件吞吐提升。使用 `--log-decode-profile` +可同时记录 `peak_messages`、`peak_bytes`、`skipped_partials` 与原有 engine profile。 + +以上源码改动尚未进入已发布的 `funasr==1.4.15`。合成解码器的传输测试不能代替 +L20 或生产流量验收，详见[压测契约](benchmark/realtime_ws_benchmark.md)。 + 长会话排障，尤其是启用 `--enable-spk` 时，可以打开周期性 session 状态日志： ```bash diff --git a/docs/vllm_guide_zh_v2.md b/docs/vllm_guide_zh_v2.md index ada173d35..27b8bd53e 100644 --- a/docs/vllm_guide_zh_v2.md +++ b/docs/vllm_guide_zh_v2.md @@ -686,11 +686,28 @@ CUDA_VISIBLE_DEVICES=0 python examples/industrial_data_pretraining/fun_asr_nano/ 只有在按生产流量测得最坏推理和排队延迟后，才设置正数 `--ws-ping-timeout`；该值应高于实测延迟，并与网关 idle timeout 策略配合。 -`websockets` 库的 `max_queue` 设置只限制入站消息的接收缓冲，不会改变 -ping/pong 的超时语义，因此增大它不能解决 keepalive timeout。只有外部网关 +`websockets` 库的接收高水位 `max_queue` 在队列满时可能暂停 socket 读取； +即使推理已经移出事件循环，后续 Ping 帧也会因此延迟处理。 +增大高水位只是推后过载触发点，不是吞吐修复。只有外部网关 已经统一负责 keepalive / reconnect 策略时，才设置 `--ws-ping-interval 0` 关闭服务端 ping。 +源码服务将消息接收与会话推理解耦，应用层 FIFO 默认限制为 +`--ws-receive-max-messages 128` 和 `--ws-receive-max-bytes 16777216`， +两项都必须为正数。限制分别统计排队消息数和载荷字节数（文本命令按 UTF-8 +字节计），并非进程总内存上限；协议缓冲、正在处理的消息与会话音频还会占用内存。 +超限时连接以 **1013** 关闭，不会作为成功的 final 返回。没有应用层恢复策略时， +不要自动重放可能已经部分处理的会话。 + +音频与命令保持原顺序。当队列中的下一条消息仍是音频时，会推迟一次本已到期的 +临时预览解码，先消费积压音频；不会丢弃或合并音频帧。VAD 完整段解码和 +`COMMIT`/`STOP` 最终解码仍处理完整输入。负载下的预览频率、上下文与回退观测可能 +变化，因此不保证转写文本逐字一致或硬件吞吐提升。使用 `--log-decode-profile` +可同时记录 `peak_messages`、`peak_bytes`、`skipped_partials` 与原有 engine profile。 + +以上源码改动尚未进入已发布的 `funasr==1.4.15`。合成解码器的传输测试不能代替 +L20 或生产流量验收，详见[压测契约](benchmark/realtime_ws_benchmark.md)。 + 长会话排障，尤其是启用 `--enable-spk` 时，可以打开周期性 session 状态日志： ```bash diff --git a/funasr/bin/realtime_ws.py b/funasr/bin/realtime_ws.py index 27f73232d..e9a1f97f5 100644 --- a/funasr/bin/realtime_ws.py +++ b/funasr/bin/realtime_ws.py @@ -1421,9 +1421,125 @@ def create_vad(vad_model, args): return DynamicStreamingVAD(session_model) +class ReceiveBufferOverflow(RuntimeError): + """A connection exceeded its bounded application receive queue.""" + + +class RealtimeReceiveBuffer: + """Receive independently of inference without mutating the ASR session.""" + + def __init__(self, websocket, max_messages=128, max_bytes=16 * 1024 * 1024): + for name, value in (("max_messages", max_messages), ("max_bytes", max_bytes)): + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive integer") + self.websocket = websocket + self.max_messages = max_messages + self.max_bytes = max_bytes + self._messages = deque() + self._bytes = 0 + self._ready = asyncio.Event() + self._done = False + self._error = None + self._reader = None + self.peak_messages = 0 + self.peak_bytes = 0 + + def start(self): + if self._reader is not None or self._done: + raise RuntimeError("receive buffer can only be started once") + self._reader = asyncio.create_task(self._receive(), name="funasr-receive") + + @property + def pending_audio(self): + return bool( + self._messages + and isinstance(self._messages[0][0], bytes) + and self._messages[0][0] + ) + + def _clear(self): + self._messages.clear() + self._bytes = 0 + + def check_connection(self): + """Recheck receive failure after an in-flight session operation finishes.""" + if self._error is not None: + raise self._error + return getattr(self.websocket, "close_code", None) is None + + async def _receive(self): + try: + async for message in self.websocket: + size = len(message) if isinstance(message, bytes) else len(message.encode("utf-8")) + if len(self._messages) >= self.max_messages or self._bytes + size > self.max_bytes: + self._error = ReceiveBufferOverflow("Realtime receive buffer exceeded") + self._clear() + self._ready.set() + await self.websocket.close(code=1013, reason="Realtime receive buffer exceeded") + return + self._messages.append((message, size)) + self._bytes += size + self.peak_messages = max(self.peak_messages, len(self._messages)) + self.peak_bytes = max(self.peak_bytes, self._bytes) + self._ready.set() + if getattr(self.websocket, "close_code", None) is not None: + self._clear() + except Exception as error: + self._error = self._error or error + self._clear() + finally: + self._done = True + self._ready.set() + + def __aiter__(self): + return self + + async def __anext__(self): + while True: + if self._error is not None: + raise self._error + if self._messages: + message, size = self._messages.popleft() + self._bytes -= size + return message + if self._done: + raise StopAsyncIteration + self._ready.clear() + await self._ready.wait() + + async def aclose(self): + try: + if self._reader is not None: + self._reader.cancel() + await asyncio.gather(self._reader, return_exceptions=True) + finally: + # Completed tasks and receive tracebacks can retain their last payload. + self._reader = None + self._error = None + self._clear() + self._done = True + self._ready.set() + + async def run_session_work(_args, operation, *operation_args, **operation_kwargs): """Run one session off-loop; shared ASR calls are serialized by the batcher.""" - return await asyncio.to_thread(operation, *operation_args, **operation_kwargs) + worker = asyncio.create_task( + asyncio.to_thread(operation, *operation_args, **operation_kwargs) + ) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + # Cancelling an asyncio waiter cannot stop a thread mutating the session. + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + except Exception: + break + if not worker.cancelled(): + worker.exception() + raise def log_session_stats(session): @@ -1464,9 +1580,16 @@ async def handle_client(websocket, args): last_decode_time = 0 stats_interval = getattr(args, "log_session_stats_interval", 0.0) last_stats_time = time.time() + receive = RealtimeReceiveBuffer( + websocket, + max_messages=getattr(args, "ws_receive_max_messages", 128), + max_bytes=getattr(args, "ws_receive_max_bytes", 16 * 1024 * 1024), + ) + skipped_partials = 0 + receive.start() try: - async for message in websocket: + async for message in receive: if isinstance(message, str): cmd = message.strip() if cmd.upper() == "START": @@ -1521,6 +1644,8 @@ async def handle_client(websocket, args): else: commit_started = time.perf_counter() result = await run_session_work(args, session.commit) + if not receive.check_connection(): + break await websocket.send(json.dumps(result)) last_decode_time = time.time() elapsed_ms = (time.perf_counter() - commit_started) * 1000 @@ -1543,6 +1668,8 @@ async def handle_client(websocket, args): result = await run_session_work( args, session.decode, is_final=True ) + if not receive.check_connection(): + break await websocket.send(json.dumps(result)) logger.info( "Final: %d sentences", len(result.get("sentences", [])) @@ -1551,6 +1678,8 @@ async def handle_client(websocket, args): await websocket.send(json.dumps({"event": "stopped"})) elif isinstance(message, bytes) and session.is_active: await run_session_work(args, session.add_audio, message) + if not receive.check_connection(): + break now = time.time() if ( stats_interval @@ -1560,20 +1689,41 @@ async def handle_client(websocket, args): log_session_stats(session) last_stats_time = now if now - last_decode_time >= decode_interval and session.should_decode(): + if receive.pending_audio: + skipped_partials += 1 + continue result = await run_session_work(args, session.decode, is_final=False) + if not receive.check_connection(): + break await websocket.send(json.dumps(result)) last_decode_time = now except ConnectionClosed: logger.info("Client disconnected") + except ReceiveBufferOverflow: + logger.warning("Client exceeded the bounded receive queue; connection closed with 1013") except Exception as e: logger.error(f"Error: {e}", exc_info=True) + finally: + await receive.aclose() + if getattr(args, "log_decode_profile", False): + logger.info( + "Realtime receive profile: peak_messages=%d peak_bytes=%d skipped_partials=%d", + receive.peak_messages, receive.peak_bytes, skipped_partials, + ) def _positive_or_none(value): return None if value <= 0 else value +def _positive_int(value): + number = int(value) + if number <= 0: + raise argparse.ArgumentTypeError("must be a positive integer") + return number + + def build_websocket_serve_kwargs(args): return { "max_size": args.ws_max_size, @@ -1701,6 +1851,14 @@ def build_arg_parser(): help="WebSocket close handshake timeout in seconds.") parser.add_argument("--ws-max-size", type=int, default=10 * 1024 * 1024, help="Maximum incoming WebSocket message size in bytes.") + parser.add_argument( + "--ws-receive-max-messages", type=_positive_int, default=128, + help="Maximum queued application messages per connection; overflow closes with 1013.", + ) + parser.add_argument( + "--ws-receive-max-bytes", type=_positive_int, default=16 * 1024 * 1024, + help="Maximum queued application payload bytes per connection; overflow closes with 1013.", + ) parser.add_argument("--log-session-stats-interval", type=float, default=0.0, help="Log bounded long-session state every N seconds; <=0 disables.") return parser diff --git a/tests/test_realtime_receive_configuration.py b/tests/test_realtime_receive_configuration.py new file mode 100644 index 000000000..09686b980 --- /dev/null +++ b/tests/test_realtime_receive_configuration.py @@ -0,0 +1,53 @@ +import asyncio +import threading + +import pytest + +from test_realtime_ws_service import load_service_module + + +def test_receive_buffer_defaults_are_finite(): + args = load_service_module().build_arg_parser().parse_args([]) + assert getattr(args, "ws_receive_max_messages", None) == 128 + assert getattr(args, "ws_receive_max_bytes", None) == 16 * 1024 * 1024 + + +@pytest.mark.parametrize("flag", ["--ws-receive-max-messages", "--ws-receive-max-bytes"]) +@pytest.mark.parametrize("value", ["0", "-1"]) +def test_receive_buffer_limits_reject_nonpositive_values(flag, value): + parser = load_service_module().build_arg_parser() + # A supported positive value must work before testing rejection. + assert vars(parser.parse_args([flag, "1"]))[flag[2:].replace("-", "_")] == 1 + with pytest.raises(SystemExit): + parser.parse_args([flag, value]) + + +def test_cancelling_session_work_waits_for_the_owned_worker(): + module = load_service_module() + entered, release, finished = threading.Event(), threading.Event(), threading.Event() + + def work(): + entered.set() + assert release.wait(3) + finished.set() + + async def exercise(): + task = asyncio.create_task(module.run_session_work(None, work)) + try: + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + task.cancel() + await asyncio.sleep(0.05) + assert not task.done(), "cancellation abandoned a still-mutating worker" + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(exercise()) diff --git a/tests/test_realtime_ws_receive_pressure.py b/tests/test_realtime_ws_receive_pressure.py new file mode 100644 index 000000000..af0d6d535 --- /dev/null +++ b/tests/test_realtime_ws_receive_pressure.py @@ -0,0 +1,683 @@ +"""Bounded receive ownership, with synthetic acoustics and real WebSockets.""" + +import asyncio +import gc +import json +import threading +import types +import weakref + +import pytest + +from test_realtime_ws_service import load_service_module + + +async def until(predicate, timeout=3): + async def wait(): + while not predicate(): + await asyncio.sleep(0.005) + + await asyncio.wait_for(wait(), timeout) + + +class Socket: + remote_address = ("127.0.0.1", 1) + + def __init__(self, messages=(), *, error=None, close_code=None, hold=False): + self.messages = iter(messages) + self.error = error + self.close_code = None + self.end_close_code = close_code + self.hold = hold + self.exhausted = False + self.cancelled = False + self.closed = [] + self.sent = [] + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.messages) + except StopIteration: + self.exhausted = True + self.close_code = self.end_close_code + if self.hold: + try: + await asyncio.Future() + except asyncio.CancelledError: + self.cancelled = True + raise + if self.error: + raise self.error + raise StopAsyncIteration + + async def close(self, code=1000, reason=""): + self.closed.append((code, reason)) + self.close_code = code + + async def send(self, message): + self.sent.append(json.loads(message)) + + +def buffer_class(module): + assert hasattr(module, "RealtimeReceiveBuffer"), "missing bounded receive owner" + assert hasattr(module, "ReceiveBufferOverflow"), "missing explicit overload failure" + return module.RealtimeReceiveBuffer + + +def test_fifo_exact_limits_utf8_and_command_head(): + module = load_service_module() + factory = buffer_class(module) + + async def exercise(): + messages = [b"ab", "中", b"", b"z", "STOP"] + socket = Socket(messages) + buffer = factory(socket, max_messages=5, max_bytes=10) + buffer.start() + try: + await until(lambda: socket.exhausted) + for expected, pending in zip(messages, [True, False, False, True, False]): + assert buffer.pending_audio is pending + assert await asyncio.wait_for(buffer.__anext__(), 1) == expected + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(buffer.__anext__(), 1) + assert socket.closed == [] + finally: + await buffer.aclose() + + asyncio.run(exercise()) + + +@pytest.mark.parametrize( + "messages,limits", + [([b"a", "STOP"], {"max_messages": 1, "max_bytes": 100}), + (["中"], {"max_messages": 8, "max_bytes": 2}), + ([b"abc"], {"max_messages": 8, "max_bytes": 2})], +) +def test_overflow_discards_queue_and_closes_1013(messages, limits): + module = load_service_module() + factory = buffer_class(module) + + async def exercise(): + socket = Socket(messages) + buffer = factory(socket, **limits) + buffer.start() + try: + await until(lambda: socket.closed) + assert socket.closed[0][0] == 1013 + assert 0 < len(socket.closed[0][1].encode("utf-8")) <= 123 + assert buffer.pending_audio is False + with pytest.raises(module.ReceiveBufferOverflow): + await asyncio.wait_for(buffer.__anext__(), 1) + finally: + await buffer.aclose() + + asyncio.run(exercise()) + + +@pytest.mark.parametrize("name", ["max_messages", "max_bytes"]) +@pytest.mark.parametrize("value", [0, -1, True, 1.5, "2", None]) +def test_buffer_rejects_nonpositive_or_noninteger_limits(name, value): + module = load_service_module() + factory = buffer_class(module) + with pytest.raises((ValueError, TypeError)): + factory(Socket(), **{name: value}) + + +@pytest.mark.parametrize("mode", ["peer_close", "receive_error"]) +def test_disconnect_discards_accepted_work(mode): + module = load_service_module() + factory = buffer_class(module) + + async def exercise(): + socket = Socket([b"pcm", "COMMIT"], close_code=1000 if mode == "peer_close" else None, + error=RuntimeError("receive failed") if mode == "receive_error" else None) + buffer = factory(socket) + buffer.start() + try: + await until(lambda: socket.exhausted) + assert buffer.pending_audio is False + try: + value = await asyncio.wait_for(buffer.__anext__(), 1) + except (StopAsyncIteration, RuntimeError): + pass + else: + pytest.fail(f"disconnected input was dispatched: {value!r}") + finally: + await buffer.aclose() + + asyncio.run(exercise()) + + +def test_aclose_joins_reader_and_releases_pending_input(): + module = load_service_module() + factory = buffer_class(module) + + async def exercise(): + socket = Socket([b"pcm"], hold=True) + buffer = factory(socket) + buffer.start() + await until(lambda: socket.exhausted) + await asyncio.wait_for(buffer.aclose(), 1) + assert socket.cancelled + assert buffer.pending_audio is False + await asyncio.wait_for(buffer.aclose(), 1) + + asyncio.run(exercise()) + + +def test_aclose_releases_queued_payload_references(): + module = load_service_module() + factory = buffer_class(module) + + class TextPayload(str): + pass + + async def exercise(): + payload = TextPayload("queued command payload") + reference = weakref.ref(payload) + socket = Socket([payload], hold=True) + del payload + buffer = factory(socket) + buffer.start() + try: + await until(lambda: socket.exhausted) + assert reference() is not None + finally: + await asyncio.wait_for(buffer.aclose(), 1) + # Let completed-task callbacks release their transient references. + await asyncio.sleep(0) + gc.collect() + assert reference() is None, "closed receive owner retained queued payload" + + asyncio.run(exercise()) + + +def test_cli_receive_defaults_do_not_change_transport_defaults(): + module = load_service_module() + args = module.build_arg_parser().parse_args([]) + assert getattr(args, "ws_receive_max_messages", None) == 128 + assert getattr(args, "ws_receive_max_bytes", None) == 16777216 + assert args.partial_window_sec == 8.0 + kwargs = module.build_websocket_serve_kwargs(args) + assert "max_queue" not in kwargs + assert kwargs["ping_interval"] == 20.0 + assert kwargs["ping_timeout"] is None + + +@pytest.mark.parametrize("flag", ["--ws-receive-max-messages", "--ws-receive-max-bytes"]) +def test_cli_accepts_positive_limits_and_rejects_zero(flag): + parser = load_service_module().build_arg_parser() + args = parser.parse_args([flag, "7"]) + assert getattr(args, flag[2:].replace("-", "_")) == 7 + for value in ["0", "-1"]: + with pytest.raises(SystemExit): + parser.parse_args([flag, value]) + + +def test_cancellation_does_not_return_before_worker_stops_mutating(): + module = load_service_module() + entered, release, finished = threading.Event(), threading.Event(), threading.Event() + + def work(): + entered.set() + assert release.wait(5), "test failed to release owned worker" + finished.set() + + async def exercise(): + task = asyncio.create_task(module.run_session_work(None, work)) + try: + await until(entered.is_set) + task.cancel() + await asyncio.sleep(0.05) + assert not task.done(), "cancellation returned while thread could still mutate session" + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + await until(finished.is_set) + + asyncio.run(exercise()) + + +@pytest.mark.parametrize("worker_error", [False, True]) +def test_repeated_cancellation_joins_worker_even_when_worker_fails(worker_error): + module = load_service_module() + entered, release, finished = threading.Event(), threading.Event(), threading.Event() + + def work(): + entered.set() + try: + assert release.wait(5) + if worker_error: + raise RuntimeError("synthetic worker failure during cancellation") + finally: + finished.set() + + async def exercise(): + errors = [] + loop = asyncio.get_running_loop() + previous = loop.get_exception_handler() + loop.set_exception_handler(lambda loop, context: errors.append(context)) + task = asyncio.create_task(module.run_session_work(None, work)) + try: + await until(entered.is_set) + for _ in range(2): + task.cancel() + await asyncio.sleep(0.02) + assert not task.done(), "a repeated cancellation abandoned the owned thread" + finally: + release.set() + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + await until(finished.is_set) + del task + gc.collect() + await asyncio.sleep(0) + assert errors == [], "worker exception was not retrieved" + finally: + loop.set_exception_handler(previous) + + asyncio.run(exercise()) + + +def test_dequeue_returns_message_and_byte_credits(): + module = load_service_module() + factory = buffer_class(module) + + async def exercise(): + class Controlled(Socket): + def __init__(self): + super().__init__() + self.incoming = asyncio.Queue() + self.requests = 0 + + async def __anext__(self): + self.requests += 1 + return await self.incoming.get() + + socket = Controlled() + buffer = factory(socket, max_messages=1, max_bytes=3) + buffer.start() + try: + for number, message in enumerate(["中", b"abc", "xyz"], start=1): + socket.incoming.put_nowait(message) + await until(lambda: socket.requests >= number + 1 or socket.closed) + assert not socket.closed + assert await asyncio.wait_for(buffer.__anext__(), 1) == message + finally: + await buffer.aclose() + + asyncio.run(exercise()) + + +def session_fixture(monkeypatch, module, *, block=None, entered=None, release=None): + sessions = [] + + class Session: + def __init__(self, engine, config, vad, **kwargs): + self.is_active = False + self.asr_kwargs = dict(config) + self.total_samples = 0 + self.audio_buffer_start_sample = 0 + self.audio = [] + self.completed = [] + self.config_at_audio = [] + self.previews = [] + self.finals = [] + sessions.append(self) + + def reset(self): + self.total_samples = 0 + self.audio_buffer_start_sample = 0 + + def wait_if_needed(self, operation): + if operation == block and self.total_samples == 1: + entered.set() + assert release.wait(8), "owned synthetic decoder was not released" + + def add_audio(self, message): + self.audio.append(message) + self.total_samples += 1 + self.config_at_audio.append(dict(self.asr_kwargs)) + self.wait_if_needed("add_audio") + # This boundary stands in for completed-segment work inside add_audio. + self.completed.append(message) + + def should_decode(self): + return True + + def decode(self, is_final=False): + self.wait_if_needed("decode") + (self.finals if is_final else self.previews).append(self.total_samples) + return {"event": "final" if is_final else "partial", + "count": self.total_samples, "sentences": []} + + def commit(self): + self.finals.append(self.total_samples) + self.audio_buffer_start_sample = self.total_samples + return {"event": "final", "count": self.total_samples, "sentences": []} + + monkeypatch.setattr(module, "load_models", lambda args: (object(), {}, object(), None)) + monkeypatch.setattr(module, "create_vad", lambda *args: object()) + monkeypatch.setattr(module, "create_speaker_tracker", lambda *args: None) + monkeypatch.setattr(module, "RealtimeASRSession", Session) + return sessions + + +def handler_args(**overrides): + values = dict(device="cpu", decode_interval=0.0, partial_window_sec=8.0, + endpoint_mode="client", ws_receive_max_messages=128, + ws_receive_max_bytes=16777216) + values.update(overrides) + return types.SimpleNamespace(**values) + + +def test_handler_defers_only_obsolete_previews_preserves_commands_and_finals(monkeypatch): + module = load_service_module() + sessions = session_fixture(monkeypatch, module) + socket = Socket(["START", b"a", b"b", "HOTWORDS:alpha,beta", "LANGUAGE:en", + "COMMIT", b"c", "STOP"]) + asyncio.run(asyncio.wait_for(module.handle_client(socket, handler_args()), 3)) + session = sessions[0] + assert session.audio == [b"a", b"b", b"c"] + assert session.completed == session.audio + assert session.config_at_audio == [{}, {}, {"hotwords": ["alpha", "beta"], "language": "en"}] + assert session.previews == [2, 3], "skip stale audio preview, not the preview before a command" + assert session.finals == [2, 3] + assert [item["event"] for item in socket.sent] == [ + "started", "partial", "hotwords_set", "language_set", "final", "partial", "final", "stopped"] + + +def test_no_backlog_keeps_each_due_preview(monkeypatch): + module = load_service_module() + sessions = session_fixture(monkeypatch, module) + + class Paced(Socket): + async def __anext__(self): + message = await super().__anext__() + required = {b"b": 1, "STOP": 2}.get(message, 0) + if required: + await until(lambda: sum(x["event"] == "partial" for x in self.sent) >= required) + return message + + socket = Paced(["START", b"a", b"b", "STOP"]) + asyncio.run(asyncio.wait_for(module.handle_client(socket, handler_args()), 4)) + assert sessions[0].previews == [1, 2] + assert sessions[0].finals == [2] + assert socket.sent[-1] == {"event": "stopped"} + + +def test_handler_failure_closes_owned_reader(monkeypatch): + module = load_service_module() + session_fixture(monkeypatch, module) + + class SendFailure(Socket): + async def send(self, message): + await until(lambda: self.exhausted) + raise RuntimeError("synthetic send failure") + + socket = SendFailure(["START"], hold=True) + asyncio.run(asyncio.wait_for(module.handle_client(socket, handler_args()), 4)) + assert socket.cancelled, "consumer failure leaked its receive task" + + +def test_handler_cancellation_waits_worker_and_joins_receiver(monkeypatch): + module = load_service_module() + entered, release = threading.Event(), threading.Event() + sessions = session_fixture(monkeypatch, module, block="add_audio", entered=entered, release=release) + + async def exercise(): + socket = Socket(["START", b"a"], hold=True) + task = asyncio.create_task(module.handle_client(socket, handler_args())) + try: + await until(entered.is_set) + task.cancel() + await asyncio.sleep(0.05) + assert not task.done(), "handler returned with an active session-mutating worker" + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + assert sessions[0].completed == [b"a"] + assert socket.cancelled + + asyncio.run(exercise()) + + +def test_handler_overflow_does_not_dispatch_queued_stop_or_success(monkeypatch): + module = load_service_module() + entered, release = threading.Event(), threading.Event() + sessions = session_fixture(monkeypatch, module, block="add_audio", entered=entered, release=release) + + class Burst(Socket): + async def __anext__(self): + message = await super().__anext__() + if message == b"b": + await until(entered.is_set) + return message + + async def exercise(): + socket = Burst(["START", b"a", b"b", b"c", "STOP"]) + task = asyncio.create_task(module.handle_client(socket, handler_args(ws_receive_max_messages=2))) + try: + await until(entered.is_set) + await until(lambda: socket.closed, 1) + assert socket.closed[0][0] == 1013 + finally: + release.set() + await asyncio.wait_for(task, 3) + assert sessions[0].audio == [b"a"] + assert sessions[0].completed == [b"a"] + assert sessions[0].previews == [] + assert sessions[0].finals == [] + assert socket.sent == [{"event": "started"}] + + asyncio.run(exercise()) + + +@pytest.mark.parametrize("operation", ["partial", "commit", "stop-client", "stop-server"]) +@pytest.mark.parametrize("failure", ["overflow", "peer-close", "receive-error"]) +def test_inflight_result_is_not_success_after_receive_failure(monkeypatch, operation, failure): + module = load_service_module() + sessions = session_fixture(monkeypatch, module) + entered, release = threading.Event(), threading.Event() + session_type = module.RealtimeASRSession + method = "commit" if operation in {"commit", "stop-client"} else "decode" + original = getattr(session_type, method) + + def blocked(self, *args, **kwargs): + entered.set() + assert release.wait(8), "owned final/preview worker was not released" + return original(self, *args, **kwargs) + + monkeypatch.setattr(session_type, method, blocked) + command = [] if operation == "partial" else ["COMMIT" if operation == "commit" else "STOP"] + + class FailureAfterWork(Socket): + failed = False + + async def __anext__(self): + message = await super().__anext__() + if message == "ABORT-BARRIER": + await until(entered.is_set) + if failure == "peer-close": + self.close_code = 1000 + self.failed = True + raise StopAsyncIteration + if failure == "receive-error": + self.failed = True + raise RuntimeError("synthetic receive error during worker") + return b"queued" + return message + + async def exercise(): + socket = FailureAfterWork(["START", b"a", *command, "ABORT-BARRIER", + *([b"queued"] * 8), "STOP"]) + args = handler_args(ws_receive_max_messages=8, + decode_interval=0 if operation == "partial" else 1e30, + endpoint_mode="server" if operation == "stop-server" else "client") + task = asyncio.create_task(module.handle_client(socket, args)) + try: + await until(entered.is_set) + await until(lambda: socket.failed or socket.closed) + if failure == "overflow": + assert socket.closed[0][0] == 1013 + finally: + release.set() + await asyncio.wait_for(task, 3) + assert sessions[0].audio == [b"a"] + assert sessions[0].completed == [b"a"] + assert socket.sent == [{"event": "started"}], "failed session emitted successful result/STOP" + if operation == "partial": + assert sessions[0].previews == [1] + assert sessions[0].finals == [] + else: + assert sessions[0].previews == [] + assert sessions[0].finals == [1] + + asyncio.run(exercise()) + + +async def receive_until_stopped(client): + values = [] + while True: + value = json.loads(await asyncio.wait_for(client.recv(), 3)) + values.append(value) + if value.get("event") == "stopped": + return values + + +@pytest.mark.parametrize("overflow", [False, True]) +def test_real_websocket_burst_pong_or_explicit_overload(monkeypatch, overflow): + import websockets + from websockets.exceptions import ConnectionClosed + + module = load_service_module() + module.ConnectionClosed = ConnectionClosed + entered, release = threading.Event(), threading.Event() + sessions = session_fixture(monkeypatch, module, block="decode", entered=entered, release=release) + args = handler_args(ws_receive_max_messages=8 if overflow else 128) + + async def exercise(): + handler_tasks = [] + + async def handler(socket, path=None): + task = asyncio.current_task() + handler_tasks.append(task) + await module.handle_client(socket, args) + + async with websockets.serve(handler, "127.0.0.1", 0, ping_interval=None, + close_timeout=0.2) as server: + port = server.sockets[0].getsockname()[1] + try: + async with websockets.connect(f"ws://127.0.0.1:{port}", ping_interval=None, + close_timeout=0.2) as client: + pong = None + try: + await client.send("START") + assert json.loads(await asyncio.wait_for(client.recv(), 2)) == {"event": "started"} + first = b"first" + bytes(6395) + await client.send(first) + await until(entered.is_set) + frames = [i.to_bytes(2, "little") + bytes(6398) for i in range(48)] + for frame in frames: + try: + await client.send(frame) + except ConnectionClosed: + assert overflow + break + await asyncio.sleep(0.003) + if overflow: + with pytest.raises(ConnectionClosed) as caught: + await asyncio.wait_for(client.recv(), 1) + assert caught.value.rcvd is not None + assert caught.value.rcvd.code == 1013 + release.set() + await until(lambda: all(t.done() for t in handler_tasks)) + assert sessions[0].audio == [first] + assert sessions[0].finals == [] + else: + pong = await client.ping(b"receive-pressure-test") + await asyncio.wait_for(asyncio.shield(pong), 0.5) + assert not release.is_set(), "Pong must arrive while decoder is blocked" + await client.send("STOP") + release.set() + received = await receive_until_stopped(client) + assert sessions[0].audio == [first, *frames] + assert sessions[0].completed == sessions[0].audio + assert sessions[0].finals == [49] + assert [r["count"] for r in received if r.get("event") == "final"] == [49] + finally: + release.set() + if pong is not None: + await asyncio.wait_for(asyncio.gather(pong, return_exceptions=True), 2) + finally: + release.set() + if handler_tasks: + await asyncio.wait_for(asyncio.gather(*handler_tasks, return_exceptions=True), 4) + + asyncio.run(exercise()) + + +def test_real_connections_do_not_share_queue_or_block_each_other(monkeypatch): + import websockets + from websockets.exceptions import ConnectionClosed + + module = load_service_module() + module.ConnectionClosed = ConnectionClosed + sessions = session_fixture(monkeypatch, module) + entered, release = threading.Event(), threading.Event() + original = module.RealtimeASRSession.decode + + def decode(self, is_final=False): + if self.audio == [b"slow"]: + entered.set() + assert release.wait(8) + return original(self, is_final) + + monkeypatch.setattr(module.RealtimeASRSession, "decode", decode) + + async def exercise(): + handlers = [] + + async def handler(socket, path=None): + handlers.append(asyncio.current_task()) + await module.handle_client(socket, handler_args()) + + async with websockets.serve(handler, "127.0.0.1", 0, ping_interval=None, + close_timeout=0.2) as server: + uri = f"ws://127.0.0.1:{server.sockets[0].getsockname()[1]}" + try: + async with websockets.connect(uri, ping_interval=None, close_timeout=0.2) as slow: + try: + await slow.send("START") + await asyncio.wait_for(slow.recv(), 2) + await slow.send(b"slow") + await until(entered.is_set) + async with websockets.connect(uri, ping_interval=None, close_timeout=0.2) as fast: + await fast.send("START") + await asyncio.wait_for(fast.recv(), 2) + await fast.send(b"fast-a") + await fast.send(b"fast-b") + await fast.send("STOP") + await receive_until_stopped(fast) + assert not release.is_set() + assert sessions[1].audio == [b"fast-a", b"fast-b"] + assert sessions[1].finals == [2] + release.set() + await slow.send("STOP") + await receive_until_stopped(slow) + assert sessions[0].audio == [b"slow"] + assert sessions[0].finals == [1] + finally: + release.set() + finally: + release.set() + if handlers: + await asyncio.wait_for(asyncio.gather(*handlers, return_exceptions=True), 4) + + asyncio.run(exercise())