jmri_withrottle.server

  1"""``WiThrottleServer`` — accept TCP connections, run a session each.
  2
  3Usage from a host app::
  4
  5    server = WiThrottleServer(connection=conn, roster_provider=lambda: [...])
  6    await server.start(host="0.0.0.0", port=12090)
  7    ...
  8    await server.stop()
  9
 10The server is independent of FastAPI — it holds its own asyncio
 11listener — but it's typically mounted inside the same process as
 12:mod:`jmri_server` so it shares the bean managers and roster.
 13"""
 14from __future__ import annotations
 15
 16import asyncio
 17import logging
 18from collections.abc import Awaitable, Callable
 19from dataclasses import dataclass, field
 20
 21from jmri_core.beans.turnout import Turnout, TurnoutStateChange
 22from jmri_core.hal.connection import SystemConnection
 23
 24from jmri_withrottle import protocol
 25from jmri_withrottle.session import WiThrottleSession
 26
 27_log = logging.getLogger(__name__)
 28
 29DEFAULT_PORT = 12090
 30
 31
 32@dataclass
 33class WiThrottleServer:
 34    """TCP server speaking the WiThrottle line protocol."""
 35
 36    connection: SystemConnection
 37    # Callable so the server can be told about roster + power state
 38    # without us hard-wiring to jmri-server. Returns a list of dicts
 39    # in the roster JSON shape (id, dccAddress, longAddress).
 40    roster_provider: Callable[[], list[dict]] = field(default_factory=lambda: list)
 41    # Callable returning the current fast-clock ``PFT`` sync line, or
 42    # ``None`` if no fast clock is attached. Called on connect and on
 43    # every heartbeat so Engine Driver / WiThrottle iOS clocks stay in
 44    # sync without needing a separate broadcast path.
 45    fast_clock_provider: Callable[[], str | None] = field(
 46        default_factory=lambda: (lambda: None),
 47    )
 48    # Callable returning the current turnout inventory as a list of
 49    # ``{"system_name", "user_name", "state"}`` dicts. Sent once per
 50    # session (after the roster) so Engine Driver / WiThrottle iOS can
 51    # populate their Turnouts tab.
 52    turnout_provider: Callable[[], list[dict]] = field(
 53        default_factory=lambda: list,
 54    )
 55    # Handler for turnout throw/close/toggle commands from the client.
 56    # Called as ``await handler(system_name, action)`` where action is
 57    # one of ``"T"``, ``"C"``, or ``"2"``. If None, turnout actions from
 58    # the client are silently dropped.
 59    turnout_action_handler: (
 60        Callable[[str, str], Awaitable[None]] | None
 61    ) = None
 62    railroad: str = "Python JMRI"
 63    heartbeat_s: int = protocol.DEFAULT_HEARTBEAT_S
 64    web_port: int = 8000
 65    # How often to push a fresh PFT to every connected client. Engine
 66    # Driver's clock extrapolates from the last sync using the rate;
 67    # 5 s keeps it within a couple of fast-seconds of the server after
 68    # rate/state changes without spamming the wire.
 69    fast_clock_broadcast_s: float = 5.0
 70    _server: asyncio.Server | None = None
 71    _sessions: set[WiThrottleSession] = field(default_factory=set)
 72    _pft_task: asyncio.Task[None] | None = None
 73    # Disconnect callables from ``state_changed`` subscriptions and
 74    # the turnout manager's ``added``/``removed`` signals. Filled on
 75    # ``start()``, drained on ``stop()``.
 76    _turnout_unsubs: list[Callable[[], None]] = field(default_factory=list)
 77
 78    # ----- lifecycle -----------------------------------------------------
 79
 80    async def start(self, host: str = "0.0.0.0", port: int = DEFAULT_PORT) -> None:
 81        if self._server is not None:
 82            return
 83        self._server = await asyncio.start_server(
 84            self._handle, host=host, port=port,
 85        )
 86        self._pft_task = asyncio.create_task(
 87            self._pft_broadcast_loop(), name="withrottle-pft-broadcast",
 88        )
 89        # Wire turnout state broadcasts: every existing turnout gets a
 90        # ``state_changed`` listener that fires a PTA line to every
 91        # connected client. ``added``/``removed`` on the manager keep
 92        # the subscription set in sync so turnouts created later (via
 93        # the JSON API, web UI, or a driver monitor) also broadcast.
 94        # Must run in ``start()`` (post-``conn.open()``), not at
 95        # construction — drivers swap in their own turnout manager
 96        # during ``open()`` and any earlier subscription would leak to
 97        # the base manager.
 98        for t in self.connection.turnouts:
 99            self._turnout_unsubs.append(self._subscribe_turnout(t))
