Matija Kecman
84950fa371
This change introduces parallel chunked downloading for large model files (.mjb) directly into WASM linear memory, enables model loading without full page reloads and reduces memory overhead. Key changes: - Implement a chunked model endpoint in the Python web server to support range requests. - Add parallel chunked fetching in the frontend with retry logic and a single-fetch fallback. - Increase initial WASM memory to 3 GB to accommodate large models and prevent heap fragmentation. - Support in-place model reloading in the C++ client, including texture cache invalidation when the Filament context is recreated. - Display a model download progress bar and model parsing/loading banner to the UI. - Fixes model drag and drop (caused by typo in sessionId, corrected to session_id). PiperOrigin-RevId: 960249236 Change-Id: Icac89e6a4ca099882aaf9b111744c1b6ab0cc6c1
1293 lines
49 KiB
Python
1293 lines
49 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)
|
|
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]))
|
|
if used == 0 or used > shm_capacity:
|
|
return None
|
|
data = bytes(shm_array[4 : 4 + used])
|
|
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
|