diff --git a/cmd_chat/server/helpers.py b/cmd_chat/server/helpers.py index 46394d1..c2816f9 100644 --- a/cmd_chat/server/helpers.py +++ b/cmd_chat/server/helpers.py @@ -24,19 +24,20 @@ def get_client_ip(request: Request) -> str: return request.ip -async def send_state(ws: Websocket, app: Sanic) -> None: +def state_frame(app: Sanic) -> str: + """Build the per-client `init` snapshot (message backlog + roster) as a JSON + string. Returned (not sent) so it can be enqueued as the connection's first + outbound frame, guaranteeing it precedes any later broadcast.""" messages = app.ctx.message_store.get_all() users = app.ctx.session_store.get_all() - await ws.send( - json.dumps( - { - "type": "init", - "messages": [asdict(m) for m in messages], - "users": [ - {"user_id": u.user_id, "username": u.username} for u in users - ], - } - ) + return json.dumps( + { + "type": "init", + "messages": [asdict(m) for m in messages], + "users": [ + {"user_id": u.user_id, "username": u.username} for u in users + ], + } ) diff --git a/cmd_chat/server/managers.py b/cmd_chat/server/managers.py index 089d40d..fd92bd1 100644 --- a/cmd_chat/server/managers.py +++ b/cmd_chat/server/managers.py @@ -3,41 +3,117 @@ from typing import Optional from sanic import Websocket +# Per-connection outbound backlog before a client is considered wedged. +# +# Every relayed frame is *enqueued* (never awaited) by the broadcaster, and a +# dedicated per-connection writer task drains the queue with `await ws.send`. +# This decouples the room's read loop from any single slow receiver: the old +# code awaited each `send` serially under a global lock, so the slowest viewer +# (e.g. a remote peer over a laggy Tailscale link) gated delivery to everyone +# *and* stalled the server's read of the next frame — so a sandbox owner's PTY +# stream backed up and then arrived "all at once". +# +# A backlog this deep absorbs a large PTY burst and transient network hiccups. +# A client that falls this far behind is effectively stuck (it would be closed +# by the websocket ping-timeout within ~40s anyway); we close it now so the room +# isn't holding unbounded memory for it. It reconnects (Ctrl-R) and the host +# replays a fresh sandbox snapshot, which is cleaner than feeding it a corrupt, +# partial terminal byte-stream with gaps. +SEND_QUEUE_MAX = 4096 + + +class _Conn: + """A live websocket plus its outbound queue and the task draining it.""" + + __slots__ = ("ws", "queue", "task") + + def __init__(self, ws: Websocket): + self.ws = ws + self.queue: asyncio.Queue[str] = asyncio.Queue(maxsize=SEND_QUEUE_MAX) + self.task: Optional[asyncio.Task] = None + + class ConnectionManager: def __init__(self): - self.active_connections: dict[str, Websocket] = {} + self.active_connections: dict[str, _Conn] = {} self._lock = asyncio.Lock() - async def connect(self, user_id: str, websocket: Websocket) -> None: + async def connect( + self, user_id: str, websocket: Websocket, initial: Optional[str] = None + ) -> None: + """Register a connection and start its writer task. + + `initial` (the per-client init snapshot) is enqueued *before* the + connection becomes visible to broadcasters, guaranteeing it is the first + frame the client receives — no later broadcast can jump ahead of it. + """ + conn = _Conn(websocket) + if initial is not None: + conn.queue.put_nowait(initial) + old = None async with self._lock: - self.active_connections[user_id] = websocket + old = self.active_connections.get(user_id) + self.active_connections[user_id] = conn + if old is not None and old.task is not None: + old.task.cancel() + conn.task = asyncio.create_task(self._writer(conn)) + + async def _writer(self, conn: _Conn) -> None: + """Drain one connection's queue in FIFO order, one `send` at a time.""" + try: + while True: + msg = await conn.queue.get() + await conn.ws.send(msg) + except asyncio.CancelledError: + raise + except Exception: + # Socket died mid-send; the handler's `finally` will disconnect us. + pass async def disconnect(self, user_id: str) -> None: async with self._lock: - if user_id in self.active_connections: - del self.active_connections[user_id] + conn = self.active_connections.pop(user_id, None) + if conn is not None and conn.task is not None: + conn.task.cancel() + + @staticmethod + def _enqueue(conn: _Conn, message: str) -> bool: + """Non-blocking enqueue. False means the backlog is full (wedged peer).""" + try: + conn.queue.put_nowait(message) + return True + except asyncio.QueueFull: + return False + + async def _drop_wedged(self, user_id: str, conn: _Conn) -> None: + """Evict a client whose backlog overflowed and nudge it to reconnect.""" + await self.disconnect(user_id) + try: + # 1013 = "Try Again Later": the client surfaces a closed socket and + # can reconnect to resync from a fresh snapshot. + await conn.ws.close(code=1013) + except Exception: + pass async def broadcast(self, message: str, exclude_user: Optional[str] = None) -> None: async with self._lock: - disconnected = [] - for user_id, connection in list(self.active_connections.items()): - if exclude_user and user_id == exclude_user: - continue - try: - await connection.send(message) - except Exception: - disconnected.append(user_id) - - for user_id in disconnected: - if user_id in self.active_connections: - del self.active_connections[user_id] + targets = [ + (uid, conn) + for uid, conn in self.active_connections.items() + if not (exclude_user and uid == exclude_user) + ] + # Enqueue is synchronous and fast, so the broadcaster returns immediately + # — it never waits on any socket. Only genuinely wedged peers overflow. + wedged = [(uid, conn) for uid, conn in targets if not self._enqueue(conn, message)] + for uid, conn in wedged: + await self._drop_wedged(uid, conn) async def send_personal(self, user_id: str, message: str) -> bool: async with self._lock: - if connection := self.active_connections.get(user_id): - try: - await connection.send(message) - return True - except Exception: - return False + conn = self.active_connections.get(user_id) + if conn is None: + return False + if self._enqueue(conn, message): + return True + await self._drop_wedged(user_id, conn) return False diff --git a/cmd_chat/server/views.py b/cmd_chat/server/views.py index c483d51..16f1bc2 100644 --- a/cmd_chat/server/views.py +++ b/cmd_chat/server/views.py @@ -2,13 +2,27 @@ import hashlib import hmac import json import base64 +import socket from dataclasses import asdict from sanic import Sanic, Request, response, Websocket from sanic.response import HTTPResponse, json as json_response from .models import Message, UserSession -from .helpers import get_client_ip, send_state, utcnow +from .helpers import get_client_ip, state_frame, utcnow + + +def _disable_nagle(ws: Websocket) -> None: + """Set TCP_NODELAY on a websocket's underlying socket. Small interactive + frames (PTY echo, keystrokes, chat) must not wait on Nagle coalescing + + delayed-ACK before reaching a viewer. Best-effort: any transport without a + raw socket (or where the option can't be set) is simply left as-is.""" + try: + sock = ws.io_proto.transport.get_extra_info("socket") + if sock is not None: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + except Exception: + pass # Hard cap on a single relayed WS frame. The largest legitimate frame is one @@ -137,11 +151,12 @@ async def chat_ws(request: Request, ws: Websocket, app: Sanic) -> None: return manager = app.ctx.connection_manager - await manager.connect(user_id, ws) + _disable_nagle(ws) + # Enqueue this client's init snapshot as its first outbound frame, then + # register it so broadcasts can target it — guaranteeing init arrives first. + await manager.connect(user_id, ws, initial=state_frame(app)) try: - await send_state(ws, app) - # Announce arrival to everyone already present, then a fresh roster. await manager.broadcast( json.dumps( diff --git a/hh/src/net.rs b/hh/src/net.rs index ae776df..944eed4 100644 --- a/hh/src/net.rs +++ b/hh/src/net.rs @@ -113,6 +113,15 @@ pub async fn connect(session: &Session) -> Result { let (ws, _) = tokio_tungstenite::connect_async(&session.ws_url) .await .context("websocket connect")?; + // Disable Nagle's algorithm. Sandbox PTY echo, keystrokes, and chat are all + // small frames; Nagle holds a small segment waiting for more data until the + // prior one is ACKed, and paired with delayed-ACK (~40 ms) — amplified by + // Tailscale's RTT — that is a classic cause of "type, pause, then a burst" + // lag for the non-host viewer. The default (and Tailscale) path is plaintext + // ws, a `Plain` TcpStream, so set TCP_NODELAY there. + if let MaybeTlsStream::Plain(s) = ws.get_ref() { + let _ = s.set_nodelay(true); + } Ok(ws) } diff --git a/hh/src/sbx.rs b/hh/src/sbx.rs index 9bb613f..23d8909 100644 --- a/hh/src/sbx.rs +++ b/hh/src/sbx.rs @@ -553,7 +553,10 @@ pub fn vm_save_state(vm: &str, label: &str, local: bool) -> Result { err.lines().last().unwrap_or("").trim() ); } - return Ok(format!("snapshot '{label}' + local appliance {}", path.display())); + return Ok(format!( + "snapshot '{label}' + local appliance {}", + path.display() + )); } Ok(format!("snapshot '{label}' of {vm}")) } @@ -766,7 +769,12 @@ pub fn run_user_for(backend: Backend, owner: &str) -> String { /// `tar` inside the container/VM — uniform for files and directories, and no /// shell interpolation of the path. Blocking — run off the UI thread. Returns /// the in-sandbox destination path on success. -pub fn push(backend: Backend, name: &str, run_user: &str, local: &std::path::Path) -> Result { +pub fn push( + backend: Backend, + name: &str, + run_user: &str, + local: &std::path::Path, +) -> Result { let (base, tar) = crate::ft::tar_path(local)?; match backend { Backend::Local => anyhow::bail!( diff --git a/hh/src/theme.rs b/hh/src/theme.rs index 988d8e1..88a4397 100644 --- a/hh/src/theme.rs +++ b/hh/src/theme.rs @@ -378,9 +378,16 @@ roster_width = 24 let inks = [Color::Rgb(0xff, 0xff, 0xff), Color::Rgb(0xd0, 0xd0, 0xd0)]; let bg = legible_bg(210.0, 0.5, &inks); for ink in inks { - assert!(contrast(ink, bg) >= 4.5, "ink unreadable: {:.2}", contrast(ink, bg)); + assert!( + contrast(ink, bg) >= 4.5, + "ink unreadable: {:.2}", + contrast(ink, bg) + ); } - assert!(luminance(bg) > 0.0, "bright inks should permit a tinted, non-black bg"); + assert!( + luminance(bg) > 0.0, + "bright inks should permit a tinted, non-black bg" + ); // A darker ink set forces a correspondingly deeper surface — the // relationship runs the right way. diff --git a/tests/test_manager.py b/tests/test_manager.py new file mode 100644 index 0000000..8854c1e --- /dev/null +++ b/tests/test_manager.py @@ -0,0 +1,160 @@ +"""Unit tests for the per-connection-queue ConnectionManager. + +These pin the latency-resilience contract: the broadcaster never awaits a +socket, so one slow viewer can't stall the room (the bug behind a sandbox PTY +stream arriving "all at once" for non-host viewers), while per-client ordering +stays FIFO and a hopelessly-backed-up peer is evicted rather than buffered +without bound. +""" + +import asyncio + +import cmd_chat.server.managers as managers +from cmd_chat.server.managers import ConnectionManager + + +class FakeWS: + """Minimal stand-in for a Sanic Websocket: records sends, can be made slow + or gated, and records the close code.""" + + def __init__(self, delay: float = 0.0, gate: "asyncio.Event | None" = None): + self.sent: list[str] = [] + self.delay = delay + self.gate = gate + self.closed: "int | None" = None + + async def send(self, msg: str) -> None: + if self.gate is not None: + await self.gate.wait() + if self.delay: + await asyncio.sleep(self.delay) + self.sent.append(msg) + + async def close(self, code: int = 1000) -> None: + self.closed = code + + +async def _drain(timeout: float = 1.0) -> None: + """Yield control so writer tasks can flush their queues.""" + await asyncio.sleep(0) + await asyncio.sleep(0.02) + + +def test_broadcast_is_fifo_per_client(): + async def go(): + mgr = ConnectionManager() + a, b = FakeWS(), FakeWS() + await mgr.connect("a", a, initial="INIT_A") + await mgr.connect("b", b, initial="INIT_B") + for i in range(5): + await mgr.broadcast(f"m{i}") + await _drain() + assert a.sent == ["INIT_A", "m0", "m1", "m2", "m3", "m4"] + assert b.sent == ["INIT_B", "m0", "m1", "m2", "m3", "m4"] + await mgr.disconnect("a") + await mgr.disconnect("b") + + asyncio.run(go()) + + +def test_exclude_user_is_skipped(): + async def go(): + mgr = ConnectionManager() + a, b = FakeWS(), FakeWS() + await mgr.connect("a", a) + await mgr.connect("b", b) + await mgr.broadcast("hello", exclude_user="a") + await _drain() + assert a.sent == [] + assert b.sent == ["hello"] + await mgr.disconnect("a") + await mgr.disconnect("b") + + asyncio.run(go()) + + +def test_slow_client_does_not_block_broadcast(): + """The broadcaster must return promptly even when a peer's socket is slow — + delivery is decoupled via the per-connection writer task.""" + + async def go(): + mgr = ConnectionManager() + slow = FakeWS(delay=0.5) + fast = FakeWS() + await mgr.connect("slow", slow) + await mgr.connect("fast", fast) + + start = asyncio.get_event_loop().time() + await mgr.broadcast("x") + elapsed = asyncio.get_event_loop().time() - start + # Enqueue-only: must not wait anywhere near the slow socket's 0.5s. + assert elapsed < 0.1, f"broadcast blocked on slow peer ({elapsed:.3f}s)" + + # Fast client gets it almost immediately; slow one lags behind. + await asyncio.sleep(0.05) + assert fast.sent == ["x"] + assert slow.sent == [] + await mgr.disconnect("slow") + await mgr.disconnect("fast") + + asyncio.run(go()) + + +def test_wedged_client_is_evicted_and_closed(monkeypatch): + """A peer that overflows its backlog is dropped (closed 1013) instead of + growing memory without bound; other clients are unaffected.""" + + async def go(): + monkeypatch.setattr(managers, "SEND_QUEUE_MAX", 3) + mgr = ConnectionManager() + gate = asyncio.Event() # never set → this peer's send never completes + wedged = FakeWS(gate=gate) + healthy = FakeWS() + await mgr.connect("wedged", wedged) + await mgr.connect("healthy", healthy) + await asyncio.sleep(0) # let each writer pull its first frame + + # Flood with a yield between frames: the healthy peer drains each time, so + # it never overflows; the wedged peer's writer is stuck on the gate, so + # its backlog fills (max 3) and then overflows → eviction. + n = 20 + for i in range(n): + await mgr.broadcast(f"f{i}") + await asyncio.sleep(0) + await _drain() + + assert "wedged" not in mgr.active_connections, "wedged peer must be evicted" + assert wedged.closed == 1013 + # The healthy peer was never penalised for its neighbour: still live and + # it received everything (far more than the wedged peer's tiny backlog). + assert "healthy" in mgr.active_connections + assert healthy.sent == [f"f{i}" for i in range(n)] + await mgr.disconnect("healthy") + + asyncio.run(go()) + + +def test_reconnect_replaces_old_connection(): + async def go(): + mgr = ConnectionManager() + first = FakeWS() + await mgr.connect("u", first, initial="INIT1") + second = FakeWS() + await mgr.connect("u", second, initial="INIT2") + await mgr.broadcast("after") + await _drain() + # Only the live (second) connection receives the broadcast. + assert second.sent == ["INIT2", "after"] + assert "after" not in first.sent + assert mgr.active_connections["u"].ws is second + await mgr.disconnect("u") + + asyncio.run(go()) + + +def test_send_personal_to_missing_user_is_false(): + async def go(): + mgr = ConnectionManager() + assert await mgr.send_personal("ghost", "hi") is False + + asyncio.run(go())