100        self._turnout_unsubs.append(
101            self.connection.turnouts.added.connect(self._on_turnout_added)
102        )
103        _log.info("WiThrottle server listening on %s:%d", host, port)
104
105    def _subscribe_turnout(self, t: Turnout) -> Callable[[], None]:
106        """Attach a ``state_changed`` listener to *t* that pushes PTA."""
107        async def _listener(change: TurnoutStateChange) -> None:
108            # Only ``known`` transitions produce a wire-visible state on
109            # the Turnouts tab; commanded transitions still show the
110            # previous known state until the layout confirms.
111            if change.kind != "known":
112                return
113            await self.broadcast_turnout_state(
114                t.system_name, int(change.current),
115            )
116        return t.state_changed.connect(_listener)
117
118    async def _on_turnout_added(self, t: Turnout) -> None:
119        self._turnout_unsubs.append(self._subscribe_turnout(t))
120
121    async def stop(self) -> None:
122        if self._server is None:
123            return
124        if self._pft_task is not None:
125            self._pft_task.cancel()
126            import contextlib  # noqa: PLC0415
127            with contextlib.suppress(asyncio.CancelledError, Exception):
128                await self._pft_task
129            self._pft_task = None
130        for unsub in self._turnout_unsubs:
131            try:
132                unsub()
133            except Exception as e:  # noqa: BLE001
134                _log.debug("turnout unsubscribe failed: %s", e)
135        self._turnout_unsubs.clear()
136        for s in list(self._sessions):
137            await s.close()
138        self._server.close()
139        await self._server.wait_closed()
140        self._server = None
141        _log.info("WiThrottle server stopped")
142
143    async def broadcast_turnout_state(
144        self, system_name: str, state: int,
145    ) -> None:
146        """Push a ``PTA<state><sysName>`` line to every connected session.
147
148        The jmri_server wiring subscribes to each turnout's ``state_changed``
149        signal and calls this so the WiThrottle Turnouts tab stays in
150        sync with layout state — whether the change came from another
151        WiThrottle client, the web UI, or a physical cab.
152        """
153        if not self._sessions:
154            return
155        line = protocol.turnout_state(system_name, state)
156        for sess in list(self._sessions):
157            try:
158                await sess.send_line(line)
159            except Exception as e:  # noqa: BLE001
160                _log.debug("turnout broadcast to session failed: %s", e)
161
162    async def _pft_broadcast_loop(self) -> None:
163        """Periodically push the current PFT to every session.
164
165        Engine Driver clocks extrapolate locally from the last sync line,
166        but if the server-side clock state changes (start, stop, rate
167        change, big time jump) the client can drift until the next sync.
168        A ~15 s broadcast keeps every client aligned within one fast-clock
169        second at 8x rate, with negligible wire cost.
170        """
171        try:
172            while True:
173                await asyncio.sleep(self.fast_clock_broadcast_s)
174                pft = self.fast_clock_provider()
175                if pft is None or not self._sessions:
176                    continue
177                for sess in list(self._sessions):
178                    try:
179                        await sess.send_line(pft)
180                    except Exception as e:  # noqa: BLE001
181                        _log.debug("PFT broadcast to session failed: %s", e)
182        except asyncio.CancelledError:
183            return
184
185    @property
186    def is_running(self) -> bool:
187        return self._server is not None
188
189    @property
190    def sockets(self) -> list:
191        """Bound sockets — useful in tests to discover the assigned port."""
192        if self._server is None:
193            return []
194        return list(self._server.sockets)
195
196    # ----- per-connection handler ----------------------------------------
197
198    async def _handle(
199        self,
200        reader: asyncio.StreamReader,
201        writer: asyncio.StreamWriter,
202    ) -> None:
203        peer = writer.get_extra_info("peername")
204        _log.info("WiThrottle client connected: %s", peer)
205
206        async def send_line(line: str) -> None:
207            try:
208                writer.write((line + "\n").encode("utf-8"))
209                await writer.drain()
210            except (ConnectionResetError, BrokenPipeError):
211                pass
212
213        session = WiThrottleSession(
214            throttle_mgr=self.connection.throttles,
215            send_line=send_line,
216            consist_mgr=getattr(self.connection, "consists", None),
217            fast_clock_provider=self.fast_clock_provider,
218            turnout_action_handler=self.turnout_action_handler,
219        )
220        self._sessions.add(session)
221        try:
222            # Greeting + initial state.
223            for line in protocol.greeting(
224                name=self.railroad,
225                heartbeat_s=self.heartbeat_s,
226                web_port=self.web_port,
227            ):
228                await send_line(line)
229            roster = self.roster_provider()
230            if roster:
231                await send_line(protocol.roster_list(roster))
232            # Push the initial fast-clock reading so the client's clock
233            # doesn't sit at 12:00 PM waiting for the first heartbeat.
234            pft = self.fast_clock_provider()
235            if pft is not None:
236                await send_line(pft)
237            # Turnout state labels + inventory so the Turnouts tab renders.
238            # PTT (titles) must precede PTL (list) so the client already
239            # knows how to render each numeric state code. We skip both
240            # when the connection has no turnouts to avoid confusing
241            # clients that hide the Turnouts tab on an empty PTT.
242            turnouts = self.turnout_provider()
243            if turnouts:
244                await send_line(protocol.turnout_titles())
245                await send_line(protocol.turnout_list(turnouts))
246
247            # Read loop.
248            while True:
249                raw = await reader.readline()
250                if not raw:
251                    break
252                line = raw.decode("utf-8", errors="replace").strip()
253                if not line:
254                    continue
255                keep_going = await session.dispatch(line)
256                if not keep_going:
257                    break
258        except asyncio.CancelledError:
259            raise
260        except Exception as e:  # noqa: BLE001
261            _log.warning("WiThrottle session error from %s: %s", peer, e)
262        finally:
263            await session.close()
264            self._sessions.discard(session)
265            writer.close()
266            with _suppress_close():
267                await writer.wait_closed()
268            _log.info("WiThrottle client disconnected: %s", peer)
269
270
271class _suppress_close:
272    def __enter__(self) -> None: ...
273    def __exit__(self, exc_type, exc, tb) -> bool:  # noqa: ANN001
274        # Swallow close-time errors; the connection's already gone.
275        return exc_type is not None and issubclass(
276            exc_type, (ConnectionError, OSError),
277        )