Files
heartbeat/hbd/server.py
T
2026-02-04 12:45:35 -05:00

129 lines
4.0 KiB
Python

"""Server runtime: starts UDP listener, HTTP server and websocket stubs."""
import asyncio
import logging
from . import udp
logger = logging.getLogger(__name__)
async def _run_async(config):
loop = asyncio.get_running_loop()
# shared runtime collections and helpers
msgs = []
# prepare runtime dependencies
import threading
import time
import hbdclass
from . import http as http_mod
from . import ws as ws_mod
from . import dns as dns_mod
from . import notify as notify_mod
notify_mod.setup(config)
def log(host, m, service=None):
ts = time.time()
s = f"{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(ts))} {host or ''} {m}"
msgs.append(s)
logger.info(s)
email = notify_mod.email
pushmsg = notify_mod.pushmsg_from_config
msg_to_websockets = ws_mod.broadcast
# UDP server endpoint (handler wired to handle_datagram with context)
bind_addr = ("0.0.0.0", config.get("hb_port", 50003))
logger.info("Starting UDP server on %s:%s", *bind_addr)
def udp_handler(msg, addr, transport):
ctx = dict(
config=config,
hbdclass=hbdclass,
log=log,
email=email,
pushmsg=pushmsg,
msg_to_websockets=msg_to_websockets,
msgs=msgs,
DEBUG=config.get("debug", 0),
verbose=config.get("verbose", False),
)
udp.handle_datagram(msg, addr, transport, ctx)
transport, protocol = await loop.create_datagram_endpoint(
lambda: udp.EchoServerProtocol(config=config, handler=udp_handler),
local_addr=bind_addr,
)
# HTTP server (runs in its own thread)
try:
handler_cls = http_mod.make_handler_class(
config=config,
hbdclass=hbdclass,
msgs_getter=lambda: msgs,
log=log,
email=email,
pushmsg=pushmsg,
msg_to_websockets=msg_to_websockets,
tcss=None,
DEBUG=config.get("debug", 0),
verbose=config.get("verbose", False),
get_now=lambda: time.time(),
VER="",
)
serv = http_mod.HttpServer((config.get("hbd_host", ""), config.get("hbd_port", 50004)), handler_cls)
http_thread = threading.Thread(target=serv.serve_forever, daemon=True)
http_thread.start()
logger.info("HTTP server started on %s:%s", config.get("hbd_host", ""), config.get("hbd_port", 50004))
except Exception as e:
logger.exception("failed to start HTTP server: %s", e)
# start dns update thread
dns_mod.start_dns_thread(hbdclass, config, log=log, email=email)
logger.info("dns update thread started")
# Start the websocket servers as a background task
try:
ws_task = asyncio.create_task(
ws_mod.start(
host=config.get("hbd_host", ""),
ws_port=config.get("ws_port", 50005),
wss_port=config.get("wss_port", None),
ssl_context=None,
get_hosts=lambda: [hbdclass.Host.hosts[h].stateinfo() for h in sorted(hbdclass.Host.hosts)],
get_msgs=lambda: msgs,
verbose=config.get("verbose", False),
)
)
logger.info("WebSocket task started")
except Exception as e:
logger.exception("websocket server failed to start: %s", e)
try:
# run forever
await asyncio.Future()
finally:
transport.close()
try:
serv.shutdown()
except Exception:
pass
try:
ws_task.cancel()
except Exception:
pass
def run(config):
"""Start the hbd service (blocking).
This is a thin wrapper around asyncio.run to host the async services.
"""
logging.basicConfig(level=logging.DEBUG if config.get("debug", 0) > 0 else logging.INFO)
try:
asyncio.run(_run_async(config))
except KeyboardInterrupt:
logger.info("Shutting down (KeyboardInterrupt)")