Files
Zordon/src/utilities/zeroconf_server.py
Brett ecf836c235 Zeroconf offline-handling improvements (#77)
* Add benchmark.py

* Add cpu / disk benchmark APIs

* Add cpu_benchmark method to distributed_job_manager.py

* Do a better job of storing hostnames =

* Remove hostname from Zeroconf cache if server goes offline

* Add cpu / disk benchmark APIs

* Add cpu_benchmark method to distributed_job_manager.py

* Do a better job of storing hostnames =

* Remove hostname from Zeroconf cache if server goes offline

* Wrap main code in try finally block to always stop zeroconf

* Add missing import
2024-02-12 14:57:00 -06:00

113 lines
3.8 KiB
Python

import logging
import socket
from pubsub import pub
from zeroconf import Zeroconf, ServiceInfo, ServiceBrowser, ServiceStateChange, NonUniqueNameException
logger = logging.getLogger()
class ZeroconfServer:
service_type = None
server_name = None
server_port = None
server_ip = None
zeroconf = Zeroconf()
service_info = None
client_cache = {}
properties = {}
@classmethod
def configure(cls, service_type, server_name, server_port):
cls.service_type = service_type
cls.server_name = server_name
cls.server_port = server_port
try: # Stop any previously running instances
socket.gethostbyname(socket.gethostname())
except socket.gaierror:
cls.stop()
@classmethod
def start(cls, listen_only=False):
if not cls.service_type:
raise RuntimeError("The 'configure' method must be run before starting the zeroconf server")
if not listen_only:
cls._register_service()
cls._browse_services()
@classmethod
def stop(cls):
cls._unregister_service()
cls.zeroconf.close()
@classmethod
def _register_service(cls):
try:
cls.server_ip = socket.gethostbyname(socket.gethostname())
info = ServiceInfo(
cls.service_type,
f"{cls.server_name}.{cls.service_type}",
addresses=[socket.inet_aton(cls.server_ip)],
port=cls.server_port,
properties=cls.properties,
)
cls.service_info = info
cls.zeroconf.register_service(info)
logger.info(f"Registered zeroconf service: {cls.service_info.name}")
except (NonUniqueNameException, socket.gaierror) as e:
logger.error(f"Error establishing zeroconf: {e}")
@classmethod
def _unregister_service(cls):
if cls.service_info:
cls.zeroconf.unregister_service(cls.service_info)
logger.info(f"Unregistered zeroconf service: {cls.service_info.name}")
cls.service_info = None
@classmethod
def _browse_services(cls):
browser = ServiceBrowser(cls.zeroconf, cls.service_type, [cls._on_service_discovered])
browser.is_alive()
@classmethod
def _on_service_discovered(cls, zeroconf, service_type, name, state_change):
info = zeroconf.get_service_info(service_type, name)
hostname = name.split(f'.{cls.service_type}')[0]
logger.debug(f"Zeroconf: {hostname} {state_change}")
if service_type == cls.service_type:
if state_change == ServiceStateChange.Added or state_change == ServiceStateChange.Updated:
cls.client_cache[hostname] = info
else:
cls.client_cache.pop(hostname)
pub.sendMessage('zeroconf_state_change', hostname=hostname, state_change=state_change)
@classmethod
def found_hostnames(cls):
local_hostname = socket.gethostname()
def sort_key(hostname):
# Return 0 if it's the local hostname so it comes first, else return 1
return False if hostname == local_hostname else True
# Sort the list with the local hostname first
sorted_hostnames = sorted(cls.client_cache.keys(), key=sort_key)
return sorted_hostnames
@classmethod
def get_hostname_properties(cls, hostname):
server_info = cls.client_cache.get(hostname).properties
decoded_server_info = {key.decode('utf-8'): value.decode('utf-8') for key, value in server_info.items()}
return decoded_server_info
# Example usage:
if __name__ == "__main__":
ZeroconfServer.configure("_zordon._tcp.local.", "foobar.local", 8080)
try:
ZeroconfServer.start()
input("Server running - Press enter to end")
finally:
ZeroconfServer.stop()