diff --git a/pylabrobot/io/executor_setup_tests.py b/pylabrobot/io/executor_setup_tests.py new file mode 100644 index 00000000000..ad6a2e02697 --- /dev/null +++ b/pylabrobot/io/executor_setup_tests.py @@ -0,0 +1,193 @@ +import asyncio +import threading +import unittest +from types import SimpleNamespace +from typing import Any, Callable, Coroutine, Optional +from unittest import mock + +from pylabrobot.io import ftdi as ftdi_module +from pylabrobot.io import hid as hid_module +from pylabrobot.io import serial as serial_module +from pylabrobot.io import usb as usb_module + + +class _ClosableDevice: + def __init__(self) -> None: + self.setup_thread: Optional[int] = None + self.close_thread: Optional[int] = None + + def close(self) -> None: + self.close_thread = threading.get_ident() + + +class _DisposableDevice: + def __init__(self) -> None: + self.setup_thread: Optional[int] = None + self.dispose_thread: Optional[int] = None + + def set_configuration(self) -> None: + self.setup_thread = threading.get_ident() + raise RuntimeError("configuration failed") + + +class ExecutorSetupTests(unittest.IsolatedAsyncioTestCase): + async def _wait_for_task(self, task: "asyncio.Task[None]") -> None: + while not task.done(): + await asyncio.wait({task}, timeout=0.01) + await task + + async def _cancel_while_worker_is_running( + self, + setup: Callable[[], Coroutine[Any, Any, None]], + started: threading.Event, + release: threading.Event, + ) -> None: + task: "asyncio.Task[None]" = asyncio.create_task(setup()) + while not started.is_set(): + await asyncio.sleep(0) + + try: + task.cancel() + await asyncio.sleep(0) + self.assertFalse(task.done(), "setup returned before its worker finished") + finally: + release.set() + + with self.assertRaises(asyncio.CancelledError): + await self._wait_for_task(task) + + async def test_cancelled_ftdi_setup_closes_device_before_executor_shutdown(self) -> None: + io = ftdi_module.FTDI.__new__(ftdi_module.FTDI) + io.human_readable_device_name = "mock FTDI" + io._device_id = "mock" + io._dev = None + io._executor = None + device = _ClosableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync() -> None: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io._dev = device # type: ignore[assignment] + + io._setup_sync = setup_sync # type: ignore[method-assign] + ftdi_error = type("MockFtdiError", (Exception,), {}) + with mock.patch.object(ftdi_module, "FtdiError", ftdi_error, create=True): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.close_thread, device.setup_thread) + self.assertIsNone(io._dev) + self.assertIsNone(io._executor) + + async def test_cancelled_hid_setup_closes_device_before_executor_shutdown(self) -> None: + io = hid_module.HID.__new__(hid_module.HID) + io.human_readable_device_name = "mock HID" + io._unique_id = "mock" + io.device = None + io._executor = None + device = _ClosableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync() -> None: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io.device = device # type: ignore[assignment] + + io._setup_sync = setup_sync # type: ignore[method-assign] + with mock.patch.object(hid_module, "USE_HID", True): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.close_thread, device.setup_thread) + self.assertIsNone(io.device) + self.assertIsNone(io._executor) + + async def test_cancelled_serial_setup_closes_port_before_executor_shutdown(self) -> None: + io = serial_module.Serial.__new__(serial_module.Serial) + io.human_readable_device_name = "mock serial" + io._ser = None + io._executor = None + device = _ClosableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync() -> str: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io._ser = device # type: ignore[assignment] + return "/dev/mock" + + io._setup_sync = setup_sync # type: ignore[method-assign] + with mock.patch.object(serial_module, "HAS_SERIAL", True): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.close_thread, device.setup_thread) + self.assertIsNone(io._ser) + self.assertIsNone(io._executor) + + async def test_cancelled_usb_setup_disposes_device_before_executor_shutdown(self) -> None: + io = usb_module.USB( + id_vendor=1, + id_product=2, + human_readable_device_name="mock USB", + packet_read_timeout=1, + read_timeout=2, + ) + device = _DisposableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync(empty_buffer: bool) -> None: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io.dev = device # type: ignore[assignment] + + def dispose_resources(dev: object) -> None: + self.assertIs(dev, device) + device.dispose_thread = threading.get_ident() + + io._setup_sync = setup_sync # type: ignore[method-assign] + fake_usb = SimpleNamespace(util=SimpleNamespace(dispose_resources=dispose_resources)) + with ( + mock.patch.object(usb_module, "USE_USB", True), + mock.patch.object(usb_module, "usb", fake_usb, create=True), + ): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.dispose_thread, device.setup_thread) + self.assertIsNone(io.dev) + self.assertIsNone(io._read_executor) + self.assertIsNone(io._write_executor) + + async def test_usb_setup_error_disposes_acquired_device_on_worker(self) -> None: + io = usb_module.USB( + id_vendor=1, + id_product=2, + human_readable_device_name="mock USB", + packet_read_timeout=1, + read_timeout=2, + ) + device = _DisposableDevice() + + def dispose_resources(dev: object) -> None: + self.assertIs(dev, device) + device.dispose_thread = threading.get_ident() + + io.get_available_devices = mock.Mock(return_value=[device]) # type: ignore[method-assign] + fake_usb = SimpleNamespace(util=SimpleNamespace(dispose_resources=dispose_resources)) + with ( + mock.patch.object(usb_module, "USE_USB", True), + mock.patch.object(usb_module, "usb", fake_usb, create=True), + self.assertRaisesRegex(RuntimeError, "configuration failed"), + ): + await self._wait_for_task(asyncio.create_task(io.setup())) + + self.assertEqual(device.dispose_thread, device.setup_thread) + self.assertIsNone(io.dev) + self.assertIsNone(io._read_executor) + self.assertIsNone(io._write_executor) diff --git a/pylabrobot/io/ftdi.py b/pylabrobot/io/ftdi.py index 17df02e7603..53b9594c506 100644 --- a/pylabrobot/io/ftdi.py +++ b/pylabrobot/io/ftdi.py @@ -178,32 +178,68 @@ def _resolve_device_serial(self) -> str: device_serial_number = cast(str, usb.util.get_string(device, device.iSerialNumber)) return device_serial_number - async def setup(self): - """Initialize the FTDI device connection with device resolution.""" + def _setup_sync(self) -> None: + """Resolve and open the device. Runs on the executor that owns all device calls.""" if self._dev is not None and not self._dev.closed: self._dev.close() + self._dev = None + + # Resolve which device to connect to + self._device_id = self._resolve_device_serial() + + # Create and open device + dev = Device( + lazy_open=True, + device_id=self.device_id, + pid=self._pid, + vid=self._vid, + interface_select=self._interface_select, + ) try: - # Resolve which device to connect to - self._device_id = self._resolve_device_serial() - - # Create and open device - self._dev = Device( - lazy_open=True, - device_id=self.device_id, - pid=self._pid, - vid=self._vid, - interface_select=self._interface_select, - ) - self._dev.open() - logger.info(f"Successfully opened FTDI device: {self.device_id}") - except FtdiError as e: - raise RuntimeError( - f"Failed to open FTDI device for '{self.human_readable_device_name}': {e}. " - "Is the device connected? Is it in use by another process? " - "Try restarting the kernel." - ) from e + dev.open() + except BaseException: + try: + dev.close() + except Exception: + logger.warning("Failed to close FTDI device after setup failure", exc_info=True) + raise + self._dev = dev - self._executor = ThreadPoolExecutor(max_workers=1) + async def setup(self): + """Initialize the FTDI device connection with device resolution.""" + if self._executor is None: + self._executor = ThreadPoolExecutor(max_workers=1) + loop = asyncio.get_running_loop() + setup_future = loop.run_in_executor(self._executor, self._setup_sync) + try: + await asyncio.shield(setup_future) + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + try: + await setup_future + except BaseException: + pass + if self._dev is not None: + try: + await loop.run_in_executor(self._executor, self._dev.close) + except Exception: + logger.warning("Failed to close FTDI device after setup failure", exc_info=True) + self._dev = None + self._shutdown_executor() + if isinstance(exc, FtdiError): + raise RuntimeError( + f"Failed to open FTDI device for '{self.human_readable_device_name}': {exc}. " + "Is the device connected? Is it in use by another process? " + "Try restarting the kernel." + ) from exc + raise + logger.info(f"Successfully opened FTDI device: {self.device_id}") + + def _shutdown_executor(self) -> None: + if self._executor is not None: + # the worker is idle here, so this does not block the event loop + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None @property def device_id(self) -> str: @@ -297,20 +333,22 @@ async def request_serial(self) -> str: return self.device_id async def stop(self): + loop = asyncio.get_running_loop() if self._dev is not None: - self.dev.close() - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None + await loop.run_in_executor(self._executor, self.dev.close) + self._dev = None + self._shutdown_executor() async def write(self, data: bytes) -> int: """Write data to the device. Returns the number of bytes written.""" logger.log(LOG_LEVEL_IO, "[%s] write %s", self._device_id, data) capturer.record(FTDICommand(device_id=self.device_id, action="write", data=data.hex())) - return cast(int, self.dev.write(data)) + loop = asyncio.get_running_loop() + return cast(int, await loop.run_in_executor(self._executor, self.dev.write, data)) async def read(self, num_bytes: int = 1) -> bytes: - data = self.dev.read(num_bytes) + loop = asyncio.get_running_loop() + data = await loop.run_in_executor(self._executor, self.dev.read, num_bytes) if len(data) != 0: logger.log(LOG_LEVEL_IO, "[%s] read %s", self._device_id, data) capturer.record( @@ -323,7 +361,8 @@ async def read(self, num_bytes: int = 1) -> bytes: return cast(bytes, data) async def readline(self) -> bytes: # type: ignore # very dumb it's reading from pyserial - data = self.dev.readline() + loop = asyncio.get_running_loop() + data = await loop.run_in_executor(self._executor, self.dev.readline) if len(data) != 0: logger.log(LOG_LEVEL_IO, "[%s] readline %s", self._device_id, data) capturer.record(FTDICommand(device_id=self.device_id, action="readline", data=data.hex())) diff --git a/pylabrobot/io/hid.py b/pylabrobot/io/hid.py index 1b624da6962..1a091522021 100644 --- a/pylabrobot/io/hid.py +++ b/pylabrobot/io/hid.py @@ -43,17 +43,8 @@ def __init__( if get_capture_or_validation_active(): raise RuntimeError("Cannot create a new HID object while capture or validation is active") - async def setup(self): - """ - Sets up the HID device by enumerating connected devices, matching the specified - VID, PID, and optional serial number, and opening a connection to the device. - """ - if not USE_HID: - raise RuntimeError( - "hid is not installed. Install with: pip install pylabrobot[hid]. " - f"Import error: {_HID_IMPORT_ERROR}" - ) - + def _setup_sync(self) -> None: + """Enumerate, match and open the device. Runs on the executor that owns all device calls.""" # --- 1. Enumerate all HID devices --- all_devices = hid.enumerate() candidates = [ @@ -102,19 +93,55 @@ async def setup(self): self.device = hid.Device( path=chosen["path"] # safer than vid/pid/serial triple ) + + async def setup(self): + """ + Sets up the HID device by enumerating connected devices, matching the specified + VID, PID, and optional serial number, and opening a connection to the device. + """ + if not USE_HID: + raise RuntimeError( + "hid is not installed. Install with: pip install pylabrobot[hid]. " + f"Import error: {_HID_IMPORT_ERROR}" + ) + self._executor = ThreadPoolExecutor(max_workers=1) + loop = asyncio.get_running_loop() + setup_future = loop.run_in_executor(self._executor, self._setup_sync) + try: + await asyncio.shield(setup_future) + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + try: + await setup_future + except BaseException: + pass + if self.device is not None: + try: + await loop.run_in_executor(self._executor, self.device.close) + except Exception: + logger.warning("Failed to close HID device after setup failure", exc_info=True) + self.device = None + self._shutdown_executor() + raise logger.log(LOG_LEVEL_IO, "Opened HID device %s", self._unique_id) capturer.record(HIDCommand(device_id=self._unique_id, action="open", data="")) + def _shutdown_executor(self) -> None: + if self._executor is not None: + # the worker is idle here, so this does not block the event loop + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None + async def stop(self): if self.device is not None: - self.device.close() + loop = asyncio.get_running_loop() + await loop.run_in_executor(self._executor, self.device.close) + self.device = None logger.log(LOG_LEVEL_IO, "Closing HID device %s", self._unique_id) capturer.record(HIDCommand(device_id=self._unique_id, action="close", data="")) - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None + self._shutdown_executor() async def write(self, data: bytes, report_id: bytes = b"\x00"): r"""Writes data to the HID device. diff --git a/pylabrobot/io/serial.py b/pylabrobot/io/serial.py index 26ddf801d96..6ddc7eaa92d 100644 --- a/pylabrobot/io/serial.py +++ b/pylabrobot/io/serial.py @@ -147,6 +147,39 @@ async def setup(self): loop = asyncio.get_running_loop() self._executor = ThreadPoolExecutor(max_workers=1) + setup_future = loop.run_in_executor(self._executor, self._setup_sync) + try: + resolved_port = await asyncio.shield(setup_future) + self._port = resolved_port + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + try: + await setup_future + except BaseException: + pass + if self._ser is not None: + try: + await loop.run_in_executor(self._executor, self._ser.close) + except Exception: + logger.warning("Failed to close serial port after setup failure", exc_info=True) + self._ser = None + self._shutdown_executor() + raise + + assert self._ser is not None + + def _shutdown_executor(self) -> None: + if self._executor is not None: + # the worker is idle here, so this does not block the event loop + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None + + def _setup_sync(self) -> str: + """Resolve the port and open the connection. Runs on the executor that owns all device calls. + + Returns the resolved port. + """ + # 1. VID:PID specified - port maybe if self._vid is not None and self._pid is not None: matching_ports = [ @@ -183,8 +216,8 @@ async def setup(self): "Please specify the correct port address explicitly (e.g. /dev/ttyUSB0 or COM3)." ) - def _open_serial() -> serial.Serial: - return serial.Serial( + try: + self._ser = serial.Serial( port=candidate_port, baudrate=self.baudrate, bytesize=self.bytesize, @@ -196,22 +229,13 @@ def _open_serial() -> serial.Serial: dsrdtr=self.dsrdtr, xonxoff=self.xonxoff, ) - - try: - self._ser = await loop.run_in_executor(self._executor, _open_serial) - - except serial.SerialException as e: + except serial.SerialException: logger.error( f"Could not connect to device '{self.human_readable_device_name}', is it in use by a different notebook/process?" ) - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None - raise e + raise - assert self._ser is not None - - self._port = candidate_port + return candidate_port async def stop(self): """Close the serial device.""" @@ -223,9 +247,7 @@ async def stop(self): raise RuntimeError(f"Call setup() first for device '{self.human_readable_device_name}'.") await loop.run_in_executor(self._executor, self._ser.close) - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None + self._shutdown_executor() async def write(self, data: bytes): """Write data to the serial device.""" diff --git a/pylabrobot/io/usb.py b/pylabrobot/io/usb.py index 8e4469f1577..192ad659f87 100644 --- a/pylabrobot/io/usb.py +++ b/pylabrobot/io/usb.py @@ -383,21 +383,8 @@ def ctrl_transfer( return bytearray(res) - async def setup(self, empty_buffer=True): - """Initialize the USB connection to the machine.""" - - if self.dev is not None: - # previous setup did not properly finish, - # or we are re-initializing the device. - logger.warning("USB device already connected. Closing previous connection.") - await self.stop() - - if not USE_USB: - raise RuntimeError( - "pyusb/libusb is not installed. Install with: pip install pylabrobot[usb]. " - f"Import error: {_USB_IMPORT_ERROR}. " - "https://docs.pylabrobot.org/installation.html" - ) + def _setup_sync(self, empty_buffer: bool) -> None: + """Open the device, resolve the endpoints and drain stale packets. Runs off the event loop.""" logger.info("Finding USB device...") @@ -457,22 +444,67 @@ async def setup(self, empty_buffer=True): while self._read_packet() is not None: pass + async def setup(self, empty_buffer=True): + """Initialize the USB connection to the machine.""" + + if self.dev is not None: + # previous setup did not properly finish, + # or we are re-initializing the device. + logger.warning("USB device already connected. Closing previous connection.") + await self.stop() + + if not USE_USB: + raise RuntimeError( + "pyusb/libusb is not installed. Install with: pip install pylabrobot[usb]. " + f"Import error: {_USB_IMPORT_ERROR}. " + "https://docs.pylabrobot.org/installation.html" + ) + self._read_executor = ThreadPoolExecutor(max_workers=self.max_workers) self._write_executor = ThreadPoolExecutor(max_workers=self.max_workers) + loop = asyncio.get_running_loop() + setup_future = loop.run_in_executor(self._read_executor, self._setup_sync, empty_buffer) + try: + await asyncio.shield(setup_future) + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + try: + await setup_future + except BaseException: + pass + if self.dev is not None: + dev = self.dev + try: + await loop.run_in_executor(self._read_executor, lambda: usb.util.dispose_resources(dev)) + except Exception: + logger.warning("Failed to dispose USB device after setup failure", exc_info=True) + self.dev = None + self.read_endpoint = None + self.write_endpoint = None + self._shutdown_executors() + raise + + def _shutdown_executors(self) -> None: + for executor in (self._read_executor, self._write_executor): + if executor is not None: + # the workers are idle here, so this does not block the event loop + executor.shutdown(wait=False, cancel_futures=True) + self._read_executor = self._write_executor = None + async def stop(self): """Close the USB connection to the machine. Safe to call multiple times.""" if self.dev is None: + self._shutdown_executors() return logger.warning("Closing connection to USB device.") - usb.util.dispose_resources(self.dev) + loop = asyncio.get_running_loop() + dev = self.dev + await loop.run_in_executor(self._read_executor, lambda: usb.util.dispose_resources(dev)) self.dev = None - for executor in (self._read_executor, self._write_executor): - if executor is not None: - executor.shutdown(wait=True) - self._read_executor = self._write_executor = None + self._shutdown_executors() def serialize(self) -> dict: """Serialize the backend to a dictionary."""