Improve handling of failed scan tasks and allow reconnect

src/fastcs/attributes/attr_r.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Awaitable, Callable
+from collections.abc import Callable, Coroutine
 from typing import Any
 
 from fastcs.attributes.attribute import Attribute, AttributeAccessMode
@@ -10,11 +10,11 @@
 from fastcs.datatypes import DataType, DType_T
 from fastcs.logging import logger
 
-AttrIOUpdateCallback = Callable[["AttrR[DType_T, Any]"], Awaitable[None]]
+AttrIOUpdateCallback = Callable[["AttrR[DType_T, Any]"], Coroutine[None, None, None]]
 """An AttributeIO callback that takes an AttrR and updates its value"""
-AttrUpdateCallback = Callable[[], Awaitable[None]]
+AttrUpdateCallback = Callable[[], Coroutine[None, None, None]]
 """A callback to be called periodically to update an attribute"""
-AttrOnUpdateCallback = Callable[[DType_T], Awaitable[None]]
+AttrOnUpdateCallback = Callable[[DType_T], Coroutine[None, None, None]]
 """A callback to be called when the value of the attribute is updated"""
 
 
@@ -132,12 +132,8 @@ def bind_update_callback(self) -> AttrUpdateCallback:
             update_callback = self._update_callback
 
         async def update_attribute():
-            try:
-                self.log_event("Update attribute", topic=self)
-                await update_callback(self)
-            except Exception:
-                logger.error("Attribute update loop stopped", attribute=self)
-                raise
+            self.log_event("Update attribute", topic=self)
+            await update_callback(self)
 
         return update_attribute
 

src/fastcs/connections/ip_connection.py

@@ -80,6 +80,13 @@ async def send_query(self, message: str) -> str:
             return response
 
     async def close(self):
+        if self.__connection is None:
+            return
+
         async with self._connection as connection:
-            await connection.close()
-            self.__connection = None
+            try:
+                await connection.close()
+            except ConnectionResetError:
+                pass
+
+        self.__connection = None

src/fastcs/control_system.py

@@ -7,11 +7,11 @@
 
 from IPython.terminal.embed import InteractiveShellEmbed
 
-from fastcs.controllers import BaseController, Controller
+from fastcs.controllers import Controller
 from fastcs.logging import logger
 from fastcs.methods import ScanCallback
 from fastcs.tracer import Tracer
-from fastcs.transports import ControllerAPI, Transport
+from fastcs.transports import Transport
 
 tracer = Tracer()
 
@@ -57,15 +57,6 @@ async def _run_initial_coros(self):
     async def _start_scan_tasks(self):
         self._scan_tasks = {self._loop.create_task(coro()) for coro in self._scan_coros}
 
-        for task in self._scan_tasks:
-            task.add_done_callback(self._scan_done)
-
-    def _scan_done(self, task: asyncio.Task):
-        try:
-            task.result()
-        except Exception:
-            logger.exception("Exception raised in scan task")
-
     def _stop_scan_tasks(self):
         for task in self._scan_tasks:
             if not task.done():
@@ -82,9 +73,8 @@ async def serve(self, interactive: bool = True) -> None:
         await self._controller.initialise()
         self._controller.post_initialise()
 
