Hana Joo
0aad65868d
Startblock: * // Put other blockers before this line to avoid churning. * has lgtm * is approved * and then * all comments are resolved * and then PiperOrigin-RevId: 961854282 Change-Id: I9d323a8ef9f36ffe69cc62d449849eac82dc6858
1293 lines
50 KiB
Python
1293 lines
50 KiB
Python
# Copyright 2026 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Python web server for the MuJoCo Web Viewer.
|
|
|
|
This server runs in a child process with one asyncio loop on a single public
|
|
port:
|
|
|
|
* Plain HTTP GET serves static files (index.html, WASM, assets), /model.
|
|
* WebSocket /ui serves the bridge to the headless NetImgui client, which
|
|
connects over loopback TCP (see headless_ui.cc).
|
|
* WebSocket /state serves the latest-wins state payload broadcast at ~60Hz
|
|
(see state_payload.h) as binary frames, and session roster
|
|
updates as text frames.
|
|
* WebSocket /drop receives models dragged onto the page (one binary frame per
|
|
file, controller only; see drop_handler).
|
|
|
|
Because everything is served through one port, one firewall rule, port-forward,
|
|
or HTTPS tunnel exposes the whole viewer. The browser derives its WebSocket URLs
|
|
from the page origin so no client configuration is needed.
|
|
|
|
One browser at a time controls the interactive session (it owns the /ui bridge);
|
|
later browsers are rejected from /ui with WebSocket close code
|
|
_WS_CLOSE_CONTROLLER_TAKEN and can take the controller slot once it frees up.
|
|
/state connections beyond the spectator limit are closed with close code
|
|
_WS_CLOSE_SESSION_FULL.
|
|
"""
|
|
|
|
import asyncio
|
|
import ctypes
|
|
import datetime
|
|
import enum
|
|
import json
|
|
import logging
|
|
import multiprocessing
|
|
import multiprocessing.queues
|
|
import multiprocessing.sharedctypes
|
|
import os
|
|
import signal
|
|
import socket
|
|
import struct
|
|
import sys
|
|
import threading
|
|
import urllib.parse
|
|
from typing import Any, Awaitable, Callable, Optional, cast
|
|
|
|
from websockets.asyncio.server import serve
|
|
from websockets.asyncio.server import ServerConnection
|
|
from websockets.datastructures import Headers
|
|
from websockets.exceptions import ConnectionClosed
|
|
from websockets.frames import CloseCode
|
|
from websockets.http11 import Request
|
|
from websockets.http11 import Response
|
|
|
|
# Our custom WebSocket close codes in range 4xxx. The client shows a notice
|
|
# and retries slowly (see web_client_session.cc / web_client.cc).
|
|
_WS_CLOSE_CONTROLLER_TAKEN = 4001 # /ui: another browser is controlling.
|
|
_WS_CLOSE_SESSION_FULL = 4002 # /state: the spectator limit is reached.
|
|
_WS_CLOSE_INACTIVE = 4003 # /state: hidden tab kicked to free a slot.
|
|
_WS_CLOSE_NOT_CONTROLLER = 4004 # /drop: only the controller may load models.
|
|
|
|
# Upper bound on the runtime-editable spectator limit.
|
|
_MAX_SPECTATOR_HARD_CAP = 32
|
|
|
|
# How many browsers may watch in addition to the controller; the controller can
|
|
# change it at runtime (up to _MAX_SPECTATOR_HARD_CAP). The limit is about
|
|
# session behaviour — keeping the crowd and the control-request queue at a
|
|
# manageable size — not about resources: a state payload is a few KB at
|
|
# 60Hz, serialized once and sent per viewer, so even a full session costs
|
|
# only a few Mbit/s of upload. There is deliberately no config or CLI
|
|
# option (spectating exposes nothing that controlling does not).
|
|
_MAX_SPECTATORS_DEFAULT = 8
|
|
|
|
# Controller-only message prefix carrying the new spectator limit.
|
|
_MAX_SPECTATORS_PREFIX = "max_spectators="
|
|
# Sent to the page whose control claim the controller slot is reserved for.
|
|
_GRANT_MESSAGE = "grant"
|
|
# Sent by the browser after applying each state payload (keep in sync:
|
|
# web_client_session.cc). Flow control for the /state stream: at most one
|
|
# payload is in flight per client, so a slow link carries the freshest state
|
|
# it can instead of buffering seconds of stale payloads in the socket.
|
|
_STATE_ACK_MESSAGE = "state_ack"
|
|
|
|
# A lost ack must not stall the stream forever; after this long the next
|
|
# payload is sent unacked.
|
|
_STATE_ACK_TIMEOUT_SEC = 2.0
|
|
|
|
# A control grant must arrive within this window, else it moves down the queue.
|
|
_GRANT_EXPIRY_SEC = 5.0
|
|
|
|
# After a model-change restart, the controller slot stays reserved for the page
|
|
# that was controlling until it has had time to reload and reconnect; unclaimed,
|
|
# the slot then opens to everyone.
|
|
_RESTART_RESERVE_SEC = 10.0
|
|
|
|
# A controller whose tab stops running (hidden or closed without a clean
|
|
# disconnect) is released once someone is waiting for control. Detected by
|
|
# silence: a live controller tab sends input packets every frame.
|
|
_CONTROLLER_SILENT_KICK_SEC = 90.0
|
|
|
|
# Spectators are kicked after this much heartbeat silence (their tab is hidden
|
|
# or gone), freeing a viewer slot. Live tabs heartbeat every ~30s.
|
|
_SPECTATOR_SILENT_KICK_SEC = 300.0
|
|
|
|
# Default public port, and how many consecutive ports to try when it is taken
|
|
# (e.g. by another running viewer).
|
|
_DEFAULT_HTTP_PORT = 8080
|
|
_PORT_SCAN_COUNT = 20
|
|
|
|
# NetImgui's CmdVersion handshake packet is always 120 bytes.
|
|
_NETIMGUI_CMD_VERSION_SIZE = 120
|
|
|
|
# How long a browser holding the controller slot waits for the headless NetImgui
|
|
# client to (re)connect over loopback before the slot is released. The client
|
|
# reconnects within ~1s in the normal case; this only bounds the pathological
|
|
# one where it never appears (e.g., the headless UI failed to start), so a stuck
|
|
# controller cannot lock every other browser out forever.
|
|
_UI_TCP_WAIT_SEC = 30.0
|
|
|
|
# Content types for the static files the HTTP handler serves.
|
|
_CONTENT_TYPES = {
|
|
".html": "text/html; charset=utf-8",
|
|
".js": "text/javascript",
|
|
".wasm": "application/wasm",
|
|
".data": "application/octet-stream",
|
|
".json": "application/json",
|
|
".png": "image/png",
|
|
".ico": "image/x-icon",
|
|
}
|
|
|
|
|
|
class _SessionMessage(enum.StrEnum):
|
|
"""Text messages browsers send on /state."""
|
|
|
|
REQUEST_CONTROL = "request_control"
|
|
LEAVE_QUEUE = "leave_queue"
|
|
FORCE_CONTROL = "force_control"
|
|
HEARTBEAT = "heartbeat"
|
|
|
|
|
|
def _roster_line(
|
|
viewers: int,
|
|
is_controller: bool,
|
|
queue_pos: int,
|
|
queue_len: int,
|
|
max_spectators: int,
|
|
) -> str:
|
|
"""Builds a roster line, the membership broadcast every browser receives.
|
|
|
|
The roster is sent as a text frame on /state whenever the session changes (a
|
|
viewer joins or leaves, queues for control, or control moves). It tells each
|
|
browser how many viewers are connected, which role the server has currently
|
|
assigned it, and where it stands in the control queue.
|
|
|
|
Args:
|
|
viewers: The number of viewers connected to the server.
|
|
is_controller: Whether the browser is the controller.
|
|
queue_pos: The browser's position in the control queue (0 if not queued).
|
|
queue_len: The number of browsers in the queue.
|
|
max_spectators: The maximum number of spectators allowed.
|
|
|
|
Returns:
|
|
A string representing the roster line.
|
|
"""
|
|
role = "controller" if is_controller else "spectator"
|
|
return (
|
|
f"viewers={viewers};role={role};queue_pos={queue_pos};"
|
|
f"queue_len={queue_len};max_spectators={max_spectators}"
|
|
)
|
|
|
|
|
|
class _WebServerFormatter(logging.Formatter):
|
|
|
|
def format(self, record: logging.LogRecord) -> str:
|
|
level_char = record.levelname[0]
|
|
dt = datetime.datetime.fromtimestamp(record.created)
|
|
date_str = dt.strftime("%m%d")
|
|
time_str = dt.strftime("%H:%M:%S.%f")
|
|
filename = os.path.basename(record.pathname)
|
|
line = record.lineno
|
|
msg = record.getMessage()
|
|
return f"{level_char}{date_str} {time_str} {filename}:{line}] {msg}"
|
|
|
|
|
|
logger = logging.getLogger("WebServer")
|
|
|
|
|
|
class _HandshakeNoiseFilter(logging.Filter):
|
|
"""Removes handshake-failure logs from connections that simply went away.
|
|
|
|
Handshake failure are produced by two routine cases:
|
|
|
|
1. browsers speculatively open spare connections and close them unused
|
|
2. a page reload races the model-change server restart, tearing connections
|
|
down mid-handshake.
|
|
|
|
The websockets library logs each as an opening-handshake failure as at ERROR
|
|
level with a full traceback. Raising the log level would hide real handshake
|
|
errors too, hence a filter matching exactly these signatures.
|
|
"""
|
|
|
|
def filter(self, record: logging.LogRecord) -> bool:
|
|
if record.getMessage() != "opening handshake failed":
|
|
return True
|
|
exc = record.exc_info[1] if record.exc_info else None
|
|
if exc is None:
|
|
return True
|
|
if isinstance(exc, (ConnectionClosed, EOFError, OSError)):
|
|
return False # The connection died; nothing wrong with the request.
|
|
return "did not receive a valid HTTP request" not in str(exc)
|
|
|
|
|
|
def _configure_logging() -> None:
|
|
"""Configures this module's logger and reins in the websockets library's.
|
|
|
|
Called once at import time; a no-op if already configured (re-import).
|
|
"""
|
|
if logger.handlers:
|
|
return
|
|
# Info is a curated, always-on channel (browser connect/disconnect);
|
|
# connection plumbing is logged at DEBUG.
|
|
logger.setLevel(logging.INFO)
|
|
handler = logging.StreamHandler(sys.stdout)
|
|
handler.setFormatter(_WebServerFormatter())
|
|
logger.addHandler(handler)
|
|
logger.propagate = False
|
|
# The websockets library logs its own lifecycle chatter at INFO ("server
|
|
# listening", "connection rejected (200 OK)" for every HTTP request); keep
|
|
# it to warnings and errors.
|
|
logging.getLogger("websockets").setLevel(logging.WARNING)
|
|
# Filters do not inherit down the logger tree, so it must be attached to
|
|
# "websockets.server" itself.
|
|
logging.getLogger("websockets.server").addFilter(_HandshakeNoiseFilter())
|
|
|
|
|
|
_configure_logging()
|
|
|
|
|
|
def bind_public_socket(host: str, port: int = 0) -> socket.socket:
|
|
"""Binds (but does not listen on) the public HTTP listening socket.
|
|
|
|
The socket is bound in the viewer process and inherited by each server child,
|
|
so the port stays stable across server restarts (model changes) and bind
|
|
conflicts surface here, synchronously, instead of inside the child.
|
|
|
|
Args:
|
|
host: Interface to bind. "::" (the default) binds a dual-stack socket that
|
|
accepts both IPv6 and IPv4 connections, falling back to IPv4-only when the
|
|
machine has no IPv6 support.
|
|
port: A specific port, or 0 to take the first free port starting at
|
|
_DEFAULT_HTTP_PORT (so several viewers can run side by side).
|
|
|
|
Returns:
|
|
The bound socket.
|
|
|
|
Raises:
|
|
RuntimeError: If the requested port (or every scanned port) is taken.
|
|
"""
|
|
if host == "::":
|
|
try:
|
|
socket.socket(socket.AF_INET6, socket.SOCK_STREAM).close()
|
|
except OSError:
|
|
host = "0.0.0.0"
|
|
family = socket.AF_INET6 if ":" in host else socket.AF_INET
|
|
|
|
candidates = (
|
|
[port]
|
|
if port
|
|
else range(_DEFAULT_HTTP_PORT, _DEFAULT_HTTP_PORT + _PORT_SCAN_COUNT)
|
|
)
|
|
for candidate in candidates:
|
|
sock = socket.socket(family, socket.SOCK_STREAM)
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
if family == socket.AF_INET6:
|
|
# Dual-stack: also accept IPv4 connections (as IPv4-mapped addresses).
|
|
sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
|
|
try:
|
|
sock.bind((host, candidate))
|
|
return sock
|
|
except OSError:
|
|
sock.close()
|
|
if port:
|
|
raise RuntimeError(
|
|
f"Port {port} is already in use, is another web viewer running? "
|
|
"Pass a different --port, or 0 to pick one automatically."
|
|
)
|
|
raise RuntimeError(
|
|
f"Ports {_DEFAULT_HTTP_PORT}-{_DEFAULT_HTTP_PORT + _PORT_SCAN_COUNT - 1} "
|
|
"are all in use, are that many web viewers running?"
|
|
)
|
|
|
|
|
|
def bind_loopback_socket(port: int = 0) -> socket.socket:
|
|
"""Binds the loopback socket the headless NetImgui client connects to.
|
|
|
|
Defaults to an OS-assigned ephemeral port: both endpoints live in this process
|
|
tree, so no fixed number is needed and collisions between viewer instances are
|
|
impossible.
|
|
|
|
Args:
|
|
port: A specific port, or 0 to take the first free port.
|
|
|
|
Returns:
|
|
The bound socket.
|
|
"""
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.bind(("127.0.0.1", port))
|
|
return sock
|
|
|
|
|
|
def _session_id(ws: ServerConnection) -> str:
|
|
"""The page's session id (?sid=...), tying its /ui and /state together."""
|
|
path = ws.request.path if ws.request else ""
|
|
query = path.split("?", 1)[1] if "?" in path else ""
|
|
for part in query.split("&"):
|
|
if part.startswith("sid="):
|
|
return part[4:]
|
|
return f"anon-{id(ws)}"
|
|
|
|
|
|
def _terminate_process(
|
|
proc: multiprocessing.process.BaseProcess, timeout: float = 2.0
|
|
) -> None:
|
|
"""Terminates a server process, escalating to SIGKILL if it hangs.
|
|
|
|
A wedged child that outlives stop() keeps its sockets alive, which prevents
|
|
the C++ NetImgui client from ever noticing the disconnect and reconnecting to
|
|
the replacement server.
|
|
|
|
Args:
|
|
proc: The process to terminate.
|
|
timeout: The timeout for each termination step.
|
|
"""
|
|
# Grace period first: the child normally exits on its own (lifeline EOF),
|
|
# and signalling a process mid-exit makes it print noise on stderr.
|
|
proc.join(timeout=timeout)
|
|
if proc.is_alive():
|
|
proc.terminate()
|
|
proc.join(timeout=timeout)
|
|
if proc.is_alive():
|
|
logger.warning("[Http] Process %s ignored SIGTERM; killing.", proc.pid)
|
|
proc.kill()
|
|
proc.join(timeout=timeout)
|
|
|
|
|
|
def _find_static_files_dir() -> Optional[str]:
|
|
"""Locate the web viewer static files.
|
|
|
|
The `dist` directory next to this file is populated by the Emscripten build
|
|
of the `web_client` target. This is the default for packaged installations.
|
|
|
|
Development option: Set MUJOCO_WEB_VIEWER_DIST to serve from a different
|
|
directory (e.g. a local Emscripten build tree). This allows rapid iteration
|
|
on web_client changes without rebuilding and reinstalling the package.
|
|
"""
|
|
env_dir = os.environ.get("MUJOCO_WEB_VIEWER_DIST")
|
|
if env_dir:
|
|
if os.path.isdir(env_dir):
|
|
return env_dir
|
|
logger.warning(
|
|
"[Http] MUJOCO_WEB_VIEWER_DIST is not a directory: %s", env_dir
|
|
)
|
|
|
|
dist_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dist")
|
|
if os.path.isdir(dist_dir):
|
|
return dist_dir
|
|
|
|
return None
|
|
|
|
|
|
def _run_cancellable(main_loop_func: Callable[[], Awaitable[None]]) -> None:
|
|
"""Runs an asyncio loop with SIGINT/SIGTERM structured task cancellation."""
|
|
|
|
async def _wrapped() -> None:
|
|
loop = asyncio.get_running_loop()
|
|
main_task = asyncio.current_task()
|
|
|
|
def _cancel() -> None:
|
|
if main_task:
|
|
main_task.cancel()
|
|
|
|
loop.add_signal_handler(signal.SIGINT, _cancel)
|
|
loop.add_signal_handler(signal.SIGTERM, _cancel)
|
|
try:
|
|
await main_loop_func()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
finally:
|
|
# Detach the handlers (and asyncio's signal wakeup pipe) before the loop
|
|
# closes: a signal landing afterwards would try to write to the closed
|
|
# pipe and print "Exception ignored ... BrokenPipeError".
|
|
loop.remove_signal_handler(signal.SIGINT)
|
|
loop.remove_signal_handler(signal.SIGTERM)
|
|
|
|
try:
|
|
asyncio.run(_wrapped())
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
func_name = getattr(main_loop_func, "__name__", "main_loop_func")
|
|
logger.error("[%s] Unexpected error: %s", func_name, e)
|
|
|
|
|
|
def _run_server(
|
|
http_sock: socket.socket,
|
|
tcp_sock: socket.socket,
|
|
lifeline_r: int,
|
|
lifeline_w: int,
|
|
static_dir: Optional[str],
|
|
mjb_data: Optional[bytes],
|
|
shm_array: Optional[ctypes.Array],
|
|
shm_capacity: int,
|
|
generation: multiprocessing.sharedctypes.Synchronized,
|
|
drop_queue: Optional[multiprocessing.queues.Queue[Any]],
|
|
controller_sid_shared: Optional[Any],
|
|
) -> None:
|
|
"""The server process: HTTP + /ui + /state on one port, one event loop."""
|
|
|
|
# Close the inherited write end of the lifeline pipe: the read end must see
|
|
# EOF the moment the parent (the only writer) goes away — for ANY reason,
|
|
# including SIGKILL, which no signal handler or getppid poll can cover.
|
|
os.close(lifeline_w)
|
|
|
|
static_root = os.path.realpath(static_dir) if static_dir else None
|
|
|
|
def _http_headers(
|
|
content_type: str, content_length: int, cacheable: bool
|
|
) -> Headers:
|
|
headers = Headers()
|
|
headers["Content-Type"] = content_type
|
|
headers["Content-Length"] = str(content_length)
|
|
# The WASM is a pthread build: SharedArrayBuffer requires cross-origin
|
|
# isolation on every response.
|
|
headers["Cross-Origin-Opener-Policy"] = "same-origin"
|
|
headers["Cross-Origin-Embedder-Policy"] = "require-corp"
|
|
if not cacheable:
|
|
headers["Access-Control-Allow-Origin"] = "*"
|
|
headers["Cache-Control"] = "no-store"
|
|
return headers
|
|
|
|
def _serve_http(path: str) -> Response:
|
|
"""Builds the HTTP response for a non-WebSocket GET request."""
|
|
if path == "/model":
|
|
if not mjb_data:
|
|
return Response(404, "Not Found", Headers(), b"no model\n")
|
|
# The model changes on hot-swap; never serve a cached copy.
|
|
headers = _http_headers(
|
|
"application/octet-stream", len(mjb_data), cacheable=False
|
|
)
|
|
return Response(200, "OK", headers, mjb_data)
|
|
|
|
if static_root is None:
|
|
return Response(503, "Service Unavailable", Headers(), b"no dist dir\n")
|
|
|
|
rel = path.lstrip("/") or "index.html"
|
|
norm_static_root = os.path.normpath(static_root)
|
|
full = os.path.normpath(os.path.join(norm_static_root, rel))
|
|
if full != norm_static_root and not full.startswith(
|
|
norm_static_root + os.sep
|
|
):
|
|
return Response(403, "Forbidden", Headers(), b"forbidden\n")
|
|
if not os.path.isfile(full):
|
|
wasm_full = os.path.normpath(os.path.join(norm_static_root, "wasm", rel))
|
|
if os.path.isfile(wasm_full):
|
|
full = wasm_full
|
|
else:
|
|
return Response(404, "Not Found", Headers(), b"not found\n")
|
|
|
|
with open(full, "rb") as f:
|
|
body = f.read()
|
|
content_type = _CONTENT_TYPES.get(
|
|
os.path.splitext(full)[1], "application/octet-stream"
|
|
)
|
|
return Response(
|
|
200, "OK", _http_headers(content_type, len(body), cacheable=True), body
|
|
)
|
|
|
|
async def main_loop() -> None:
|
|
# The NetImgui client connection (from HeadlessUi, loopback TCP). The client
|
|
# retries every second, so after a teardown a fresh connection is fast.
|
|
tcp_reader = None
|
|
tcp_writer = None
|
|
tcp_connected = asyncio.Event()
|
|
|
|
# The browser controlling the UI (single slot) and all browsers receiving
|
|
# the state broadcast (controller + spectators), keyed by session id.
|
|
active_ui_ws = None
|
|
controller_sid: Optional[str] = None
|
|
state_clients: dict[str, ServerConnection] = {}
|
|
|
|
# Control handoff: spectators queue for the controller slot; when it frees,
|
|
# the head of the queue is granted a short exclusive claim window.
|
|
control_queue: list[str] = []
|
|
pending_grant_sid: Optional[str] = None
|
|
broadcast_roster: Any = None
|
|
grant_next: Any = None
|
|
|
|
# Strong references to fire-and-forget tasks (the grant-expiry watchdog):
|
|
# asyncio holds only a weak reference, so without this a task can be
|
|
# garbage-collected mid-sleep and never run.
|
|
background_tasks: set[asyncio.Task[Any]] = set()
|
|
max_spectators = _MAX_SPECTATORS_DEFAULT
|
|
|
|
# Liveness by absence of traffic: a hidden tab's rendering loop stops,
|
|
# so it cannot report anything — silence is the signal.
|
|
last_heartbeat: dict[str, float] = {}
|
|
last_controller_input = [0.0]
|
|
loop_time = asyncio.get_event_loop().time
|
|
|
|
def remember_controller(sid: str) -> None:
|
|
"""Records the controlling page's sid in viewer-owned shared memory.
|
|
|
|
Written on every claim (never cleared: teardown paths must not wipe it) so
|
|
that the next server, after a model-change restart, can reserve the slot
|
|
for the page that was controlling.
|
|
"""
|
|
if controller_sid_shared is not None:
|
|
controller_sid_shared.value = sid.encode("utf-8", "replace")[:63]
|
|
|
|
# Reserve the controller slot for the previous controller across a restart.
|
|
# Session ids survive page reloads (sessionStorage, see web_client.cc), so
|
|
# the reloaded controller page claims /ui with the same sid; other pages are
|
|
# rejected until then via pending_grant_sid.
|
|
if controller_sid_shared is not None and controller_sid_shared.value:
|
|
pending_grant_sid = controller_sid_shared.value.decode("utf-8")
|
|
logger.debug(
|
|
"[Session] Controller slot reserved for the previous controller"
|
|
)
|
|
|
|
async def expire_restart_reserve(reserved: str) -> None:
|
|
nonlocal pending_grant_sid, grant_next, broadcast_roster
|
|
# grant_next/broadcast_roster are defined below; by the time this timer
|
|
# fires they exist.
|
|
await asyncio.sleep(_RESTART_RESERVE_SEC)
|
|
if pending_grant_sid == reserved and active_ui_ws is None:
|
|
logger.debug("[Session] Restart reservation expired; slot open")
|
|
pending_grant_sid = None
|
|
if grant_next and broadcast_roster:
|
|
await grant_next()
|
|
await broadcast_roster()
|
|
|
|
reserve_task = asyncio.create_task(
|
|
expire_restart_reserve(pending_grant_sid)
|
|
)
|
|
background_tasks.add(reserve_task)
|
|
reserve_task.add_done_callback(background_tasks.discard)
|
|
|
|
async def broadcast_roster() -> None:
|
|
"""Sends every browser the counts, its role and its queue position."""
|
|
count = len(state_clients)
|
|
for sid, client in list(state_clients.items()):
|
|
try:
|
|
pos = control_queue.index(sid) + 1
|
|
except ValueError:
|
|
pos = 0
|
|
try:
|
|
await client.send(
|
|
_roster_line(
|
|
viewers=count,
|
|
is_controller=sid == controller_sid,
|
|
queue_pos=pos,
|
|
queue_len=len(control_queue),
|
|
max_spectators=max_spectators,
|
|
)
|
|
)
|
|
except (ConnectionClosed, ConnectionError):
|
|
pass
|
|
|
|
async def grant_next() -> None:
|
|
"""Offers the free controller slot to the pending or next queued page."""
|
|
nonlocal pending_grant_sid
|
|
if active_ui_ws is not None:
|
|
return
|
|
|
|
if pending_grant_sid is None:
|
|
while control_queue:
|
|
candidate = control_queue.pop(0)
|
|
if candidate in state_clients:
|
|
pending_grant_sid = candidate
|
|
break
|
|
if pending_grant_sid is None:
|
|
return
|
|
|
|
client = state_clients.get(pending_grant_sid)
|
|
if client is None:
|
|
pending_grant_sid = None
|
|
await grant_next()
|
|
return
|
|
|
|
try:
|
|
await client.send(_GRANT_MESSAGE)
|
|
except (ConnectionClosed, ConnectionError):
|
|
pending_grant_sid = None
|
|
await grant_next()
|
|
return
|
|
logger.info("[Session] Control granted to %s", pending_grant_sid)
|
|
await broadcast_roster()
|
|
|
|
async def expire(granted: str) -> None:
|
|
nonlocal pending_grant_sid
|
|
await asyncio.sleep(_GRANT_EXPIRY_SEC)
|
|
if pending_grant_sid == granted and active_ui_ws is None:
|
|
pending_grant_sid = None
|
|
await grant_next()
|
|
|
|
task = asyncio.create_task(expire(pending_grant_sid))
|
|
background_tasks.add(task)
|
|
task.add_done_callback(background_tasks.discard)
|
|
|
|
async def handle_session_message(sid: str, text: str) -> None:
|
|
"""A control or heartbeat message from one browser (/state text)."""
|
|
nonlocal pending_grant_sid, max_spectators
|
|
last_heartbeat[sid] = loop_time()
|
|
if text == _SessionMessage.REQUEST_CONTROL:
|
|
if sid != controller_sid and sid not in control_queue:
|
|
control_queue.append(sid)
|
|
await grant_next()
|
|
await broadcast_roster()
|
|
elif text == _SessionMessage.LEAVE_QUEUE:
|
|
if sid in control_queue:
|
|
control_queue.remove(sid)
|
|
await broadcast_roster()
|
|
elif text == _SessionMessage.FORCE_CONTROL:
|
|
if sid != controller_sid:
|
|
logger.info("[Session] %s forces control", sid)
|
|
if sid in control_queue:
|
|
control_queue.remove(sid)
|
|
pending_grant_sid = sid
|
|
if active_ui_ws is not None:
|
|
# The bridge teardown frees the slot and grant_next honors the
|
|
# pending claim; the ousted page settles into spectating.
|
|
await active_ui_ws.close(
|
|
_WS_CLOSE_CONTROLLER_TAKEN, "control taken"
|
|
)
|
|
else:
|
|
await grant_next()
|
|
elif text == _SessionMessage.HEARTBEAT:
|
|
pass # last_heartbeat is updated for every message.
|
|
elif text.startswith(_MAX_SPECTATORS_PREFIX):
|
|
if sid != controller_sid:
|
|
return # Only the controller sets the limit.
|
|
try:
|
|
value = int(text[len(_MAX_SPECTATORS_PREFIX) :])
|
|
except ValueError:
|
|
return
|
|
max_spectators = max(0, min(_MAX_SPECTATOR_HARD_CAP, value))
|
|
logger.info("[Session] Spectator limit set to %s", max_spectators)
|
|
# Enforce the new limit immediately, kicking the newest spectators
|
|
# first. A kicked page shows the session-full notice and retries,
|
|
# so it walks back in when a slot frees up.
|
|
excess = len(state_clients) - (max_spectators + 1)
|
|
for kick_sid in reversed(list(state_clients)):
|
|
if excess <= 0:
|
|
break
|
|
if kick_sid == controller_sid:
|
|
continue
|
|
client = state_clients.get(kick_sid)
|
|
if client is not None:
|
|
logger.info("[Session] Kicking spectator over the new limit")
|
|
try:
|
|
await client.close(_WS_CLOSE_SESSION_FULL, "session full")
|
|
except (ConnectionClosed, ConnectionError):
|
|
pass
|
|
excess -= 1
|
|
await broadcast_roster()
|
|
|
|
async def enforce_activity() -> None:
|
|
"""Releases a silent controller when someone waits; kicks silent tabs."""
|
|
while True:
|
|
await asyncio.sleep(5.0)
|
|
now = loop_time()
|
|
if (
|
|
active_ui_ws is not None
|
|
and last_controller_input[0] > 0
|
|
and now - last_controller_input[0] > _CONTROLLER_SILENT_KICK_SEC
|
|
and (control_queue or pending_grant_sid)
|
|
):
|
|
logger.info("[Session] Releasing control from an inactive controller")
|
|
await active_ui_ws.close(_WS_CLOSE_CONTROLLER_TAKEN, "inactive")
|
|
for sid, client in list(state_clients.items()):
|
|
if sid == controller_sid:
|
|
continue
|
|
seen = last_heartbeat.get(sid, now)
|
|
if now - seen > _SPECTATOR_SILENT_KICK_SEC:
|
|
logger.info("[Session] Kicking inactive spectator %s", sid)
|
|
try:
|
|
await client.close(_WS_CLOSE_INACTIVE, "inactive")
|
|
except (ConnectionClosed, ConnectionError):
|
|
pass
|
|
|
|
async def handle_tcp_client(
|
|
reader: asyncio.StreamReader, writer: asyncio.StreamWriter
|
|
) -> None:
|
|
nonlocal tcp_reader, tcp_writer
|
|
if tcp_writer is not None:
|
|
logger.debug("[UiBridge] Replacing previous NetImgui TCP connection")
|
|
tcp_writer.close()
|
|
tcp_reader = reader
|
|
tcp_writer = writer
|
|
# Disable Nagle's algorithm so small packets (CmdInput ~80 bytes) are
|
|
# sent immediately instead of being buffered for up to 200ms.
|
|
sock = writer.get_extra_info("socket")
|
|
if sock:
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
|
tcp_connected.set()
|
|
logger.debug("[UiBridge] NetImgui TCP connection accepted")
|
|
|
|
# --- /ui: bridge the browser to the NetImgui client ----------------------
|
|
|
|
async def ui_handler(ws: ServerConnection) -> None:
|
|
nonlocal active_ui_ws, controller_sid, pending_grant_sid
|
|
nonlocal tcp_reader, tcp_writer
|
|
if active_ui_ws is not None:
|
|
# One controller at a time. The rejected browser spectates and may retry
|
|
# via its "Take control" button once the controller leaves.
|
|
logger.debug("[UiBridge] Controller slot taken; rejecting new browser")
|
|
await ws.close(_WS_CLOSE_CONTROLLER_TAKEN, "controller slot taken")
|
|
return
|
|
sid = _session_id(ws)
|
|
if pending_grant_sid is not None and sid != pending_grant_sid:
|
|
# The slot is reserved for a granted page (queue fairness).
|
|
await ws.close(_WS_CLOSE_CONTROLLER_TAKEN, "controller slot reserved")
|
|
return
|
|
if sid == pending_grant_sid:
|
|
pending_grant_sid = None
|
|
if sid in control_queue:
|
|
control_queue.remove(sid)
|
|
active_ui_ws = ws
|
|
controller_sid = sid
|
|
remember_controller(sid)
|
|
last_controller_input[0] = loop_time()
|
|
await broadcast_roster()
|
|
|
|
# Each browser needs a fresh NetImgui session (full handshake, textures
|
|
# and draw state). Drop any connection used by a previous bridge and
|
|
# wait for the client's automatic reconnect.
|
|
if tcp_writer is not None:
|
|
tcp_writer.close()
|
|
tcp_reader = None
|
|
tcp_writer = None
|
|
tcp_connected.clear()
|
|
logger.debug("[UiBridge] Waiting for NetImgui TCP connection...")
|
|
|
|
async def wait_for_tcp() -> None:
|
|
# Loop: the connection can be replaced (and the event cleared)
|
|
# between the event firing and this task waking up.
|
|
while tcp_writer is None:
|
|
await tcp_connected.wait()
|
|
|
|
# Wait for the client, but not forever and not blind to the browser: if
|
|
# the client never appears (bounded by _UI_TCP_WAIT_SEC) or the browser
|
|
# leaves first, release the controller slot instead of holding it
|
|
# against every future browser.
|
|
tcp_task = asyncio.create_task(wait_for_tcp())
|
|
close_task = asyncio.create_task(ws.wait_closed())
|
|
done, still_pending = await asyncio.wait(
|
|
[tcp_task, close_task],
|
|
timeout=_UI_TCP_WAIT_SEC,
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
for task in still_pending:
|
|
task.cancel()
|
|
if tcp_task not in done or tcp_task.cancelled():
|
|
logger.debug("[UiBridge] No NetImgui connection; releasing slot")
|
|
if active_ui_ws is ws:
|
|
active_ui_ws = None
|
|
controller_sid = None
|
|
await grant_next()
|
|
await broadcast_roster()
|
|
await ws.close(_WS_CLOSE_CONTROLLER_TAKEN, "controller unavailable")
|
|
return
|
|
if tcp_reader is None or tcp_writer is None:
|
|
logger.debug("[UiBridge] NetImgui reader/writer None; releasing slot")
|
|
if active_ui_ws is ws:
|
|
active_ui_ws = None
|
|
controller_sid = None
|
|
await grant_next()
|
|
await broadcast_roster()
|
|
await ws.close(_WS_CLOSE_CONTROLLER_TAKEN, "controller unavailable")
|
|
return
|
|
my_reader = cast(asyncio.StreamReader, tcp_reader)
|
|
my_writer = cast(asyncio.StreamWriter, tcp_writer)
|
|
|
|
try:
|
|
# Handshake: browser CmdVersion -> client; client CmdVersion -> browser.
|
|
browser_version = await ws.recv()
|
|
my_writer.write(browser_version) # pyrefly: ignore[bad-argument-type]
|
|
await my_writer.drain()
|
|
server_version = await my_reader.readexactly(_NETIMGUI_CMD_VERSION_SIZE)
|
|
await ws.send(server_version)
|
|
logger.debug("[UiBridge] Handshake complete. Bridging.")
|
|
|
|
async def ws_to_tcp() -> None:
|
|
async for message in ws:
|
|
if isinstance(message, bytes):
|
|
last_controller_input[0] = loop_time()
|
|
my_writer.write(message)
|
|
await my_writer.drain()
|
|
|
|
async def tcp_to_ws() -> None:
|
|
while True:
|
|
data = await my_reader.read(65536)
|
|
if not data:
|
|
break
|
|
await ws.send(data)
|
|
|
|
done, pending = await asyncio.wait(
|
|
[
|
|
asyncio.create_task(ws_to_tcp()),
|
|
asyncio.create_task(tcp_to_ws()),
|
|
],
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
for task in pending:
|
|
task.cancel()
|
|
for task in done:
|
|
task.exception() # Retrieve to avoid "exception never retrieved".
|
|
except (ConnectionClosed, ConnectionError, asyncio.IncompleteReadError):
|
|
pass
|
|
finally:
|
|
logger.debug("[UiBridge] Bridge closed.")
|
|
# Tear down this bridge's TCP connection so the NetImgui client
|
|
# reconnects and re-sends everything to the next browser.
|
|
if my_writer is not None:
|
|
my_writer.close()
|
|
if tcp_writer is my_writer:
|
|
tcp_reader = None
|
|
tcp_writer = None
|
|
tcp_connected.clear()
|
|
if active_ui_ws is ws:
|
|
active_ui_ws = None
|
|
controller_sid = None
|
|
await grant_next()
|
|
await broadcast_roster()
|
|
|
|
# --- /state: latest-wins payload broadcast --------------------------------
|
|
|
|
async def state_handler(ws: ServerConnection) -> None:
|
|
sid = _session_id(ws)
|
|
# A page reconnecting after a network blip keeps its sid; its stale entry
|
|
# (dead socket, not yet cleaned up) must not count against the limit,
|
|
# since installing the new socket merely replaces it. Only a genuinely new
|
|
# sid consumes a slot.
|
|
if sid not in state_clients and len(state_clients) > max_spectators:
|
|
logger.info("[StateWS] Session full; rejecting browser")
|
|
await ws.close(_WS_CLOSE_SESSION_FULL, "session full")
|
|
return
|
|
state_clients[sid] = ws
|
|
last_heartbeat[sid] = loop_time()
|
|
logger.info("[StateWS] Browser connected (%d total)", len(state_clients))
|
|
await broadcast_roster()
|
|
|
|
def read_state() -> Optional[tuple[int, bytes]]:
|
|
# Seqlock read against update_state()'s memmove: an even generation
|
|
# unchanged across the copy means the buffer held still, so the
|
|
# payload is a single complete frame. An odd generation (write in
|
|
# flight) or a changed generation is a torn read; a few retries cover
|
|
# the microsecond-wide memmove, and skipping the tick is harmless
|
|
# under latest-wins. Returns the (generation, payload) actually read.
|
|
for _ in range(8):
|
|
gen_before = generation.value
|
|
if gen_before & 1:
|
|
continue
|
|
(used,) = struct.unpack("<I", bytes(shm_array[:4])) # pyrefly: ignore[unsupported-operation]
|
|
if used == 0 or used > shm_capacity:
|
|
return None
|
|
data = bytes(shm_array[4 : 4 + used]) # pyrefly: ignore[unsupported-operation]
|
|
if generation.value == gen_before:
|
|
return gen_before, data
|
|
return None
|
|
|
|
# See _STATE_ACK_MESSAGE: one payload in flight at a time, so the latest
|
|
# wins read below always ships the freshest state a slow link can carry
|
|
# (the camera rides this stream, so socket buffering shows up directly as
|
|
# input latency).
|
|
payload_acked = asyncio.Event()
|
|
payload_acked.set()
|
|
|
|
async def send_payloads() -> None:
|
|
last_gen = 0
|
|
while True:
|
|
await asyncio.sleep(1.0 / 60.0)
|
|
if generation.value == last_gen:
|
|
continue
|
|
try:
|
|
await asyncio.wait_for(
|
|
payload_acked.wait(), timeout=_STATE_ACK_TIMEOUT_SEC
|
|
)
|
|
except asyncio.TimeoutError:
|
|
pass # Lost ack; send anyway rather than stalling forever.
|
|
result = read_state()
|
|
if result is None:
|
|
continue
|
|
last_gen, data = result
|
|
payload_acked.clear()
|
|
await ws.send(data)
|
|
|
|
async def recv_messages() -> None:
|
|
async for message in ws:
|
|
if isinstance(message, str):
|
|
if message == _STATE_ACK_MESSAGE:
|
|
# Deliberately does NOT count as liveness: acks keep flowing from
|
|
# hidden tabs (WebSocket events are not requestAnimationFrame
|
|
# gated), and the inactivity policy must still see those tabs as
|
|
# silent.
|
|
payload_acked.set()
|
|
continue
|
|
await handle_session_message(sid, message)
|
|
|
|
try:
|
|
done, pending = await asyncio.wait(
|
|
[
|
|
asyncio.create_task(send_payloads()),
|
|
asyncio.create_task(recv_messages()),
|
|
],
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
for task in pending:
|
|
task.cancel()
|
|
for task in done:
|
|
task.exception() # Retrieve to avoid "exception never retrieved".
|
|
except (ConnectionClosed, ConnectionError):
|
|
pass
|
|
finally:
|
|
# Only tear down if this socket is still the one registered for sid: a
|
|
# reconnect may have already replaced it, and that newer connection now
|
|
# owns the queue and heartbeat entries.
|
|
if state_clients.get(sid) is ws:
|
|
del state_clients[sid]
|
|
if sid in control_queue:
|
|
control_queue.remove(sid)
|
|
last_heartbeat.pop(sid, None)
|
|
logger.info(
|
|
"[StateWS] Browser disconnected (%d total)", len(state_clients)
|
|
)
|
|
await broadcast_roster()
|
|
|
|
# --- Dispatch ----------------------------------------------------------
|
|
|
|
def process_request(
|
|
connection: ServerConnection, request: Request
|
|
) -> Optional[Response]:
|
|
del connection
|
|
|
|
# Split path and query string.
|
|
if "?" in request.path:
|
|
path, query_string = request.path.split("?", 1)
|
|
else:
|
|
path, query_string = request.path, ""
|
|
if path in ("/ui", "/state", "/drop"):
|
|
return None # Proceed with the WebSocket handshake.
|
|
|
|
# Chunked model endpoint: the client fetches /model in parallel chunks
|
|
#
|
|
# GET /model?total_bytes -> {"total_bytes": <n>}
|
|
# GET /model?offset_bytes=X&size_bytes=Y -> bytes [X, X+Y)
|
|
#
|
|
# Full model endpoint: the client fetches /model in a single request
|
|
#
|
|
# GET /model -> full model bytes
|
|
if path == "/model" and mjb_data and query_string:
|
|
params = urllib.parse.parse_qs(query_string)
|
|
if "total_bytes" in params or query_string == "total_bytes":
|
|
body = json.dumps({"total_bytes": len(mjb_data)}).encode()
|
|
headers = _http_headers(
|
|
"application/json", len(body), cacheable=False
|
|
)
|
|
return Response(200, "OK", headers, body)
|
|
try:
|
|
offset = int(params.get("offset_bytes", [0])[0])
|
|
size = int(params.get("size_bytes", [0])[0])
|
|
except (ValueError, TypeError):
|
|
return Response(400, "Bad Request", Headers(), b"bad params\n")
|
|
if size <= 0 or offset < 0 or offset >= len(mjb_data):
|
|
return Response(400, "Bad Request", Headers(), b"bad range\n")
|
|
chunk = mjb_data[offset : min(offset + size, len(mjb_data))]
|
|
headers = _http_headers(
|
|
"application/octet-stream", len(chunk), cacheable=False
|
|
)
|
|
return Response(200, "OK", headers, chunk)
|
|
|
|
return _serve_http(path)
|
|
|
|
def process_response(
|
|
connection: ServerConnection, request: Request, response: Response
|
|
) -> Optional[Response]:
|
|
"""Adds cross-origin headers to the WebSocket upgrade (101) response.
|
|
|
|
The page serves Cross-Origin-Opener-Policy: same-origin and
|
|
Cross-Origin-Embedder-Policy: require-corp so the WASM pthread build
|
|
can use SharedArrayBuffer. When accessed through an HTTPS reverse proxy,
|
|
the browser's cross-origin isolation policy blocks WebSocket connections
|
|
unless the upgrade response carries matching headers. Without them, the
|
|
handshake fails with close code CloseCode.ABNORMAL_CLOSURE.
|
|
|
|
Args:
|
|
connection: The WebSocket server connection (unused).
|
|
request: The initial HTTP upgrade request (unused).
|
|
response: The HTTP upgrade response to attach CORS headers to.
|
|
|
|
Returns:
|
|
None to allow the WebSocket upgrade handshake to proceed.
|
|
"""
|
|
del connection
|
|
del request
|
|
response.headers["Cross-Origin-Opener-Policy"] = "same-origin"
|
|
response.headers["Cross-Origin-Embedder-Policy"] = "require-corp"
|
|
response.headers["Access-Control-Allow-Origin"] = "*"
|
|
response.headers["Cross-Origin-Resource-Policy"] = "cross-origin"
|
|
return None
|
|
|
|
# --- /drop: receive a file dropped onto the browser page ---------------
|
|
|
|
async def drop_handler(ws: ServerConnection) -> None:
|
|
"""Receives dropped files and forwards them to the viewer process.
|
|
|
|
The browser opens this socket on demand and sends one binary frame per
|
|
file ([u32 path length][relative path utf-8][file bytes], see
|
|
index.html), then an empty frame as an end marker; this server closes
|
|
once it has everything. The files cross to the viewer process via
|
|
drop_queue as a dict of relative path -> bytes; WebViewer.get_drop_file()
|
|
writes them to a temporary directory for the regular drop-loading flow.
|
|
|
|
Loading a model changes the session for every connected page, so only
|
|
the controller may do it; drops from other pages are rejected.
|
|
|
|
Args:
|
|
ws: The WebSocket connection to receive from.
|
|
"""
|
|
if controller_sid is None or _session_id(ws) != controller_sid:
|
|
logger.info("[Drop] Ignored a model drop from a non-controller page")
|
|
await ws.close(_WS_CLOSE_NOT_CONTROLLER, "not the controller")
|
|
return
|
|
|
|
files: dict[str, bytes] = {}
|
|
complete = False
|
|
try:
|
|
async for message in ws:
|
|
if not isinstance(message, bytes):
|
|
continue
|
|
if not message:
|
|
complete = True # End-of-drop marker.
|
|
break
|
|
if len(message) < 4:
|
|
continue
|
|
(name_len,) = struct.unpack_from("<I", message)
|
|
if 4 + name_len > len(message):
|
|
continue
|
|
name = message[4 : 4 + name_len].decode("utf-8", "replace")
|
|
files[name] = message[4 + name_len :]
|
|
except (ConnectionClosed, ConnectionError):
|
|
pass
|
|
|
|
# Only forward a drop that arrived in full. A connection that dies before
|
|
# the end marker (e.g. a file frame exceeding the websocket size cap
|
|
# closes it with CloseCode.MESSAGE_TOO_BIG) leaves a partial file set that
|
|
# would load as a broken model, so discard it instead.
|
|
if not complete:
|
|
if files:
|
|
logger.info(
|
|
"[Drop] Incomplete drop (%d file(s) before the"
|
|
" connection closed); discarding.",
|
|
len(files),
|
|
)
|
|
elif files and drop_queue is not None:
|
|
total = sum(len(data) for data in files.values())
|
|
logger.info("[Drop] Received %d file(s), %d bytes", len(files), total)
|
|
drop_queue.put(files)
|
|
|
|
await ws.close()
|
|
|
|
async def ws_handler(ws: ServerConnection) -> None:
|
|
req = ws.request
|
|
if req is None:
|
|
await ws.close(CloseCode.POLICY_VIOLATION, "missing request")
|
|
return
|
|
path = req.path.split("?")[0]
|
|
if path == "/ui":
|
|
await ui_handler(ws)
|
|
elif path == "/state":
|
|
await state_handler(ws)
|
|
elif path == "/drop":
|
|
await drop_handler(ws)
|
|
else:
|
|
await ws.close(CloseCode.POLICY_VIOLATION, "unknown endpoint")
|
|
|
|
# Block until the lifeline pipe hits EOF, which happens exactly when the
|
|
# viewer process is gone (clean stop() close, crash, or SIGKILL). An
|
|
# orphaned server would otherwise keep serving a stale model and fight any
|
|
# replacement server for browsers.
|
|
async def watch_lifeline() -> None:
|
|
# A dedicated daemon thread rather than run_in_executor: executor threads
|
|
# are non-daemonic and are joined at interpreter exit, so a blocked
|
|
# os.read would keep this process alive for as long as the viewer holds
|
|
# the lifeline open — e.g. when the viewer is still tearing down after
|
|
# Ctrl+C hit both processes.
|
|
loop = asyncio.get_running_loop()
|
|
eof = asyncio.Event()
|
|
|
|
def wait_for_eof() -> None:
|
|
try:
|
|
os.read(lifeline_r, 1)
|
|
loop.call_soon_threadsafe(eof.set)
|
|
except (OSError, RuntimeError):
|
|
pass # fd or loop already closed, so the process is exiting anyway.
|
|
|
|
threading.Thread(target=wait_for_eof, daemon=True).start()
|
|
await eof.wait()
|
|
logger.debug("[Http] Viewer process exited; shutting down.")
|
|
|
|
tcp_server = await asyncio.start_server(handle_tcp_client, sock=tcp_sock)
|
|
|
|
# NetImgui does its own delta compression and the state payload is small;
|
|
# permessage-deflate would only add latency at 60Hz.
|
|
ws_server = await serve(
|
|
ws_handler,
|
|
sock=http_sock,
|
|
process_request=process_request,
|
|
process_response=process_response,
|
|
compression=None,
|
|
ping_interval=None,
|
|
ping_timeout=None,
|
|
open_timeout=None,
|
|
close_timeout=None,
|
|
# Big enough for model files uploaded via /drop.
|
|
max_size=2**26,
|
|
)
|
|
http_host, http_port = http_sock.getsockname()[:2]
|
|
tcp_port = tcp_sock.getsockname()[1]
|
|
logger.debug(
|
|
"[Http] Serving on http://%s:%d "
|
|
"(/, /model, /ui, /state; NetImgui TCP on 127.0.0.1:%d)",
|
|
http_host,
|
|
http_port,
|
|
tcp_port,
|
|
)
|
|
|
|
enforcer = asyncio.create_task(enforce_activity())
|
|
try:
|
|
await watch_lifeline()
|
|
finally:
|
|
enforcer.cancel()
|
|
# Close listeners explicitly to prevent hanging on open client sockets.
|
|
ws_server.close()
|
|
tcp_server.close()
|
|
|
|
_run_cancellable(main_loop)
|
|
|
|
|
|
class WebServer:
|
|
"""The Web Viewer's single-port server.
|
|
|
|
Serves the browser page, WASM, and model over HTTP; bridges the NetImgui UI
|
|
stream at /ui; and broadcasts the state payload at /state on a single public
|
|
port from a child process. State payloads are handed over through shared
|
|
memory with latest-wins semantics: update_state() may be called at any rate
|
|
from the viewer thread; browsers only ever see the newest payload.
|
|
|
|
Run the simulation with the web viewer, then visit the printed URL.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
http_sock: socket.socket,
|
|
tcp_sock: socket.socket,
|
|
static_files_dir: Optional[str] = None,
|
|
mjb_data: Optional[bytes] = None,
|
|
max_payload_size: int = 0,
|
|
drop_queue: Optional[multiprocessing.queues.Queue[Any]] = None,
|
|
controller_sid_shared: Optional[Any] = None,
|
|
) -> None:
|
|
"""Initializes the server around pre-bound listening sockets.
|
|
|
|
The sockets are bound by the viewer (see bind_public_socket /
|
|
bind_loopback_socket) and shared with the server child, so ports stay
|
|
stable across restarts and bind failures can't occur here.
|
|
|
|
Args:
|
|
http_sock: The listening socket for HTTP and WebSocket traffic.
|
|
tcp_sock: The listening socket for NetImgui traffic.
|
|
static_files_dir: The directory containing the static files to serve.
|
|
mjb_data: The model data to serve from /model.
|
|
max_payload_size: The maximum size of the state payload.
|
|
drop_queue: The queue to put dropped files onto.
|
|
controller_sid_shared: The shared value containing the controller's
|
|
session id.
|
|
"""
|
|
self.http_sock = http_sock
|
|
self.tcp_sock = tcp_sock
|
|
self.static_files_dir = static_files_dir or _find_static_files_dir()
|
|
self.mjb_data = mjb_data
|
|
|
|
# Owned by the viewer (it outlives server restarts); dropped-file bytes
|
|
# travel from the server child back to the viewer process through it.
|
|
self.drop_queue = drop_queue
|
|
|
|
# Also viewer-owned: the controlling page's session id, so a fresh
|
|
# server can reserve the controller slot for the page that was
|
|
# controlling before a model-change restart.
|
|
self.controller_sid_shared = controller_sid_shared
|
|
self._process = None
|
|
self._lifeline_w = None
|
|
|
|
# Shared memory for zero-copy payload transfer to the server process:
|
|
# [u32 length][payload bytes], so payloads can vary in size up to
|
|
# max_payload_size (see state_payload.max_state_payload_size).
|
|
self._shm_capacity = max_payload_size
|
|
self._shm_array = (
|
|
multiprocessing.RawArray(ctypes.c_char, 4 + self._shm_capacity)
|
|
if self._shm_capacity > 0
|
|
else None
|
|
)
|
|
self._generation = multiprocessing.Value("Q", 0) # uint64 counter
|
|
|
|
def update_state(self, payload: bytes) -> None:
|
|
"""Publishes the latest state payload. Called from the viewer thread."""
|
|
if self._shm_array is None:
|
|
return
|
|
|
|
if len(payload) > self._shm_capacity:
|
|
logger.error(
|
|
"[StateWS] Payload of %d bytes exceeds capacity %d; dropping.",
|
|
len(payload),
|
|
self._shm_capacity,
|
|
)
|
|
return
|
|
|
|
# Seqlock: bump the generation to odd before writing and back to even after,
|
|
# so a reader that copies the buffer while this memmove is in flight can
|
|
# detect the torn read (see send_payloads) and retry. Without it a reader
|
|
# could splice two frames and ship an invalid physics block.
|
|
with self._generation.get_lock():
|
|
self._generation.value += 1
|
|
ctypes.memmove(
|
|
self._shm_array,
|
|
struct.pack("<I", len(payload)) + payload,
|
|
4 + len(payload),
|
|
)
|
|
with self._generation.get_lock():
|
|
self._generation.value += 1
|
|
|
|
def start(self) -> None:
|
|
"""Starts the server in a background process."""
|
|
if not self.static_files_dir:
|
|
logger.warning("[Http] WARNING: Could not find web viewer static files.")
|
|
return
|
|
|
|
logger.debug("[Http] Serving web viewer from: %s", self.static_files_dir)
|
|
|
|
# Lifeline pipe: the child holds the read end and exits on EOF, which
|
|
# happens when this process closes the write end (stop()) or dies for any
|
|
# reason at all, preventing orphaned servers.
|
|
lifeline_r, self._lifeline_w = os.pipe()
|
|
|
|
# Sockets and pipe fds must be inherited, so fork explicitly (the default
|
|
# start method is fork on Linux today, but is changing upstream).
|
|
self._process = multiprocessing.get_context("fork").Process(
|
|
target=_run_server,
|
|
args=(
|
|
self.http_sock,
|
|
self.tcp_sock,
|
|
lifeline_r,
|
|
self._lifeline_w,
|
|
self.static_files_dir,
|
|
self.mjb_data,
|
|
self._shm_array,
|
|
self._shm_capacity,
|
|
self._generation,
|
|
self.drop_queue,
|
|
self.controller_sid_shared,
|
|
),
|
|
daemon=True,
|
|
)
|
|
self._process.start()
|
|
os.close(lifeline_r)
|
|
|
|
def stop(self) -> None:
|
|
"""Stops the server, escalating to SIGKILL if it hangs."""
|
|
if self._lifeline_w is not None:
|
|
os.close(self._lifeline_w) # EOF on the child's lifeline: clean exit.
|
|
self._lifeline_w = None
|
|
if self._process:
|
|
_terminate_process(self._process)
|
|
self._process = None
|