diff --git a/pylabrobot/visualizer/visualizer.py b/pylabrobot/visualizer/visualizer.py index 466165db77e..9033b19859f 100644 --- a/pylabrobot/visualizer/visualizer.py +++ b/pylabrobot/visualizer/visualizer.py @@ -1,4 +1,6 @@ import asyncio +import concurrent.futures +import errno import functools import http.server import inspect @@ -8,7 +10,6 @@ import os import re import threading -import time import webbrowser from typing import Any, Dict, List, Optional, Tuple @@ -27,6 +28,14 @@ logger = logging.getLogger(__name__) +_SERVER_PORT_ATTEMPTS = 100 +_SERVER_STARTUP_TIMEOUT = 10.0 + + +def _is_address_in_use(error: OSError) -> bool: + """Return whether a server failed because its requested address is already in use.""" + return error.errno == errno.EADDRINUSE + @functools.lru_cache(maxsize=None) def _get_public_methods(cls: type) -> list: @@ -127,9 +136,9 @@ def __init__( Args: host: The hostname of the file and websocket server. ws_port: The port of the websocket server. If this port is in use, the port will be - incremented until a free port is found. + incremented until a free port is found. Use `0` to let the operating system choose a port. fs_port: The port of the file server. If this port is in use, the port will be incremented - until a free port is found. + until a free port is found. Use `0` to let the operating system choose a port. open_browser: If `True`, the visualizer will open a browser window when it is started. name: A custom name to display in the browser header. If ``None``, the filename of the calling script or notebook is detected automatically. @@ -267,27 +276,27 @@ async def _socket_handler( """Handle a new websocket connection. Save the websocket connection store received messages in `self.received`.""" - while True: - try: + try: + while True: message = await websocket.recv() - except websockets.exceptions.ConnectionClosed: - return - except asyncio.CancelledError: - return - - data = json.loads(message) - if data.get("id") in self._pending_response_ids: - self.received.append(data) - - # If the event is "ready", then we can save the connection and send the saved messages. - if data.get("event") == "ready": - self._websocket = websocket - await self._send_resources_and_state() - - if "event" in data: - await self.handle_event(data.get("event"), data) - else: - logger.warning("Unhandled message: %s", message) + data = json.loads(message) + if data.get("id") in self._pending_response_ids: + self.received.append(data) + + # If the event is "ready", then we can save the connection and send the saved messages. + if data.get("event") == "ready": + self._websocket = websocket + await self._send_resources_and_state() + + if "event" in data: + await self.handle_event(data.get("event"), data) + else: + logger.warning("Unhandled message: %s", message) + except (websockets.exceptions.ConnectionClosed, asyncio.CancelledError): + return + finally: + if self._websocket is websocket: + self._websocket = None def _assemble_command( self, @@ -485,7 +494,7 @@ async def setup(self): raise RuntimeError("The visualizer has already been started.") await self._run_ws_server() - self._run_file_server() + await self._run_file_server() self.setup_finished = True async def _run_ws_server(self): @@ -496,33 +505,39 @@ async def _run_ws_server(self): async def run_server(): self._stop_ = self.loop.create_future() - while True: + for attempt in range(_SERVER_PORT_ATTEMPTS): try: - async with websockets.asyncio.server.serve(self._socket_handler, self.host, self.ws_port): + async with websockets.asyncio.server.serve( + self._socket_handler, self.host, self.ws_port + ) as server: + if self.ws_port == 0: + self.ws_port = next(iter(server.sockets)).getsockname()[1] print(f"Websocket server started at http://{self.host}:{self.ws_port}") - lock.release() + startup.set_result(None) await self.stop_ - break + return except asyncio.CancelledError: - pass - except OSError: - # If the port is in use, try the next port. + raise + except OSError as error: + if not _is_address_in_use(error) or attempt == _SERVER_PORT_ATTEMPTS - 1: + raise self.ws_port += 1 def start_loop(): - self.loop.run_until_complete(run_server()) + try: + self.loop.run_until_complete(run_server()) + except BaseException as error: + if not startup.done(): + startup.set_exception(error) - # Acquire a lock to prevent setup from returning until the server is running. - lock = threading.Lock() - lock.acquire() + startup: concurrent.futures.Future[None] = concurrent.futures.Future() self._loop = asyncio.new_event_loop() self._t = threading.Thread(target=start_loop, daemon=True) self.t.start() - while lock.locked(): - time.sleep(0.001) + await asyncio.wait_for(asyncio.wrap_future(startup), timeout=_SERVER_STARTUP_TIMEOUT) - def _run_file_server(self): + async def _run_file_server(self): """Start a simple webserver to serve static files.""" dirname = os.path.dirname(__file__) @@ -532,7 +547,7 @@ def _run_file_server(self): "Could not find Visualizer files. Please run from the root of the repository." ) - def start_server(lock): + def run_server(): ws_port, fs_port, source_filename = self.ws_port, self.fs_port, self._source_filename favicon_path = self._favicon_path liquid_color = self._liquid_color @@ -579,36 +594,44 @@ def do_GET(self) -> None: else: return super().do_GET() - while True: + for attempt in range(_SERVER_PORT_ATTEMPTS): try: self._httpd = http.server.HTTPServer( (self.host, self.fs_port), QuietSimpleHTTPRequestHandler, ) + if self.fs_port == 0: + self.fs_port = self._httpd.server_port print( f"File server started at http://{self.host}:{self.fs_port} . " "Open this URL in your browser." ) - lock.release() + startup.set_result(None) break - except OSError: + except OSError as error: + if not _is_address_in_use(error) or attempt == _SERVER_PORT_ATTEMPTS - 1: + raise self.fs_port += 1 self.httpd.serve_forever() - lock = threading.Lock() - lock.acquire() + def start_server(): + try: + run_server() + except BaseException as error: + if not startup.done(): + startup.set_exception(error) + + startup: concurrent.futures.Future[None] = concurrent.futures.Future() self._fst = threading.Thread( name="visualizer_fs", target=start_server, - args=(lock,), daemon=True, ) self.fst.start() # Wait for the server to start before opening the browser so that we can get the correct port. - while lock.locked(): - time.sleep(0.001) + await asyncio.wait_for(asyncio.wrap_future(startup), timeout=_SERVER_STARTUP_TIMEOUT) if self.open_browser: webbrowser.open(f"http://{self.host}:{self.fs_port}") @@ -630,12 +653,18 @@ async def stop(self): self._fst = None # -- websocket -- - if self.has_connection(): + had_connection = self.has_connection() + if had_connection: # send stop event to the browser await self.send_command("stop", wait_for_response=False) - # must be thread safe, because event loop is running in a separate thread + # Must be thread safe because the server event loop runs in a separate thread. The server + # must also be stopped when no browser has connected yet. + server_thread = self.t + if not self.stop_.done(): self.loop.call_soon_threadsafe(self.stop_.set_result, "done") + if not had_connection: + await asyncio.to_thread(server_thread.join, _SERVER_STARTUP_TIMEOUT) # Clear all relevant attributes. self.received.clear() diff --git a/pylabrobot/visualizer/visualizer_tests.py b/pylabrobot/visualizer/visualizer_tests.py index 1a74b467be9..3467a5d20e2 100644 --- a/pylabrobot/visualizer/visualizer_tests.py +++ b/pylabrobot/visualizer/visualizer_tests.py @@ -1,4 +1,5 @@ import asyncio +import errno import json import time import unittest @@ -6,7 +7,6 @@ import urllib.request from typing import Optional -import pytest import websockets from pylabrobot.__version__ import STANDARD_FORM_JSON_VERSION @@ -93,12 +93,11 @@ def test_short_hex_raises(self): class VisualizerSetupStopTests(unittest.IsolatedAsyncioTestCase): """Tests for the setup and stop methods of the visualizer backend.""" - @pytest.mark.timeout(20) async def test_setup_stop(self): """Test that the thread is started and stopped correctly.""" r = Resource(size_x=100, size_y=100, size_z=100, name="root") - vis = Visualizer(r, open_browser=False) + vis = Visualizer(r, ws_port=0, fs_port=0, open_browser=False) async def setup_stop_single(): await vis.setup() @@ -113,13 +112,52 @@ async def setup_stop_single(): await setup_stop_single() +class VisualizerServerStartupFailureTests(unittest.IsolatedAsyncioTestCase): + """Tests for failures while starting the visualizer servers.""" + + def setUp(self): + resource = Resource(size_x=100, size_y=100, size_z=100, name="root") + self.vis = Visualizer(resource, open_browser=False) + + async def test_websocket_permission_error_is_propagated(self): + error = PermissionError(errno.EACCES, "permission denied") + with unittest.mock.patch( + "pylabrobot.visualizer.visualizer.websockets.asyncio.server.serve", + side_effect=error, + ): + with self.assertRaises(PermissionError): + await asyncio.wait_for(self.vis._run_ws_server(), timeout=1) + + async def test_websocket_address_retries_are_bounded(self): + error = OSError(errno.EADDRINUSE, "address already in use") + with ( + unittest.mock.patch( + "pylabrobot.visualizer.visualizer.websockets.asyncio.server.serve", + side_effect=error, + ) as serve, + unittest.mock.patch("pylabrobot.visualizer.visualizer._SERVER_PORT_ATTEMPTS", 3), + ): + with self.assertRaises(OSError): + await asyncio.wait_for(self.vis._run_ws_server(), timeout=1) + self.assertEqual(serve.call_count, 3) + + async def test_file_server_permission_error_is_propagated(self): + error = PermissionError(errno.EACCES, "permission denied") + with unittest.mock.patch( + "pylabrobot.visualizer.visualizer.http.server.HTTPServer", + side_effect=error, + ): + with self.assertRaises(PermissionError): + await asyncio.wait_for(self.vis._run_file_server(), timeout=1) + + class VisualizerServerTests(unittest.IsolatedAsyncioTestCase): """Tests for servers (ws/fs).""" async def asyncSetUp(self): await super().asyncSetUp() self.r = Resource(size_x=100, size_y=100, size_z=100, name="root") - self.vis = Visualizer(self.r, open_browser=False) + self.vis = Visualizer(self.r, ws_port=0, fs_port=0, open_browser=False) await self.vis.setup() ws_port = self.vis.ws_port # port may change if port is already in use @@ -128,8 +166,8 @@ async def asyncSetUp(self): async def asyncTearDown(self): await super().asyncTearDown() - await self.vis.stop() await self.client.close() + await self.vis.stop() def test_get_index_html(self): """Test that the index.html file is returned.""" @@ -182,7 +220,7 @@ class VisualizerShowMachineToolsTests(unittest.IsolatedAsyncioTestCase): async def test_show_machine_tools_at_start_false(self): """When show_machine_tools_at_start=False, the show_machine_tools event should not be sent.""" r = Resource(size_x=100, size_y=100, size_z=100, name="root") - vis = Visualizer(r, open_browser=False, show_machine_tools_at_start=False) + vis = Visualizer(r, ws_port=0, fs_port=0, open_browser=False, show_machine_tools_at_start=False) vis.send_command = unittest.mock.AsyncMock() # type: ignore[method-assign] await vis.setup() @@ -207,7 +245,7 @@ async def asyncSetUp(self): await super().asyncSetUp() self.maxDiff = None self.r = Resource(size_x=100, size_y=100, size_z=100, name="root") - self.vis = Visualizer(self.r, open_browser=False) + self.vis = Visualizer(self.r, ws_port=0, fs_port=0, open_browser=False) # mock the send_command method to catch the events self.send_command_mock = unittest.mock.AsyncMock() @@ -215,6 +253,10 @@ async def asyncSetUp(self): await self.vis.setup() + async def asyncTearDown(self): + await self.vis.stop() + await super().asyncTearDown() + async def _wait_for_event(self, event: str, data_key: Optional[str] = None, timeout: float = 5.0): """Wait until the most recent send_command call is ``event`` (optionally carrying ``data_key`` in its data), yielding to the loop. diff --git a/pyproject.toml b/pyproject.toml index 7a54b559bd7..54a40e54b60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,6 @@ xarm = ["xarm-python-sdk"] all = ["PyLabRobot[serial,usb,ftdi,hid,modbus,opentrons,sila,pico,xarm]"] test = [ "pytest", - "pytest-timeout", ] dev = [ "PyLabRobot[all,test]",