refactor(devices): vendor term/trail socket controller from ophyd_devices
This commit is contained in:
@@ -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.")
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user