io: keep a cancelled read's bytes, and release a device a failed setup claimed

pylabrobot/io/ftdi.py

@@ -33,8 +33,9 @@
 class FTDICommand(Command):
   data: str
 
-  def __init__(self, device_id: str, action: str, data: str):
-    super().__init__(module="ftdi", device_id=device_id, action=action)
+  def __init__(self, device_id: str, action: str, data: str, module: str = "ftdi"):
+    super().__init__(module=module, device_id=device_id, action=action)
+    self.data = data
 
 
 class FTDI(IOBase):
@@ -83,6 +84,9 @@ def __init__(
     # Will be resolved in setup()
     self._dev: Optional[Device] = None
     self._executor: Optional[ThreadPoolExecutor] = None
+    # Bytes off the wire that no read has taken yet, and the read still in flight, if any.
+    self._unread = bytearray()
+    self._pending_read: Optional["asyncio.Future"] = None
 
     if get_capture_or_validation_active():
       raise RuntimeError(
@@ -250,6 +254,7 @@ async def set_dtr(self, level: bool):
   async def usb_reset(self):
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(self._executor, lambda: self.dev.ftdi_fn.ftdi_usb_reset())
+    self._discard_reads()
     logger.log(LOG_LEVEL_IO, "[%s] usb_reset", self._device_id)
     capturer.record(FTDICommand(device_id=self.device_id, action="usb_reset", data=""))
 
@@ -288,6 +293,7 @@ async def set_flowctrl(self, flowctrl: int):
   async def usb_purge_rx_buffer(self):
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(self._executor, lambda: self.dev.ftdi_fn.ftdi_usb_purge_rx_buffer())
+    self._discard_reads()
     logger.log(LOG_LEVEL_IO, "[%s] usb_purge_rx_buffer", self._device_id)
     capturer.record(FTDICommand(device_id=self.device_id, action="usb_purge_rx_buffer", data=""))
 
@@ -318,35 +324,101 @@ async def stop(self):
       await loop.run_in_executor(self._executor, self.dev.close)
       self._dev = None
     self._shutdown_executor()
+    self._discard_reads()
+
+  def _discard_reads(self) -> None:
+    """Drop read data buffered here and in flight, once the caller has declared it stale."""
+    self._unread.clear()
+    self._pending_read = None
+
+  def _keep(self, chunk) -> None:
+    """Add what came off the wire to the buffer that read() serves from.
+
+    pylibftdi returns str, decoded with latin-1, when the device was opened in text mode. setup()
+    opens it in byte mode, but latin-1 round-trips every byte value, so encoding back gives
+    exactly what was on the wire either way.
+    """
+    self._unread.extend(chunk if isinstance(chunk, bytes) else chunk.encode("latin-1"))
+
+  def _buffer_read(self, future: "asyncio.Future") -> None:
+    """Keep the bytes of a read whose caller was cancelled, for the next read to take."""
+    if future is not self._pending_read:
+      return  # a purge, reset or close has since declared everything in flight stale
+    self._pending_read = None
+    if future.cancelled() or future.exception() is not None:
+      return
+    self._keep(future.result())
 
   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()))
     loop = asyncio.get_running_loop()
-    return cast(int, await loop.run_in_executor(self._executor, self.dev.write, data))
+    # shield: a call handed to the worker cannot be recalled, so cancelling the caller would
+    # otherwise leave the device holding a partial command, or none, with no way to tell which.
+    return cast(
+      int, await asyncio.shield(loop.run_in_executor(self._executor, self.dev.write, data))
+    )
 
   async def read(self, num_bytes: int = 1) -> bytes:
-    loop = asyncio.get_running_loop()
-    data = await loop.run_in_executor(self._executor, self.dev.read, num_bytes)
+    if not self._unread:
+      pending = self._pending_read
+      if pending is None or pending.done():
+        loop = asyncio.get_running_loop()
+        pending = loop.run_in_executor(self._executor, self.dev.read, num_bytes)
+        self._pending_read = pending
+      try:
+        # shield: the worker keeps going after the caller is gone, so its bytes are held for the
+        # next read instead of being dropped part-way through a response.
+        chunk = await asyncio.shield(pending)
+      except asyncio.CancelledError:
+        pending.add_done_callback(self._buffer_read)
+        raise
+      if pending is self._pending_read:
+        self._pending_read = None
+        self._keep(chunk)
+    data = bytes(self._unread[:num_bytes])
+    del self._unread[:num_bytes]
     if len(data) != 0:
       logger.log(LOG_LEVEL_IO, "[%s] read %s", self._device_id, data)
-      capturer.record(
-        FTDICommand(
-          device_id=self.device_id,
-          action="read",
-          data=data if isinstance(data, str) else data.hex(),
-        )
-      )
-    return cast(bytes, data)
+      capturer.record(FTDICommand(device_id=self.device_id, action="read", data=data.hex()))
+    return data
+
+  async def readline(  # type: ignore # very dumb it's reading from pyserial
+    self, terminator: bytes = b"\n", timeout: Optional[float] = None
+  ) -> bytes:
+    """Read until `terminator`, returning the line with the terminator still on it.
+
+    Assembled from single-byte reads: pylibftdi's own `readline` raises `TypeError` unless the
+    device was opened in text mode, and `setup` opens it in byte mode.
+
+    Args:
+      terminator: byte sequence that ends the line.
+      timeout: seconds to wait for a complete line. None waits indefinitely.
 
-  async def readline(self) -> bytes:  # type: ignore # very dumb it's reading from pyserial
+    Raises:
+      ValueError: if `terminator` is empty.
+      TimeoutError: if no complete line arrives within `timeout`.
+    """
+    if not terminator:
+      raise ValueError("terminator must be at least one byte")
     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()))
-    return cast(bytes, data)
+    deadline = None if timeout is None else loop.time() + timeout
+    line = bytearray()
+    while not line.endswith(terminator):
+      chunk = await self.read(1)
+      if chunk:
+        line.extend(chunk)
+        continue
+      if deadline is not None and loop.time() >= deadline:
+        raise TimeoutError(
+          f"'{self.human_readable_device_name}' sent no complete line within {timeout} s; "
+          f"received {bytes(line)!r} so far."
+        )
+      # An empty read already cost a latency-timer period, so yield without adding delay.
+      await asyncio.sleep(0)
+    logger.log(LOG_LEVEL_IO, "[%s] readline %s", self._device_id, bytes(line))
+    return bytes(line)
 
   def serialize(self):
     return {
@@ -493,19 +565,11 @@ async def read(self, num_bytes: int = 1) -> bytes:
       next_command.module == "ftdi"
       and next_command.device_id == self._device_id
       and next_command.action == "read"
-      and len(next_command.data) == num_bytes
+      # data is hex, so two characters per byte, and a short read is normal.
+      and len(next_command.data) <= num_bytes * 2
     ):
       raise ValidationError(f"Next line is {next_command}, expected FTDI read {self._device_id}")
     return bytes.fromhex(next_command.data)
 
-  async def readline(self) -> bytes:  # type: ignore # very dumb it's reading from pyserial
-    next_command = FTDICommand(**self.cr.next_command())
-    if not (
-      next_command.module == "ftdi"
-      and next_command.device_id == self._device_id
-      and next_command.action == "readline"
-    ):
-      raise ValidationError(
-        f"Next line is {next_command}, expected FTDI readline {self._device_id}"
-      )
-    return bytes.fromhex(next_command.data)
+  # readline is inherited: it assembles the line from read(), so a replay goes through the
+  # recorded reads rather than a "readline" command that is no longer recorded.

pylabrobot/io/ftdi_tests.py

@@ -1,9 +1,17 @@
+import asyncio
 import logging
+import tempfile
+import time
 import unittest
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Any, List, Union
 from unittest import mock
 
+from pylabrobot.io import capture as capture_module
 from pylabrobot.io import ftdi as ftdi_module
-from pylabrobot.io.ftdi import FTDI, HAS_PYLIBFTDI, HAS_PYUSB
+from pylabrobot.io.capture import CaptureReader, capturer
+from pylabrobot.io.ftdi import FTDI, HAS_PYLIBFTDI, HAS_PYUSB, FTDICommand, FTDIValidator
 from pylabrobot.io.validation_utils import LOG_LEVEL_IO
 
 
@@ -53,7 +61,8 @@ async def test_empty_read_is_not_logged_or_captured(self) -> None:
   async def test_empty_readline_is_not_logged_or_captured(self) -> None:
     dev = self._ftdi(b"")
     with mock.patch.object(ftdi_module.capturer, "record") as mock_record:
