From cdd208618ed5cbdc45ecbd7715220592e7d2acbe Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Wed, 5 Aug 2026 13:16:11 +0200 Subject: [PATCH] refactor(devices): vendor term/trail socket controller from ophyd_devices --- .../devices/utils/term_trail_controller.py | 270 ++++++++++++++++++ .../test_term_trail_controller.py | 150 ++++++++++ 2 files changed, 420 insertions(+) create mode 100644 debye_bec/devices/utils/term_trail_controller.py create mode 100644 tests/tests_devices/test_term_trail_controller.py diff --git a/debye_bec/devices/utils/term_trail_controller.py b/debye_bec/devices/utils/term_trail_controller.py new file mode 100644 index 0000000..fcbc50c --- /dev/null +++ b/debye_bec/devices/utils/term_trail_controller.py @@ -0,0 +1,270 @@ +"""Socket controller base with a configurable wire protocol. + +On top of the stock ``Controller`` this adds: + - configurable outgoing line termination (``term``) and reply terminators + (``trail``, a single string or several alternatives), + - a buffered ``socket_get`` that reassembles replies split across TCP reads and + rejects runaway replies (``max_reply_length``), + - a per-controller ``socket_timeout`` for individual send/recv operations, + - reconnect-before-retry on communication errors, discarding stale replies. +""" + +import functools +import time +import traceback + +from bec_lib import bec_logger +from ophyd_devices.utils.controller import Controller, ControllerCommunicationError, threadlocked + +logger = bec_logger.logger + + +def retry_once_reconnected(fcn): + """Decorator to rerun a function once if a communication error was raised. + + Reconnects first to discard any stale/desynced reply left on the wire from the + failed attempt -- without this, a late-arriving reply to the *first* attempt could + be misread as the reply to the retry. + """ + + @functools.wraps(fcn) + def wrapper(self, *args, **kwargs): + try: + val = fcn(self, *args, **kwargs) + except Exception: + content = traceback.format_exc() + logger.warning( + f"Communication error occurred. Reconnecting and retrying the command. Traceback: {content}" + ) + self._reconnect() + val = fcn(self, *args, **kwargs) + return val + + return wrapper + + +class TermTrailController(Controller): + """Socket controller with configurable term/trail strings and buffered replies. + + Args: + term (str, optional): Termination string appended to each outgoing socket request. + Defaults to the class attribute _term ("\\n"). + trail (str | list[str], optional): Termination string(s) marking the end of a + socket reply. Defaults to the class attribute _trail ("\\r\\n"). + max_reply_length (int, optional): Max accepted length of a socket reply in bytes. + Defaults to 1024 bytes. + socket_timeout (int | float, optional): Timeout for each socket operation in + seconds. Defaults to 2 seconds. + + Subclasses that need a different wire protocol should override the _term and _trail + class attributes. The constructor arguments only take effect on the first + instantiation per host:port; the controller is a singleton and later conflicting + values are ignored with a warning. + """ + + _term = "\n" # termination string appended to each outgoing request + _trail = "\r\n" # termination string(s) stripped from the end of each reply + _max_reply_length = 1024 # max accepted length of a socket reply in bytes + _socket_timeout = 2 # timeout for each socket operation in seconds + + def __init__( + self, + *, + socket_cls, + socket_host, + socket_port, + device_manager, + name: str = "", + attr_name="", + parent=None, + labels=None, + kind=None, + term: str | None = None, + trail: str | list[str] | tuple[str, ...] | None = None, + max_reply_length: int | None = None, + socket_timeout: int | float | None = None, + ): + if term is not None and not isinstance(term, str): + raise TypeError(f"term must be a string, got {type(term).__name__}") + if trail is not None: + if isinstance(trail, str): + pass + elif isinstance(trail, (list, tuple)) and all(isinstance(t, str) for t in trail): + if not trail: + raise ValueError("trail must not be an empty list/tuple.") + else: + raise TypeError( + f"trail must be a string or a list/tuple of strings, got {type(trail).__name__}" + ) + if max_reply_length is not None and not isinstance(max_reply_length, int): + raise TypeError( + f"max_reply_length must be an int, got {type(max_reply_length).__name__}" + ) + if socket_timeout is not None and not isinstance(socket_timeout, (int, float)): + raise TypeError(f"socket_timeout must be a number, got {type(socket_timeout).__name__}") + + first_init = not self._initialized + super().__init__( + socket_cls=socket_cls, + socket_host=socket_host, + socket_port=socket_port, + device_manager=device_manager, + name=name, + attr_name=attr_name, + parent=parent, + labels=labels, + kind=kind, + ) + if first_init: + if term is not None: + self._term = term + if trail is not None: + self._trail = trail + self._trail_options: tuple[str, ...] = ( + (self._trail,) if isinstance(self._trail, str) else tuple(self._trail) + ) + if max_reply_length is not None: + self._max_reply_length = max_reply_length + if socket_timeout is not None: + self._socket_timeout = socket_timeout + elif ( + (term is not None and term != self._term) + or (trail is not None and trail != self._trail) + or (max_reply_length is not None and max_reply_length != self._max_reply_length) + or (socket_timeout is not None and socket_timeout != self._socket_timeout) + ): + logger.warning( + f"Controller {self._socket_host}:{self._socket_port} is already initialized with " + f"term={self._term!r}, trail={self._trail!r}, max_reply_length={self._max_reply_length!r}, " + f"socket_timeout={self._socket_timeout!r}; ignoring conflicting values " + f"term={term!r}, trail={trail!r}, max_reply_length={max_reply_length!r}, " + f"socket_timeout={socket_timeout!r}." + ) + + @threadlocked + def _reconnect(self): + """ + Close and reopen the socket connection. + + Required after any communication error, in particular a recv() timeout: + this protocol has no per-message IDs, so there is no way to know whether + bytes that show up on the wire *after* a timeout belong to the request that + timed out or to whatever is sent next. A stale reply arriving late would + otherwise be read as the answer to a new command (or concatenated with it). + Closing and reopening the TCP connection discards any such reply-in-flight, + so the next command starts from a guaranteed-clean slate. + """ + try: + if self.sock is not None: + self.sock.close() + except Exception: + logger.warning("Error closing socket during reconnect.", exc_info=True) + finally: + self.sock = None + self.connected = False + self.on() + + @threadlocked + def socket_put(self, val: str): + """ + Send a command to the controller through the socket. + + Args: + val (str): Command to send + """ + self.command_history.append(f"[PUT]: time:{time.time()}, cmd:{val + self._term}") + self.sock.put(f"{val}{self._term}".encode()) + + @threadlocked + def socket_get(self): + """ + Receive a single, complete reply from the controller. + + Loops on `recv()` until any one of `self._trail_options` is seen, since a + reply can arrive split across multiple TCP reads, and different commands on + the same controller (e.g. ACS SETVAR vs GETVAR) can use different terminators. + `self._max_reply_length` guards against a malformed/runaway reply with no + matching trail. Does not protect against stale replies from a prior + timed-out request; that's handled by reconnecting the socket on + communication errors (see `_reconnect`). + + Returns: + str: The decoded reply, including its trailing terminator. + + Raises: + ControllerCommunicationError: If the connection closes while waiting + for a reply, or the reply exceeds `self._max_reply_length`. + """ + buf = b"" + while True: + for trail in self._trail_options: + if trail.encode() in buf: + response = buf.decode() + self.command_history.append(f"[GET]: time:{time.time()}, rep:{response}") + return response + + chunk = self.sock.receive() + if not chunk: + raise ControllerCommunicationError( + "Socket connection closed by remote host while waiting for a reply." + ) + buf += chunk + if len(buf) > self._max_reply_length: + raise ControllerCommunicationError( + f"Reply exceeded max_reply_length ({self._max_reply_length} bytes): {buf!r}" + ) + + @retry_once_reconnected + @threadlocked + def socket_put_and_receive(self, val: str, remove_trailing_chars=True) -> str: + """ + Send a command to the controller and receive the response. + Override this method in the derived class if necessary, especially if the response + needs to be parsed differently. + """ + try: + self.socket_put(val) + if remove_trailing_chars: + return self._remove_trailing_characters(self.socket_get()) + return self.socket_get() + except Exception as exc: + logger.error( + f"Error in socket_put_and_receive: {exc}. Command history: {list(self.command_history)}" + ) + raise ControllerCommunicationError( + f"Failed to communicate with the controller. The last {self._command_history_length} commands were: " + f"{list(self.command_history)}" + ) from exc + + def _remove_trailing_characters(self, var) -> str: + """Strip whichever configured trail terminator is present at the end of a + reply; mid-reply occurrences are kept.""" + for trail in self._trail_options: + if var.endswith(trail): + return var.removesuffix(trail) + return var + + def on(self, timeout: int = 10) -> None: + """ + Open a new socket connection to the controller + + Args: + timeout (int): Time in seconds to wait for the connection itself to + be established (passed to `SocketIO.open`). This is separate from + `self._socket_timeout`, which governs how long each individual + send/recv call is allowed to take once connected. + """ + if not self.connected or self.sock is None: + try: + self.sock = self._socket_cls( + host=self._socket_host, + port=self._socket_port, + socket_timeout=self._socket_timeout, + ) + except TypeError: + # socket classes without a socket_timeout parameter (e.g. test mocks) + self.sock = self._socket_cls(host=self._socket_host, port=self._socket_port) + self.sock.open(timeout=timeout) + self.connected = True + else: + logger.info("The connection has already been established.") diff --git a/tests/tests_devices/test_term_trail_controller.py b/tests/tests_devices/test_term_trail_controller.py new file mode 100644 index 0000000..b40def4 --- /dev/null +++ b/tests/tests_devices/test_term_trail_controller.py @@ -0,0 +1,150 @@ +"""Tests for the vendored TermTrailController (term/trail wire-protocol support).""" + +from unittest import mock + +import pytest +from bec_server.device_server.tests.utils import DMMock +from ophyd_devices.tests.utils import SocketMock +from ophyd_devices.utils.controller import Controller, ControllerCommunicationError + +from debye_bec.devices.utils.term_trail_controller import TermTrailController + +# pylint: disable=protected-access + + +@pytest.fixture +def make_controller(): + """Factory fixture to build a fresh controller singleton with custom term/trail settings.""" + + def _make(controller_cls=TermTrailController, **kwargs): + Controller._reset_controller() + controller = controller_cls( + name="controller", + socket_cls=SocketMock, + socket_host="localhost", + socket_port=8080, + device_manager=DMMock(), + **kwargs, + ) + controller.on() + return controller + + yield _make + Controller._reset_controller() + + +def test_socket_put_appends_custom_term(make_controller): + controller = make_controller(term="\r") + controller.socket_put("get") + assert controller.sock.buffer_put == [b"get\r"] + + +@pytest.mark.parametrize( + ["trail", "reply", "expected"], + [ + (None, "value\r\n", "value"), + (None, "\r\n", ""), + (None, "line1\r\nline2\r\n", "line1\r\nline2"), + ("\n", "value\n", "value"), + ("\n", "\n", ""), + ("\n", "line1\nline2\n", "line1\nline2"), + ("", "value\r\n", "value\r\n"), + ], +) +def test_remove_trailing_characters_strips_suffix_only(make_controller, trail, reply, expected): + """The trail is stripped from the end of a reply only; mid-reply occurrences are preserved, + a bare terminator reduces to an empty payload, and an empty trail strips nothing.""" + kwargs = {} if trail is None else {"trail": trail} + controller = make_controller(**kwargs) + assert controller._remove_trailing_characters(reply) == expected + + +def test_multiple_trail_options_strip_whichever_matches(make_controller): + """With several trail alternatives (e.g. ACS GETVAR vs SETVAR replies), the one + that terminates the reply is stripped.""" + controller = make_controller(term="\r", trail=["\r:\r", ":\r"]) + controller.sock.buffer_recv = [b"8000.5\r:\r"] + assert controller.socket_put_and_receive("?GETVAR(1)") == "8000.5" + controller.sock.buffer_recv = [b":\r"] + assert controller.socket_put_and_receive("SETVAR(1,1)") == "" + + +def test_socket_get_reassembles_chunked_reply(make_controller): + controller = make_controller() + controller.sock.buffer_recv = [b"val", b"ue\r", b"\n"] + assert controller.socket_put_and_receive("get") == "value" + + +def test_socket_get_raises_on_runaway_reply(make_controller): + controller = make_controller(max_reply_length=8) + controller.sock.buffer_recv = [b"0123456789abcdef"] + with pytest.raises(ControllerCommunicationError): + controller.socket_put_and_receive("get") + + +def test_put_and_receive_reconnects_and_retries(make_controller): + """A communication error triggers exactly one reconnect, then the retry succeeds.""" + controller = make_controller() + controller.sock.buffer_recv = [b""] # remote closed the connection + + def _restock(): + controller.sock.buffer_recv = [b"value\r\n"] + + with mock.patch.object(controller, "_reconnect", side_effect=_restock) as mock_reconnect: + assert controller.socket_put_and_receive("get") == "value" + mock_reconnect.assert_called_once() + + +@pytest.mark.parametrize( + "kwargs", [{"term": b"\n"}, {"trail": b"\r\n"}, {"term": 13}, {"trail": 0}] +) +def test_term_and_trail_must_be_strings(make_controller, kwargs): + with pytest.raises(TypeError): + make_controller(**kwargs) + + +def test_second_construction_with_conflicting_term_trail_warns(make_controller): + """A second construction for the same host:port keeps the first term/trail and warns.""" + controller = make_controller(term="\r", trail="\n") + with mock.patch("debye_bec.devices.utils.term_trail_controller.logger") as mock_logger: + second = TermTrailController( + name="controller", + socket_cls=SocketMock, + socket_host="localhost", + socket_port=8080, + device_manager=DMMock(), + term="\n", + ) + assert second is controller + mock_logger.warning.assert_called_once() + assert controller._term == "\r" + assert controller._trail == "\n" + + +def test_second_construction_without_conflict_does_not_warn(make_controller): + controller = make_controller(term="\r") + with mock.patch("debye_bec.devices.utils.term_trail_controller.logger") as mock_logger: + for kwargs in ({"term": "\r"}, {}): + second = TermTrailController( + name="controller", + socket_cls=SocketMock, + socket_host="localhost", + socket_port=8080, + device_manager=DMMock(), + **kwargs, + ) + assert second is controller + mock_logger.warning.assert_not_called() + + +def test_subclass_overrides_term_trail_as_class_attributes(make_controller): + """Subclasses can set the wire protocol with class attributes, without any constructor plumbing.""" + + class CRTermController(TermTrailController): + _term = "\r" + _trail = "\n" + + controller = make_controller(controller_cls=CRTermController) + controller.sock.buffer_recv = [b"value\n"] + assert controller.socket_put_and_receive("get") == "value" + assert controller.sock.buffer_put == [b"get\r"]