led-controller/tests/async_tcp_server.py

#!/usr/bin/env python3
# Standalone async TCP server (stdlib only). Multiple simultaneous clients.
# No watchdog: runs on a full host (e.g. Raspberry Pi); ESP32 clients may use WDT.
# For RTT latency, clients may send lines like ``rtt 12345`` (ticks); they are echoed back.
#
# Run from anywhere (default: all IPv4 interfaces, port 9000):
#   python3 async_tcp_server.py
#   python3 async_tcp_server.py --port 9000
# Localhost only:
#   python3 async_tcp_server.py --host 127.0.0.1
#
# Or from this directory:
#   chmod +x async_tcp_server.py && ./async_tcp_server.py

import argparse
import asyncio
import time


class _ClientRegistry:
    """Track writers and broadcast newline-terminated lines to all clients."""

    def __init__(self) -> None:
        self._writers: set[asyncio.StreamWriter] = set()

    def add(self, writer: asyncio.StreamWriter) -> None:
        self._writers.add(writer)

    def remove(self, writer: asyncio.StreamWriter) -> None:
        self._writers.discard(writer)

    def count(self) -> int:
        return len(self._writers)

    async def broadcast_line(self, line: str) -> None:
        data = (line.rstrip("\r\n") + "\n").encode("utf-8")
        for writer in list(self._writers):
            try:
                writer.write(data)
                await writer.drain()
            except Exception as e:
                print(f"[tcp] broadcast failed, dropping client: {e}")
                self._writers.discard(writer)
                try:
                    writer.close()
                    await writer.wait_closed()
                except Exception:
                    pass


async def _periodic_broadcast(
    registry: _ClientRegistry,
    interval_sec: float,
    message: str,
) -> None:
    while True:
        await asyncio.sleep(interval_sec)
        if registry.count() == 0:
            continue
        line = message.format(t=time.time())
        print(f"[tcp] broadcast to {registry.count()} client(s): {line!r}")
        await registry.broadcast_line(line)


async def _handle_client(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    registry: _ClientRegistry,
) -> None:
    peer = writer.get_extra_info("peername")
    print(f"[tcp] connected: {peer}")
    registry.add(writer)
    try:
        while not reader.at_eof():
            data = await reader.readline()
            if not data:
                break
            message = data.decode("utf-8", errors="replace").rstrip("\r\n")
            # Echo newline-delimited lines (simple test harness behaviour).
            # Clients may send ``rtt <ticks>`` for round-trip timing; echo unchanged.
            t0 = time.perf_counter()
            writer.write((message + "\n").encode("utf-8"))
            await writer.drain()
            if message.startswith("rtt "):
                server_ms = (time.perf_counter() - t0) * 1000.0
                print(
                    f"[tcp] echoed rtt from {peer} "
                    f"(host write+drain ~{server_ms:.2f} ms)"
                )
    finally:
        registry.remove(writer)
        writer.close()
        await writer.wait_closed()
        print(f"[tcp] disconnected: {peer}")


def _make_client_handler(registry: _ClientRegistry):
    async def _handler(
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        await _handle_client(reader, writer, registry)

    return _handler


async def _run(
    host: str,
    port: int,
    broadcast_interval: float | None,
    broadcast_message: str,
) -> None:
    registry = _ClientRegistry()
    handler = _make_client_handler(registry)
    server = await asyncio.start_server(handler, host, port)
    print(f"[tcp] listening on {host}:{port} (Ctrl+C to stop)")
    if broadcast_interval is not None and broadcast_interval > 0:
        print(
            f"[tcp] periodic broadcast every {broadcast_interval}s "
            f"(use {{t}} in --message for unix time)"
        )
    async with server:
        tasks = []
        if broadcast_interval is not None and broadcast_interval > 0:
            tasks.append(
                asyncio.create_task(
                    _periodic_broadcast(registry, broadcast_interval, broadcast_message),
                    name="broadcast",
                )
            )
        try:
            if tasks:
                await asyncio.gather(server.serve_forever(), *tasks)
            else:
                await server.serve_forever()
        finally:
            for t in tasks:
                t.cancel()
            for t in tasks:
                try:
                    await t
                except asyncio.CancelledError:
                    pass


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Standalone asyncio TCP server (multiple connections).",
    )
    parser.add_argument(
        "--host",
        default="0.0.0.0",
        help="bind address (default: all IPv4 interfaces)",
    )
    parser.add_argument("--port", type=int, default=9000, help="bind port")
    parser.add_argument(
        "--interval",
        type=float,
        default=5.0,
        metavar="SEC",
        help="seconds between broadcast lines to all clients (default: 5)",
    )
    parser.add_argument(
        "--message",
        default="ping {t:.0f}",
        help='broadcast line (newline added); use "{t}" for time.time() (default: %(default)s)',
    )
    parser.add_argument(
        "--no-broadcast",
        action="store_true",
        help="disable periodic broadcast (echo-only)",
    )
    args = parser.parse_args()
    interval = None if args.no_broadcast else args.interval
    try:
        asyncio.run(_run(args.host, args.port, interval, args.message))
    except KeyboardInterrupt:
        print("\n[tcp] stopped")


if __name__ == "__main__":
    main()