Fix visualizer server startup failures

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()

pylabrobot/visualizer/visualizer_tests.py

@@ -1,12 +1,12 @@
 import asyncio
+import errno
 import json
 import time
 import unittest
 import unittest.mock
 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,14 +245,18 @@ 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()
     self.vis.send_command = self.send_command_mock  # type: ignore[method-assign]
 
     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.

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]",