242 lines
10 KiB
Python
242 lines
10 KiB
Python
"""Shared-session routing and backpressure, exercised through real OS pipes."""
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import queue
|
||
import socket
|
||
import threading
|
||
from contextlib import ExitStack, suppress
|
||
|
||
import pytest
|
||
|
||
from tui_gateway import server
|
||
from tui_gateway.transport import FanoutTransport, StdioTransport
|
||
from tui_gateway.ws import WSTransport
|
||
|
||
|
||
class PipeClient:
|
||
def __init__(self, stack, *, reading=True):
|
||
read_fd, write_fd = os.pipe()
|
||
self.reader = stack.enter_context(os.fdopen(read_fd, "r", encoding="utf-8"))
|
||
self.writer = os.fdopen(write_fd, "w", encoding="utf-8")
|
||
stack.callback(self._cleanup)
|
||
self.transport = StdioTransport(lambda: self.writer, threading.Lock())
|
||
self._closed = False
|
||
self.writes = 0
|
||
self.frames = queue.Queue()
|
||
if reading:
|
||
self.thread = threading.Thread(target=self._read, daemon=True)
|
||
self.thread.start()
|
||
|
||
def _cleanup(self):
|
||
with suppress(BrokenPipeError):
|
||
self.writer.close()
|
||
if hasattr(self, "thread"):
|
||
self.thread.join(timeout=5)
|
||
assert not self.thread.is_alive()
|
||
|
||
def _read(self):
|
||
for line in self.reader:
|
||
self.frames.put(json.loads(line))
|
||
|
||
def write(self, obj):
|
||
self.writes += 1
|
||
return self.transport.write(obj)
|
||
|
||
def close(self):
|
||
self._closed = True
|
||
|
||
def receive(self):
|
||
return self.frames.get(timeout=5)
|
||
|
||
|
||
class SocketClient(WSTransport):
|
||
"""Real WSTransport with its ASGI send backed by a kernel socketpair."""
|
||
def __init__(self, stack, *, reading=True):
|
||
self.reader, self.writer = socket.socketpair()
|
||
self.writer.setblocking(False)
|
||
loop = asyncio.new_event_loop()
|
||
super().__init__(self, loop)
|
||
self.writes = 0
|
||
self.loop_thread = threading.Thread(target=loop.run_forever, daemon=True)
|
||
self.loop_thread.start()
|
||
self.frames = queue.Queue()
|
||
self.read_thread = None
|
||
if reading:
|
||
self.read_thread = threading.Thread(target=self._read, daemon=True)
|
||
self.read_thread.start()
|
||
stack.callback(self._cleanup)
|
||
|
||
async def send_text(self, payload):
|
||
self.writes += 1
|
||
await self._loop.sock_sendall(self.writer, (payload + "\n").encode())
|
||
|
||
def _read(self):
|
||
with self.reader.makefile("r", encoding="utf-8") as stream:
|
||
for line in stream:
|
||
self.frames.put(json.loads(line))
|
||
|
||
def receive(self):
|
||
return self.frames.get(timeout=5)
|
||
|
||
def _cleanup(self):
|
||
async def cancel_sends():
|
||
self.close()
|
||
tasks = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
|
||
for task in tasks:
|
||
task.cancel()
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
asyncio.run_coroutine_threadsafe(cancel_sends(), self._loop).result(timeout=5)
|
||
self.writer.close()
|
||
if self.read_thread:
|
||
self.read_thread.join(timeout=5)
|
||
assert not self.read_thread.is_alive()
|
||
self.reader.close()
|
||
self._loop.call_soon_threadsafe(self._loop.stop)
|
||
self.loop_thread.join(timeout=5)
|
||
self._loop.close()
|
||
|
||
|
||
def _session(transport):
|
||
return dict(transport=transport, agent=None, session_key="fanout-invariant",
|
||
history=[], history_lock=threading.Lock(), history_version=0,
|
||
running=False, attached_images=[])
|
||
|
||
|
||
@pytest.mark.parametrize("attachment", ["direct", "flattened"])
|
||
def test_membership_preserves_terminal_delivery_and_revokes_departed_peers(monkeypatch, attachment):
|
||
with ExitStack() as stack:
|
||
a, b, stranger = [PipeClient(stack) for _ in range(3)]
|
||
session = _session(a)
|
||
monkeypatch.setitem(server._sessions, "fanout-invariant", session)
|
||
newcomer = b if attachment == "direct" else FanoutTransport(a, b)
|
||
assert server._attach_session_transport(session, newcomer)
|
||
assert server._attach_session_transport(session, b)
|
||
for kind in ("message.delta", "reasoning.delta", "message.complete"):
|
||
server._emit(kind, "fanout-invariant", {"text": "α"})
|
||
first, second = a.receive(), b.receive()
|
||
assert first == second
|
||
assert first["params"]["type"] == kind
|
||
for client, allowed in ((a, True), (b, True), (stranger, False)):
|
||
token = server.bind_transport(client)
|
||
try:
|
||
assert (server._current_session_steer_authority("fanout-invariant")[0] is client) == allowed
|
||
finally:
|
||
server.reset_transport(token)
|
||
# RPC replies stay on their request transport, never the subscriber set.
|
||
token = server.bind_transport(a)
|
||
try:
|
||
assert server.write_json({"jsonrpc": "2.0", "id": "private", "result": "owner only"})
|
||
finally:
|
||
server.reset_transport(token)
|
||
assert a.receive()["id"] == "private"
|
||
assert b.frames.empty()
|
||
from gateway import browser_control_broker as broker_module
|
||
monkeypatch.setattr(broker_module, "browser_control_enabled", lambda: True)
|
||
for client in (a, b, stranger):
|
||
client.auth_identity = {"user_id": "fanout-owner", "provider": "fixture"}
|
||
session["profile"] = "default"
|
||
|
||
def controller(client, action, **params):
|
||
return server.dispatch({"jsonrpc": "2.0", "id": 1,
|
||
"method": "browser.controller." + action,
|
||
"params": {"session_id": "fanout-invariant", **params}}, client)
|
||
|
||
registered = controller(a, "register", controller_id="invariant",
|
||
browser_profile_id="fixture", capabilities=["controller.noop"],
|
||
protocol_version=broker_module.BROWSER_CONTROL_PROTOCOL_VERSION)
|
||
assert "result" in registered, registered
|
||
try:
|
||
assert controller(a, "heartbeat")["result"] == {"ok": True}
|
||
assert controller(b, "heartbeat")["error"]["message"] == "controller is not owned by this transport"
|
||
assert controller(stranger, "heartbeat")["error"]["message"] == "session is not owned by this transport"
|
||
finally:
|
||
controller(a, "detach")
|
||
session["viewers"] = {b: object(), stranger: object()}
|
||
assert server._detach_session_transport(session, b)
|
||
assert not session["viewers"]
|
||
assert not server._session_transport_contains(session, b)
|
||
server._emit("message.complete", "fanout-invariant", {"text": "only A"})
|
||
assert a.receive()["params"]["payload"]["text"] == "only A"
|
||
assert b.frames.empty()
|
||
assert server._attach_session_transport(session, b)
|
||
server._emit("message.complete", "fanout-invariant", {"text": "reattached"})
|
||
assert a.receive() == b.receive()
|
||
b.close()
|
||
# A stale queued envelope must not restore a dead peer's authority.
|
||
assert not server._attach_session_transport(session, b)
|
||
assert not server._session_transport_contains(session, b)
|
||
assert not server._attach_session_transport(session, server._stdio_transport)
|
||
assert server._close_sessions_for_transport(b) == (0, 0)
|
||
assert server._close_sessions_for_transport(a) == (0, 1)
|
||
assert session["transport"] is server._detached_ws_transport
|
||
|
||
|
||
@pytest.mark.linux_only
|
||
@pytest.mark.parametrize("client_type", [PipeClient, SocketClient])
|
||
@pytest.mark.parametrize("slow_first", [True, False])
|
||
@pytest.mark.parametrize("on_loop", [True, False])
|
||
def test_backpressure_never_blocks_later_frames_or_other_subscribers(slow_first, on_loop, client_type, monkeypatch):
|
||
monkeypatch.setattr("tui_gateway.ws._WS_WRITE_TIMEOUT_S", 0.01)
|
||
monkeypatch.setattr("tui_gateway.ws._TOKEN_COALESCE_S", 0)
|
||
with ExitStack() as stack:
|
||
healthy, slow = client_type(stack), client_type(stack, reading=False)
|
||
# Fill the actual kernel pipe, not a fake wait in a transport.write().
|
||
fd = slow.writer.fileno()
|
||
os.set_blocking(fd, False)
|
||
try:
|
||
while True:
|
||
os.write(fd, b"x" * 4096)
|
||
except BlockingIOError:
|
||
pass
|
||
finally:
|
||
os.set_blocking(fd, client_type is PipeClient)
|
||
fan = FanoutTransport(*((slow, healthy) if slow_first else (healthy, slow)))
|
||
returned = threading.Event()
|
||
errors = []
|
||
|
||
def emit():
|
||
try:
|
||
for kind in ("message.start", "message.delta", "message.complete"):
|
||
assert fan.write({"params": {"type": kind}})
|
||
except BaseException as exc:
|
||
errors.append(exc)
|
||
finally:
|
||
returned.set()
|
||
|
||
async def loop_emit():
|
||
emit()
|
||
|
||
worker = threading.Thread(target=(lambda: asyncio.run(loop_emit())) if on_loop else emit, daemon=True)
|
||
worker.start()
|
||
try:
|
||
assert returned.wait(3), "slow subscriber blocked the emitting turn"
|
||
assert not errors
|
||
assert [healthy.receive()["params"]["type"] for _ in range(3)] == [
|
||
"message.start", "message.delta", "message.complete"]
|
||
# Exhaust only the slow peer's bounded backlog; pace the healthy
|
||
# reader by receipts so scheduler latency cannot make it overflow.
|
||
for n in range(1024):
|
||
frame = {"params": {"type": "message.delta", "n": n}}
|
||
assert fan.write(frame)
|
||
assert healthy.receive() == frame
|
||
if not fan.contains(slow):
|
||
break
|
||
assert not fan.contains(slow), "slow backlog grew without bound"
|
||
for _ in range(16):
|
||
assert fan.attach(slow)
|
||
assert fan.write({"reattach": True})
|
||
assert healthy.receive() == {"reattach": True}
|
||
fan.detach(slow)
|
||
assert slow.writes == 1, "reattach spawned more writers behind blocked I/O"
|
||
assert fan.contains(healthy)
|
||
assert fan.write({"params": {"type": "message.complete"}})
|
||
assert healthy.receive()["params"]["type"] == "message.complete"
|
||
finally:
|
||
# Closing the real reader releases any blocked writer even on RED.
|
||
slow.reader.close()
|
||
worker.join(timeout=15)
|
||
fan.close()
|
||
assert not worker.is_alive()
|
||
assert not fan.write({"after": "close"})
|