fix(robot-server): Blink the OT-2's button light while the persistenc…

…e layer initializes (#14388)

robot-server/robot_server/app_setup.py

@@ -10,7 +10,15 @@
 from opentrons import __version__
 
 from .errors import exception_handlers
-from .hardware import start_initializing_hardware, clean_up_hardware
+from .hardware import (
+    fbl_init,
+    fbl_mark_hardware_init_complete,
+    fbl_mark_persistence_init_complete,
+    start_initializing_hardware,
+    clean_up_hardware,
+    fbl_start_blinking,
+    fbl_clean_up,
+)
 from .persistence import start_initializing_persistence, clean_up_persistence
 from .router import router
 from .service import initialize_logging
@@ -71,15 +79,27 @@ async def on_startup() -> None:
 
     initialize_logging()
     initialize_task_runner(app_state=app.state)
+    fbl_init(app_state=app.state)
     start_initializing_hardware(
         app_state=app.state,
         callbacks=[
+            # Flex light control:
             (start_light_control_task, True),
             (mark_light_control_startup_finished, False),
+            # OT-2 light control:
+            (fbl_start_blinking, True),
+            (fbl_mark_hardware_init_complete, False),
         ],
     )
     start_initializing_persistence(
-        app_state=app.state, persistence_directory_root=persistence_directory
+        app_state=app.state,
+        persistence_directory_root=persistence_directory,
+        done_callbacks=[
+            # For OT-2 light control only. The Flex status bar isn't handled here
+            # because it's currently tied to hardware and run status, not to
+            # initialization of the persistence layer.
+            fbl_mark_persistence_init_complete
+        ],
     )
     initialize_notification_client(
         app_state=app.state,
@@ -89,7 +109,12 @@ async def on_startup() -> None:
 @app.on_event("shutdown")
 async def on_shutdown() -> None:
     """Handle app shutdown."""
+    # FIXME(mm, 2024-01-31): Cleaning up everything concurrently like this is prone to
+    # race conditions, e.g if we clean up hardware before we clean up the background
+    # task that's blinking the front button light (which uses the hardware).
+    # Startup and shutdown should be in FILO order.
     shutdown_results = await asyncio.gather(
+        fbl_clean_up(app.state),
         clean_up_hardware(app.state),
         clean_up_persistence(app.state),
         clean_up_task_runner(app.state),

robot-server/robot_server/hardware.py

@@ -4,17 +4,18 @@
 from pathlib import Path
 from fastapi import Depends, status
 from typing import (
-    Callable,
     TYPE_CHECKING,
     cast,
     Awaitable,
+    Callable,
     Iterator,
     Iterable,
+    Optional,
     Tuple,
 )
 from uuid import uuid4  # direct to avoid import cycles in service.dependencies
 from traceback import format_exception_only, TracebackException
-from contextlib import contextmanager
+from contextlib import contextmanager, suppress
 
 from opentrons_shared_data import deck
 from opentrons_shared_data.robot.dev_types import RobotType, RobotTypeEnum
@@ -72,6 +73,9 @@
 
 _hw_api_accessor = AppStateAccessor[ThreadManagedHardware]("hardware_api")
 _init_task_accessor = AppStateAccessor["asyncio.Task[None]"]("hardware_init_task")
+_front_button_light_blinker_accessor = AppStateAccessor["_FrontButtonLightBlinker"](
+    "front_button_light_blinker"
+)
 _postinit_task_accessor = AppStateAccessor["asyncio.Task[None]"](
     "hardware_postinit_task"
 )
@@ -127,6 +131,113 @@ async def clean_up_hardware(app_state: AppState) -> None:
         thread_manager.clean_up()
 
 
+# TODO(mm, 2024-01-30): Consider merging this with the Flex's LightController.
+class _FrontButtonLightBlinker:
+    def __init__(self) -> None:
+        self._hardware_and_task: Optional[
+            Tuple[HardwareControlAPI, "asyncio.Task[None]"]
+        ] = None
+        self._hardware_init_complete = False
+        self._persistence_init_complete = False
+
+    async def set_hardware(self, hardware: HardwareControlAPI) -> None:
+        assert self._hardware_and_task is None, "hardware should only be set once."
+
+        async def blink_forever() -> None:
+            while True:
+                await hardware.set_lights(button=True)
+                await asyncio.sleep(0.5)
+                await hardware.set_lights(button=False)
+                await asyncio.sleep(0.5)
+
+        task = asyncio.create_task(blink_forever())
+
+        self._hardware_and_task = (hardware, task)
+
+    async def mark_hardware_init_complete(self) -> None:
+        self._hardware_init_complete = True
+        await self._maybe_stop_blinking()
+
+    async def mark_persistence_init_complete(self) -> None:
+        self._persistence_init_complete = True
+        await self._maybe_stop_blinking()
+
+    async def clean_up(self) -> None:
+        if self._hardware_and_task is not None:
+            _, task = self._hardware_and_task
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
+
+    async def _maybe_stop_blinking(self) -> None:
+        if self._hardware_and_task is not None and self._all_complete():
+            # We're currently blinking, but we should stop.
+            hardware, task = self._hardware_and_task
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
+            await hardware.set_lights(button=True)
+
+    def _all_complete(self) -> bool:
+        return self._persistence_init_complete and self._hardware_init_complete
+
+
+def fbl_init(app_state: AppState) -> None:
+    """Prepare to blink the OT-2's front button light.
+
+    This should be called once during server startup.
+    """
+    if should_use_ot3():
+        # This is only for the OT-2's front button light.
+        # The Flex's status bar is handled elsewhere -- see LightController.
+        return
+    _front_button_light_blinker_accessor.set_on(app_state, _FrontButtonLightBlinker())
+
+
+async def fbl_start_blinking(app_state: AppState, hardware: HardwareControlAPI) -> None:
+    """Start blinking the OT-2's front button light.
+
+    This should be called once during server startup, as soon as the hardware is
+    initialized enough to support the front button light.
+
+    Note that this is preceded by two other visually indistinguishable stages of
+    blinking:
+    1. A separate system process blinks the light while this process's Python
+       interpreter is initializing.
+    2. build_hardware_controller() blinks the light internally while it's doing hardware
+       initialization.
+
+    Blinking will continue until `fbl_mark_hardware_init_complete()` and
+    `fbl_mark_persistence_init_complete()` have both been called.
+    """
+    blinker = _front_button_light_blinker_accessor.get_from(app_state)
+    if blinker is not None:  # May be None on a Flex.
+        await blinker.set_hardware(hardware)
+
+
+async def fbl_mark_hardware_init_complete(
+    app_state: AppState, hardware: HardwareControlAPI
+) -> None:
+    """See `fbl_start_blinking()`."""
+    blinker = _front_button_light_blinker_accessor.get_from(app_state)
+    if blinker is not None:  # May be None on a Flex.
+        await blinker.mark_hardware_init_complete()
+
+
+async def fbl_mark_persistence_init_complete(app_state: AppState) -> None:
+    """See `fbl_start_blinking()`."""
+    blinker = _front_button_light_blinker_accessor.get_from(app_state)
+    if blinker is not None:  # May be None on a Flex.
+        await blinker.mark_persistence_init_complete()
+
+
+async def fbl_clean_up(app_state: AppState) -> None:
+    """Clean up the background task that blinks the OT-2's front button light."""
+    blinker = _front_button_light_blinker_accessor.get_from(app_state)
+    if blinker is not None:
+        await blinker.clean_up()
+
+
 # TODO(mm, 2022-10-18): Deduplicate this background initialization infrastructure
 # with similar code used for initializing the persistence layer.
 async def get_thread_manager(
@@ -281,28 +392,9 @@ async def _postinit_ot2_tasks(
     callbacks: Iterable[PostInitCallback],
 ) -> None:
     """Tasks to run on an initialized OT-2 before it is ready to use."""
-
-    async def _blink() -> None:
-        while True:
-            await hardware.set_lights(button=True)
-            await asyncio.sleep(0.5)
-            await hardware.set_lights(button=False)
-            await asyncio.sleep(0.5)
-
-    # While the hardware was initializing in _create_hardware_api(), it blinked the
-    # front button light. But that blinking stops when the completed hardware object
-    # is returned. Do our own blinking here to keep it going while we home the robot.
-    blink_task = asyncio.create_task(_blink())
-
     try:
         await _home_on_boot(hardware.wrapped())
-        await hardware.set_lights(button=True)
     finally:
-        blink_task.cancel()
-        try:
-            await blink_task
-        except asyncio.CancelledError:
-            pass
         for callback in callbacks:
             if not callback[1]:
                 await callback[0](app_state, hardware.wrapped())
@@ -324,7 +416,6 @@ async def _home_on_boot(hardware: HardwareControlAPI) -> None:
 async def _do_updates(
     hardware: "OT3API", update_manager: FirmwareUpdateManager
 ) -> None:
-
     update_handles = [
         await update_manager.start_update_process(
             str(uuid4()), SubSystem.from_hw(subsystem), utc_now()

robot-server/robot_server/persistence/_fastapi_dependencies.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from pathlib import Path
-from typing import Optional
+from typing import Awaitable, Callable, Iterable, Optional
 from typing_extensions import Literal
 
 from sqlalchemy.engine import Engine as SQLEngine
@@ -57,7 +57,9 @@ class DatabaseFailedToInitialize(ErrorDetails):
 
 
 def start_initializing_persistence(  # noqa: C901
-    app_state: AppState, persistence_directory_root: Optional[Path]
+    app_state: AppState,
+    persistence_directory_root: Optional[Path],
+    done_callbacks: Iterable[Callable[[AppState], Awaitable[None]]],
 ) -> None:
     """Initialize the persistence layer to get it ready for use by endpoint functions.
 
@@ -136,6 +138,15 @@ async def init_sql_engine() -> SQLEngine:
         app_state=app_state, value=sql_engine_init_task
     )
 
+    async def wait_until_done_then_trigger_callbacks() -> None:
+        try:
+            await sql_engine_init_task
+        finally:
+            for callback in done_callbacks:
+                await callback(app_state)
+
+    asyncio.create_task(wait_until_done_then_trigger_callbacks())
+
 
 async def clean_up_persistence(app_state: AppState) -> None:
     """Clean up the persistence layer.

robot-server/robot_server/runs/light_control_task.py

@@ -1,4 +1,4 @@
-"""Background task to drive the status bar."""
+"""Background task to drive the Flex's status bar."""
 from typing import Optional, List
 from logging import getLogger
 import asyncio
@@ -81,7 +81,7 @@ def _active_updates_to_status_bar(
 
 
 class LightController:
-    """LightController sets the status bar to match the protocol status."""
+    """LightController sets the Flex's status bar to match the protocol status."""
 
     def __init__(
         self, api: HardwareControlAPI, engine_store: Optional[EngineStore]

robot-server/robot_server/service/task_runner.py

@@ -73,7 +73,7 @@ async def cancel_all_and_clean_up(self) -> None:
 def initialize_task_runner(app_state: AppState) -> None:
     """Create a new `TaskRunner` and store it on `app_state`
 
-    Intended to be called just once, when the server starts up.s
+    Intended to be called just once, when the server starts up.
     """
     _task_runner_accessor.set_on(app_state, TaskRunner())