-        self.controller_api = build_controller_api(self._controller)
-        self._scan_coros, self._initial_coros = (
-            self.controller_api.get_scan_and_initial_coros()
+        self.controller_api, self._scan_coros, self._initial_coros = (
+            self._controller.create_api_and_tasks()
         )
 
         context = {
@@ -168,21 +158,3 @@ async def interactive_shell(
 
     def __del__(self):
         self._stop_scan_tasks()
-
-
-def build_controller_api(controller: Controller) -> ControllerAPI:
-    return _build_controller_api(controller, [])
-
-
-def _build_controller_api(controller: BaseController, path: list[str]) -> ControllerAPI:
-    return ControllerAPI(
-        path=path,
-        attributes=controller.attributes,
-        command_methods=controller.command_methods,
-        scan_methods=controller.scan_methods,
-        sub_apis={
-            name: _build_controller_api(sub_controller, path + [name])
-            for name, sub_controller in controller.sub_controllers.items()
-        },
-        description=controller.description,
-    )

src/fastcs/controllers/__init__.py

@@ -1,3 +1,4 @@
 from .base_controller import BaseController as BaseController
 from .controller import Controller as Controller
+from .controller_api import ControllerAPI as ControllerAPI
 from .controller_vector import ControllerVector as ControllerVector

src/fastcs/controllers/base_controller.py

@@ -3,7 +3,7 @@
 from collections import Counter
 from collections.abc import Sequence
 from copy import deepcopy
-from typing import (  # type: ignore
+from typing import (
     TypeVar,
     _GenericAlias,  # type: ignore
     get_args,
@@ -12,6 +12,7 @@
 )
 
 from fastcs.attributes import AnyAttributeIO, Attribute, AttrR, AttrW, HintedAttribute
+from fastcs.controllers.controller_api import ControllerAPI
 from fastcs.logging import logger
 from fastcs.methods import Command, Method, Scan, UnboundCommand, UnboundScan
 from fastcs.tracer import Tracer
@@ -388,3 +389,16 @@ def add_scan(self, name: str, scan: Scan):
     @property
     def scan_methods(self) -> dict[str, Scan]:
         return self.__scan_methods
+
+    def _build_api(self, path: list[str]) -> ControllerAPI:
+        return ControllerAPI(
+            path=path,
+            attributes=self.attributes,
+            command_methods=self.command_methods,
+            scan_methods=self.scan_methods,
+            sub_apis={
+                name: sub_controller._build_api(path + [name])  # noqa: SLF001
+                for name, sub_controller in self.sub_controllers.items()
+            },
+            description=self.description,
+        )

src/fastcs/controllers/controller.py

@@ -1,7 +1,15 @@
+import asyncio
+from collections import defaultdict
 from collections.abc import Sequence
 
 from fastcs.attributes import AnyAttributeIO
+from fastcs.attributes.attr_r import AttrR
+from fastcs.attributes.attribute_io_ref import AttributeIORef
 from fastcs.controllers.base_controller import BaseController
+from fastcs.controllers.controller_api import ControllerAPI
+from fastcs.logging import logger
+from fastcs.methods import ScanCallback
+from fastcs.util import ONCE
 
 
 class Controller(BaseController):
@@ -13,6 +21,7 @@ def __init__(
         ios: Sequence[AnyAttributeIO] | None = None,
     ) -> None:
         super().__init__(description=description, ios=ios)
+        self._connected = False
 
     def add_sub_controller(self, name: str, sub_controller: BaseController):
         if name.isdigit():
@@ -23,7 +32,105 @@ def add_sub_controller(self, name: str, sub_controller: BaseController):
         return super().add_sub_controller(name, sub_controller)
 
     async def connect(self) -> None:
-        pass
+        """Hook to perform initial connection to device
+
+        This should set ``_connected`` to ``True`` if the connection was successful to
+        enable scan tasks.
+
+        """
+        self._connected = True
+
+    async def reconnect(self):
+        """Hook to reconnect to device after an error
+
+        This should set ``_connected`` to ``True`` if the connection was successful to
+        enable scan tasks.
+
+        If the connection cannot be re-established it should log an error with the
+        reason. It should not raise an exception.
+
+        """
+        self._connected = True
 
     async def disconnect(self) -> None:
+        """Hook to tidy up resources before stopping the application"""
         pass
+
+    def create_api_and_tasks(
+        self,
+    ) -> tuple[ControllerAPI, list[ScanCallback], list[ScanCallback]]:
+        """Create api for transports tasks for FastCS backend
+
+        Creates a tuple of
+            - The `ControllerAPI` for this controller
+            - Initial coroutines to be run once on startup
+            - Periodic coroutines to run as background tasks
+
+        Returns:
+            tuple[ControllerAPI, list[ScanCallback], list[ScanCallback]]
+
+        """
+        controller_api = self._build_api([])
+
+        scan_dict: dict[float, list[ScanCallback]] = defaultdict(list)
+        initial_coros: list[ScanCallback] = []
+
+        for api in controller_api.walk_api():
+            for method in api.scan_methods.values():
+                if method.period is ONCE:
+                    initial_coros.append(method.fn)
+                else:
+                    scan_dict[method.period].append(method.fn)
+
+            for attribute in api.attributes.values():
+                match attribute:
+                    case AttrR(_io_ref=AttributeIORef(update_period=update_period)):
+                        if update_period is ONCE:
+                            initial_coros.append(attribute.bind_update_callback())
+                        elif update_period is not None:
+                            scan_dict[update_period].append(
+                                attribute.bind_update_callback()
+                            )
+
+        periodic_scan_coros: list[ScanCallback] = []
+        for period, methods in scan_dict.items():
+            periodic_scan_coros.append(self._create_periodic_scan_coro(period, methods))
+
+        return controller_api, periodic_scan_coros, initial_coros
+
+    def _create_periodic_scan_coro(
+        self, period: float, scans: Sequence[ScanCallback]
+    ) -> ScanCallback:
+        """Create a coroutine to run scans at a given period
+
+        This returns a coroutine that runs scans at a given period. If an exception is
+        raised in a callback it is caught and the updates for the controller are
+        paused, waiting for `_connected` to be set back to true via the `reconnect`
+        method.
+
+        Args:
+            period: The period to run the scans at
+            scans: A list of `ScanCallback` to run periodically
+
+        Returns:
+            A wrapper `ScanCallback` that runs all of the callbacks at a given period
+        """
+
+        async def scan_coro() -> None:
+            while True:
+                if not self._connected:
+                    await asyncio.sleep(1)
+                    continue
+
+                try:
+                    await asyncio.gather(
+                        asyncio.sleep(period), *[scan() for scan in scans]
+                    )
+                except Exception:
+                    logger.exception("Exception in scan task", period=period)
+                    self._connected = False
+
+                    await asyncio.sleep(1)  # Wait so this message appears last
+                    logger.error("Pausing scan tasks and waiting for reconnect")
+
+        return scan_coro

src/fastcs/controllers/controller_api.py

@@ -0,0 +1,38 @@
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+
+from fastcs.attributes import Attribute
+from fastcs.methods import Command, Scan
+
+
+@dataclass
+class ControllerAPI:
+    """Attributes, Methods and sub APIs of a `Controller` to expose in a transport"""
+
+    path: list[str] = field(default_factory=list)
+    """Path within controller tree (empty if this is the root)"""
+    attributes: dict[str, Attribute] = field(default_factory=dict)
+    command_methods: dict[str, Command] = field(default_factory=dict)
+    scan_methods: dict[str, Scan] = field(default_factory=dict)
+    sub_apis: dict[str, "ControllerAPI"] = field(default_factory=dict)
+    """APIs of the sub controllers of the `Controller` this API was built from"""
+    description: str | None = None
+
+    def walk_api(self) -> Iterator["ControllerAPI"]:
+        """Walk through all the nested `ControllerAPI` s of this `ControllerAPI`.
+
+        Yields the `ControllerAPI` s from a depth-first traversal of the tree,
+        including self.
+
+        """
+        yield self
+        for api in self.sub_apis.values():
+            yield from api.walk_api()
+
+    def __repr__(self):
+        return (
+            f"ControllerAPI("
+            f"path={self.path}, "
+            f"sub_apis=[{', '.join(self.sub_apis.keys())}]"
+            f")"
+        )

src/fastcs/demo/controllers.py

@@ -95,6 +95,12 @@ async def cancel_all(self) -> None:
     async def connect(self) -> None:
         await self.connection.connect(self._settings.ip_settings)
 
+    async def reconnect(self):
+        await self.connection.close()
+        await self.connection.connect(self._settings.ip_settings)
+
+        self._connected = True
+
     async def close(self) -> None:
         await self.connection.close()
 

src/fastcs/transports/__init__.py

@@ -1,4 +1,3 @@
-from .controller_api import ControllerAPI as ControllerAPI
 from .transport import Transport as Transport
 
 try: