STAR: v1 driver, first wave

pylabrobot/hamilton/__init__.py

@@ -0,0 +1,5 @@
+"""Hamilton devices."""
+
+from .star.device import STAR, STARDevice, STARLet, STARPlus
+from .star.driver.master import STARDriver
+from .star.driver.simulator import STARSimulationDriver

pylabrobot/hamilton/protocol/text/router.py

@@ -63,13 +63,16 @@ def __init__(
     self._parse_id = parse_id
     self._raise_for_error = raise_for_error
 
+    # TODO: not used by the reader yet.
     self.packet_read_timeout = packet_read_timeout
     self.read_timeout = read_timeout
 
     self.id_ = 0
     self._reading_thread: Optional[threading.Thread] = None
     self._reading_thread_stop = threading.Event()
+    # The reading thread and the event loop both change this, so every change goes through the lock.
     self._waiting_tasks: List[HamiltonTask] = []
+    self._waiting_tasks_lock = threading.Lock()
 
   def start(self) -> None:
     """Begin reading replies. The caller opens the transport first."""
@@ -83,11 +86,12 @@ def stop(self) -> None:
     if self._reading_thread is not None:
       self._reading_thread.join(timeout=10)
       self._reading_thread = None
-    for task in self._waiting_tasks:
+    with self._waiting_tasks_lock:
+      waiting, self._waiting_tasks = self._waiting_tasks, []
+    for task in waiting:
       task.loop.call_soon_threadsafe(
         task.fut.set_exception, RuntimeError("Stopping the reply router.")
       )
-    self._waiting_tasks.clear()
 
   def next_id(self) -> int:
     """continuously generate unique ids 0 <= x < 10000."""
@@ -115,17 +119,28 @@ async def send(
     Returns:
       The reply, or None when `wait` is False.
     """
-    await self.io.write(cmd.encode(), timeout=write_timeout)
-
     if not wait:
+      await self.io.write(cmd.encode(), timeout=write_timeout)
       return None
 
     if read_timeout is None:
       read_timeout = self.read_timeout
 
+    # Wait for the reply before asking for it. The reading thread runs alongside this one, and a
+    # device can answer before a task registered after the write would exist - the reply then
+    # arrives with nothing waiting for it, is dropped, and the command times out with its answer
+    # already gone past. Registering first closes that window; the task is taken back off again if
+    # the write never happens.
     loop = asyncio.get_event_loop()
     fut: asyncio.Future[str] = loop.create_future()
-    self._start_reading(id_, loop, fut, cmd, read_timeout)
+    task = self._start_reading(id_, loop, fut, cmd, read_timeout)
+    try:
+      await self.io.write(cmd.encode(), timeout=write_timeout)
+    except BaseException:
+      with self._waiting_tasks_lock:
+        if task in self._waiting_tasks:
+          self._waiting_tasks.remove(task)
+      raise
     return await fut
 
   async def send_raw(
@@ -151,19 +166,25 @@ def _start_reading(
     fut: asyncio.Future,
     cmd: str,
     timeout: int,
-  ) -> None:
-    """Submit a task to the reading thread."""
+  ) -> HamiltonTask:
+    """Submit a task to the reading thread, and hand it back so a caller can take it off again.
+
+    Returns:
+      The task now waiting for a reply.
+    """
 
     timeout_time = time.time() + timeout
-    self._waiting_tasks.append(
-      HamiltonTask(id_=id_, loop=loop, fut=fut, cmd=cmd, timeout_time=timeout_time)
-    )
+    task = HamiltonTask(id_=id_, loop=loop, fut=fut, cmd=cmd, timeout_time=timeout_time)
+    with self._waiting_tasks_lock:
+      self._waiting_tasks.append(task)
 
     if self._reading_thread is None or not self._reading_thread.is_alive():
       self._reading_thread_stop.clear()
       self._reading_thread = threading.Thread(target=self._reading_thread_main, daemon=True)
       self._reading_thread.start()
 
+    return task
+
   def _reading_thread_main(self) -> None:
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
@@ -182,17 +203,20 @@ async def _continuously_read(self) -> None:
     """
 
     while not self._reading_thread_stop.is_set():
-      for idx in range(len(self._waiting_tasks) - 1, -1, -1):  # reverse order to allow deletion
-        task = self._waiting_tasks[idx]
-        if time.time() > task.timeout_time:
-          logger.warning("Timeout while waiting for response to command %s.", task.cmd)
-          task.loop.call_soon_threadsafe(
-            task.fut.set_exception,
-            TimeoutError(f"Timeout while waiting for response to command {task.cmd}."),
-          )
-          del self._waiting_tasks[idx]
-
-      if len(self._waiting_tasks) == 0:
+      with self._waiting_tasks_lock:
+        now = time.time()
+        timed_out = [task for task in self._waiting_tasks if now > task.timeout_time]
+        for task in timed_out:
+          self._waiting_tasks.remove(task)
+        waiting = len(self._waiting_tasks)
+      for task in timed_out:
+        logger.warning("Timeout while waiting for response to command %s.", task.cmd)
+        task.loop.call_soon_threadsafe(
+          task.fut.set_exception,
+          TimeoutError(f"Timeout while waiting for response to command {task.cmd}."),
+        )
+
+      if waiting == 0:
         await asyncio.sleep(0.01)
         continue
 
@@ -212,17 +236,30 @@ async def _continuously_read(self) -> None:
         continue
 
       module_and_command = resp[: self.module_id_length + 2]
-      for idx in range(len(self._waiting_tasks)):
-        task = self._waiting_tasks[idx]
-        # if the command has no id, we have to check the command itself
-        if response_id == task.id_ or (
-          task.id_ is None and task.cmd.startswith(module_and_command)
-        ):
-          try:
-            self._raise_for_error(resp)
-          except Exception as e:
-            task.loop.call_soon_threadsafe(task.fut.set_exception, e)
-          else:
-            task.loop.call_soon_threadsafe(task.fut.set_result, resp)
-          del self._waiting_tasks[idx]
-          break
+      with self._waiting_tasks_lock:
+        matched = next(
+          (
+            task
+            for task in self._waiting_tasks
+            # if the command has no id, we have to check the command itself
+            if response_id == task.id_
+            or (task.id_ is None and task.cmd.startswith(module_and_command))
+          ),
+          None,
+        )
+        if matched is not None:
+          self._waiting_tasks.remove(matched)
+
+      if matched is None:
+        # Deliberately not guessed at. Handing it to whichever outstanding command shares its
+        # module and code answers that command with something that was never its answer, and a
+        # reply arriving with nothing waiting for it is worth seeing rather than repairing.
+        logger.warning("nothing was waiting for this reply, and it was dropped: %s", resp)
+        continue
+
+      try:
+        self._raise_for_error(resp)
+      except Exception as e:
+        matched.loop.call_soon_threadsafe(matched.fut.set_exception, e)
+      else:
+        matched.loop.call_soon_threadsafe(matched.fut.set_result, resp)

pylabrobot/hamilton/protocol/text/router_tests.py

@@ -0,0 +1,108 @@
+import asyncio
+import queue
+import unittest
+from typing import List, Optional
+
+from pylabrobot.hamilton.protocol.text.framing import read_id
+from pylabrobot.hamilton.protocol.text.router import ReplyRouter
+
+
+class FakeTransport:
+  """A link whose replies are queued by the test, and whose writes can be made to fail."""
+
+  def __init__(self, fail_writes: bool = False):
+    self.replies: "queue.Queue[bytes]" = queue.Queue()
+    self.written: List[bytes] = []
+    self.fail_writes = fail_writes
+
+  async def write(self, data: bytes, timeout: Optional[int] = None) -> None:
+    if self.fail_writes:
+      raise ConnectionError("the write did not go out")
+    self.written.append(data)
+
+  async def read(self) -> bytes:
+    try:
+      return self.replies.get(timeout=0.01)
+    except queue.Empty:
+      raise TimeoutError from None
+
+
+def router(transport: FakeTransport, read_timeout: int = 5) -> ReplyRouter:
+  def raise_for_error(reply: str) -> None:
+    if "er99" in reply:
+      raise RuntimeError(reply)
+
+  return ReplyRouter(
+    io=transport,  # type: ignore[arg-type]
+    module_id_length=2,
+    parse_id=read_id,
+    raise_for_error=raise_for_error,
+    read_timeout=read_timeout,
+  )
+
+
+class TestReplyRouter(unittest.IsolatedAsyncioTestCase):
+  """Each reply reaches the command waiting for it, and nothing is left waiting that cannot be."""
+
+  async def test_a_reply_reaches_the_command_with_its_id(self):
+    transport = FakeTransport()
+    r = router(transport)
+    r.start()
+    try:
+      first = asyncio.ensure_future(r.send("C0RTid0001", id_=1))
+      second = asyncio.ensure_future(r.send("C0RWid0002", id_=2))
+      await asyncio.sleep(0.05)
+      transport.replies.put(b"C0RWid0002er00/00rw000")
+      transport.replies.put(b"C0RTid0001er00/00rt0")
+      self.assertEqual(await asyncio.wait_for(first, 2), "C0RTid0001er00/00rt0")
+      self.assertEqual(await asyncio.wait_for(second, 2), "C0RWid0002er00/00rw000")
+      self.assertEqual(r._waiting_tasks, [])
+    finally:
+      r.stop()
+
+  async def test_an_error_reply_raises_on_its_command(self):
+    transport = FakeTransport()
+    r = router(transport)
+    r.start()
+    try:
+      sent = asyncio.ensure_future(r.send("C0ZAid0003", id_=3))
+      await asyncio.sleep(0.05)
+      transport.replies.put(b"C0ZAid0003er99/00")
+      with self.assertRaises(RuntimeError):
+        await asyncio.wait_for(sent, 2)
+    finally:
+      r.stop()
+
+  async def test_an_unanswered_command_times_out_and_is_taken_off(self):
+    r = router(FakeTransport(), read_timeout=0)
+    r.start()
+    try:
+      with self.assertRaises(TimeoutError):
+        await asyncio.wait_for(r.send("C0QWid0004", id_=4), 2)
+      self.assertEqual(r._waiting_tasks, [])
+    finally:
+      r.stop()
+
+  async def test_a_failed_write_leaves_nothing_waiting(self):
+    r = router(FakeTransport(fail_writes=True))
+    r.start()
+    try:
+      with self.assertRaises(ConnectionError):
+        await r.send("C0QWid0005", id_=5)
+      self.assertEqual(r._waiting_tasks, [])
+    finally:
+      r.stop()
+
+  async def test_stop_fails_every_command_still_waiting(self):
+    r = router(FakeTransport())
+    r.start()
+    sent = asyncio.ensure_future(r.send("C0QWid0006", id_=6))
+    await asyncio.sleep(0.05)
+    r.stop()
+    with self.assertRaises(RuntimeError):
+      await asyncio.wait_for(sent, 2)
+    self.assertEqual(r._waiting_tasks, [])
+
+
+if __name__ == "__main__":
+  unittest.main()

pylabrobot/hamilton/star/conftest.py

@@ -0,0 +1,19 @@
+"""Fixtures the STAR tests share.
+
+Here rather than beside the driver because nothing the package ships uses them: they exist to
+build devices to test against, not to describe a device anything runs on.
+"""
+
+from pylabrobot.hamilton.star.driver.features.x_arm import XArmConfiguration
+
+# An arm that carries nothing: the geometry of a STAR arm with none of its feature bits set. The
+# firmware requires the two drives' bits to be disjoint, so a device with two arms has its features
+# on one of them and an arm like this as the other. Its firmware is the arm's, since it stands in
+# for a real one.
+BARE_X_ARM = XArmConfiguration(
+  firmware_version="1.4S 2012-04-25",
+  width=354.0,
+  x_range=(95.0, 1340.2),
+  workspace_x_range=(-323.2, 1517.2),
+  wrap_size=595.2,
+)