-      self.assertEqual(await dev.readline(), b"")
+      with self.assertRaises(TimeoutError):
+        await dev.readline(timeout=0.05)
     self.assertEqual(self._handler.records, [])
     mock_record.assert_not_called()
 
@@ -66,5 +75,181 @@ async def test_nonempty_read_is_logged_and_captured(self) -> None:
     mock_record.assert_called_once()
 
 
+class _BlockingDevice:
+  """A pylibftdi device stand-in whose calls block, standing in for a bulk transfer.
+
+  Args:
+    to_read: bytes the device hands out, oldest first.
+    block_s: how long every call sleeps before returning.
+  """
+
+  def __init__(self, to_read: bytes = b"", block_s: float = 0.0, text_mode: bool = False) -> None:
+    self._to_read = bytearray(to_read)
+    self._block_s = block_s
+    self._text_mode = text_mode
+    self.written = bytearray()
+    self.reads = 0
+    self.ftdi_fn = mock.Mock()
+
+  def read(self, num_bytes: int) -> Union[bytes, str]:
+    data = bytes(self._to_read[:num_bytes])
+    del self._to_read[:num_bytes]
+    self.reads += 1
+    time.sleep(self._block_s)
+    # pylibftdi decodes with latin-1 when the device was opened in text mode.
+    return data.decode("latin-1") if self._text_mode else data
+
+  def write(self, data: bytes) -> int:
+    time.sleep(self._block_s)
+    self.written.extend(data)
+    return len(data)
+
+
+@unittest.skipUnless(HAS_PYLIBFTDI and HAS_PYUSB, "pylibftdi/pyusb not installed")
+class FTDICancellationTests(unittest.IsolatedAsyncioTestCase):
+  """Running read/write on the executor makes them cancellable for the first time, and a call
+  already handed to the worker cannot be recalled. Locks the guarantee that neither end loses
+  data when its caller goes away: a cancelled read's bytes are already out of the chip's buffer,
+  so dropping them would leave the next read part-way through a response."""
+
+  def _ftdi(self, device: _BlockingDevice) -> FTDI:
+    # Bypass setup() (no hardware) by driving the underlying device directly.
+    io = FTDI(human_readable_device_name="test", device_id="test")
+    io._dev = device
+    io._executor = ThreadPoolExecutor(max_workers=1)
+    self.addCleanup(io._executor.shutdown)
+    return io
+
+  async def test_cancelled_read_keeps_its_bytes_for_the_next_read(self) -> None:
+    device = _BlockingDevice(to_read=b"ABCD", block_s=0.2)
+    io = self._ftdi(device)
+
+    with self.assertRaises(asyncio.TimeoutError):
+      await asyncio.wait_for(io.read(4), timeout=0.05)
+
+    self.assertEqual(await io.read(4), b"ABCD")
+    # The second read waited for the one in flight rather than issuing another.
+    self.assertEqual(device.reads, 1)
+
+  async def test_cancelled_read_is_captured_once(self) -> None:
+    device = _BlockingDevice(to_read=b"ABCD", block_s=0.2)
+    io = self._ftdi(device)
+    records: List[Any] = []
+
+    with mock.patch.object(ftdi_module.capturer, "record", records.append):
+      with self.assertRaises(asyncio.TimeoutError):
+        await asyncio.wait_for(io.read(4), timeout=0.05)
+      await io.read(4)
+
+    self.assertEqual([(r.action, r.data) for r in records], [("read", "41424344")])
+
+  async def test_cancelled_write_still_reaches_the_device(self) -> None:
+    device = _BlockingDevice(block_s=0.2)
+    io = self._ftdi(device)
+    records: List[Any] = []
+
+    with mock.patch.object(ftdi_module.capturer, "record", records.append):
+      with self.assertRaises(asyncio.TimeoutError):
+        await asyncio.wait_for(io.write(b"\x01\x02"), timeout=0.05)
+      await asyncio.sleep(0.3)  # let the worker finish
+
+    self.assertEqual(bytes(device.written), b"\x01\x02")
+    self.assertEqual([(r.action, r.data) for r in records], [("write", "0102")])
+
+  async def test_purge_discards_bytes_held_for_the_next_read(self) -> None:
+    device = _BlockingDevice(to_read=b"ABCD", block_s=0.2)
+    io = self._ftdi(device)
+
+    with self.assertRaises(asyncio.TimeoutError):
+      await asyncio.wait_for(io.read(4), timeout=0.05)
+    await io.usb_purge_rx_buffer()
+
+    # A purge means the caller has decided everything in flight is stale.
+    self.assertEqual(await io.read(4), b"")
+
+  async def test_a_text_mode_device_still_reads(self) -> None:
+    """pylibftdi returns str, decoded with latin-1, if the device was opened in text mode. setup()
+    opens it in byte mode, so this only guards against losing the tolerance upstream had."""
+    io = self._ftdi(_BlockingDevice(to_read=b"\xff\x01", text_mode=True))
+
+    self.assertEqual(await io.read(2), b"\xff\x01")
+
+  async def test_second_concurrent_reader_gets_an_empty_read(self) -> None:
+    """Two callers on one io stay unsupported, but the second one now backs off rather than
+    taking bytes out of the middle of the first one's response."""
+    io = self._ftdi(_BlockingDevice(to_read=b"ABCD", block_s=0.01))
+    received: dict = {"first": [], "second": []}
+
+    async def reader(name: str) -> None:
+      for _ in range(4):
+        received[name].append(await io.read(1))
+
+    await asyncio.gather(reader("first"), reader("second"))
+
+    self.assertEqual(received["first"], [b"A", b"B", b"C", b"D"])
+    self.assertEqual(received["second"], [b"", b"", b"", b""])
+
+
+@unittest.skipUnless(HAS_PYLIBFTDI and HAS_PYUSB, "pylibftdi/pyusb not installed")
+class FTDICaptureTests(unittest.IsolatedAsyncioTestCase):
+  """A capture keeps its payload, so FTDIValidator can replay it. FTDICommand declared `data`
+  without assigning it, so every recorded FTDI command was written out without its bytes."""
+
+  def setUp(self) -> None:
+    # CaptureReader.done() leaves the capture/validation flag set, which blocks the next FTDI().
+    self.addCleanup(setattr, capture_module, "_capture_or_validation_active", False)
+
+  def test_command_carries_its_data(self) -> None:
+    # The capture file is written from __dict__, so an unassigned field is a lost payload.
+    command = FTDICommand(device_id="test", action="write", data="dead")
+    self.assertEqual(command.__dict__["data"], "dead")
+
+  async def test_capture_replays_through_the_validator(self) -> None:
+    io = FTDI(human_readable_device_name="test", device_id="test")
+    io._dev = _BlockingDevice(to_read=b"\x03\x04")
+
+    with tempfile.TemporaryDirectory() as directory:
+      path = Path(directory) / "capture.json"
+      capturer.start(path)
+      self.addCleanup(lambda: capturer.capture_active and capturer.stop())
+      await io.set_baudrate(115200)
+      await io.write(b"\x01\x02")
+      await io.read(2)
+      capturer.stop()
+
+      reader = CaptureReader(str(path))
+      validator = FTDIValidator(reader, "test", "test")
+      await validator.set_baudrate(115200)
+      await validator.write(b"\x01\x02")
+      self.assertEqual(await validator.read(2), b"\x03\x04")
+      reader.done()
+
+
+@unittest.skipUnless(HAS_PYLIBFTDI and HAS_PYUSB, "pylibftdi/pyusb not installed")
+class FTDIReadlineTests(unittest.IsolatedAsyncioTestCase):
+  """pylibftdi's readline raises TypeError unless the device was opened in text mode, and setup()
+  opens it in byte mode, so the line is assembled here instead."""
+
+  def _ftdi(self, to_read: bytes) -> FTDI:
+    io = FTDI(human_readable_device_name="test", device_id="test")
+    io._dev = _BlockingDevice(to_read=to_read)
+    return io
+
+  async def test_line_is_returned_with_its_terminator(self) -> None:
+    io = self._ftdi(b"ok\ntrailing")
+    self.assertEqual(await io.readline(), b"ok\n")
+    self.assertEqual(await io.readline(terminator=b"g"), b"trailing")
+
+  async def test_incomplete_line_times_out(self) -> None:
+    io = self._ftdi(b"partial")
+    with self.assertRaises(TimeoutError):
+      await io.readline(timeout=0.05)
+
+  async def test_empty_terminator_is_rejected(self) -> None:
+    io = self._ftdi(b"")
+    with self.assertRaises(ValueError):
+      await io.readline(terminator=b"")
+
+
 if __name__ == "__main__":
   unittest.main()