Skip to content

Commit

Permalink
Migrate to AsyncZeroconf (LedFx#581)
Browse files Browse the repository at this point in the history
* WIP

* Refactor device discovery using AsyncZeroConf

* Fix old comments

* Actually finish the comments (its late)

* Remove debug asserts

* Add GET method to find_devices

* Protect POST and cleanup exception handling

* Check we have used zeroconf before trying to shut it down
  • Loading branch information
shauneccles committed Jan 4, 2024
1 parent 5a80bac commit cc8b9b2
Show file tree
Hide file tree
Showing 5 changed files with 176 additions and 62 deletions.
30 changes: 25 additions & 5 deletions ledfx/api/find_devices.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,11 @@
_LOGGER = logging.getLogger(__name__)


def handle_exception(future):
# Ignore exceptions, these will be raised when a device is found that already exists
exc = future.exception()


class FindDevicesEndpoint(RestEndpoint):
"""REST end-point for detecting and adding wled devices"""

Expand All @@ -21,14 +26,29 @@ async def post(self, request) -> web.Response:
except JSONDecodeError:
return await self.json_decode_error()

set_name_to_icon(data.get("name_to_icon"))
name_to_icon = data.get("name_to_icon")

if name_to_icon is None:
response = {
"status": "failed",
"reason": 'Required attribute "name_to_icon" was not provided',
}
return web.json_response(data=response, status=400)
set_name_to_icon(name_to_icon)

async_fire_and_forget(
self._ledfx.zeroconf.discover_wled_devices(),
loop=self._ledfx.loop,
exc_handler=handle_exception,
)

def handle_exception(future):
# Ignore exceptions, these will be raised when a device is found that already exists
exc = future.exception()
response = {"status": "success"}
return web.json_response(data=response, status=200)

async def get(self) -> web.Response:
"""Handle HTTP GET requests"""
async_fire_and_forget(
self._ledfx.devices.find_wled_devices(),
self._ledfx.zeroconf.discover_wled_devices(),
loop=self._ledfx.loop,
exc_handler=handle_exception,
)
Expand Down
14 changes: 6 additions & 8 deletions ledfx/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
currently_frozen,
)
from ledfx.virtuals import Virtuals
from ledfx.zeroconf import ZeroConfRunner

_LOGGER = logging.getLogger(__name__)
if currently_frozen():
Expand Down Expand Up @@ -259,17 +260,14 @@ async def async_start(self, open_ui=False):
self.devices.create_from_config(self.config["devices"])
await self.devices.async_initialize_devices()

# sync_mode = WLED_CONFIG_SCHEMA(self.config["wled_preferences"])[
# "wled_preferred_mode"
# ]
# if sync_mode:
# await self.devices.set_wleds_sync_mode(sync_mode)

self.zeroconf = ZeroConfRunner(ledfx=self)
self.virtuals.create_from_config(self.config["virtuals"])
self.integrations.create_from_config(self.config["integrations"])

if self.config["scan_on_startup"]:
async_fire_and_forget(self.devices.find_wled_devices(), self.loop)
async_fire_and_forget(
self.zeroconf.discover_wled_devices(), self.loop
)

async_fire_and_forget(
self.integrations.activate_integrations(), self.loop
Expand Down Expand Up @@ -310,7 +308,7 @@ async def async_stop(self, exit_code):
# Fire a shutdown event and flush the loop
self.events.fire_event(LedFxShutdownEvent())
await asyncio.sleep(0)

await self.zeroconf.async_close()
_LOGGER.info("Stopping HttpServer...")
await self.http.stop()

Expand Down
47 changes: 0 additions & 47 deletions ledfx/devices/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import serial
import serial.tools.list_ports
import voluptuous as vol
import zeroconf

from ledfx.config import save_config
from ledfx.events import DeviceCreatedEvent, DeviceUpdateEvent, Event
Expand Down Expand Up @@ -581,11 +580,9 @@ def __init__(self, ledfx):
super().__init__(ledfx, Device, self.PACKAGE_NAME)

def on_shutdown(e):
self._zeroconf.close()
self.deactivate_devices()

self._ledfx.events.add_listener(on_shutdown, Event.LEDFX_SHUTDOWN)
self._zeroconf = zeroconf.Zeroconf()

def create_from_config(self, config):
for device in config:
Expand Down Expand Up @@ -807,47 +804,3 @@ async def set_wleds_sync_mode(self, mode):
device.wled.set_sync_mode(mode)
await device.wled.flush_sync_settings()
device.update_config({"sync_mode": mode})

async def find_wled_devices(self):
# Scan the LAN network that match WLED using zeroconf - Multicast DNS
# Service Discovery Library
_LOGGER.info("Scanning for WLED devices...")
wled_listener = WLEDListener(self._ledfx)
self._zeroconf.add_service_listener("_wled._tcp.local.", wled_listener)
try:
await asyncio.sleep(30)
finally:
_LOGGER.info("WLED device scan finished!")
self._zeroconf.remove_service_listener(wled_listener)


class WLEDListener(zeroconf.ServiceBrowser):
def __init__(self, _ledfx):
self._ledfx = _ledfx

def remove_service(self, zeroconf_obj, type, name):
_LOGGER.info(f"Service {name} removed")

def add_service(self, zeroconf_obj, type, name):
info = zeroconf_obj.get_service_info(type, name)

if info:
hostname = str(info.server).rstrip(".")
_LOGGER.info(f"Found WLED device: {hostname}")

device_type = "wled"
device_config = {"ip_address": hostname}

def handle_exception(future):
# Ignore exceptions, these will be raised when a device is found that already exists
exc = future.exception()

async_fire_and_forget(
self._ledfx.devices.add_new_device(device_type, device_config),
loop=self._ledfx.loop,
exc_handler=handle_exception,
)

def update_service(self, zeroconf_obj, type, name):
"""Callback when a service is updated."""
pass
2 changes: 0 additions & 2 deletions ledfx/virtuals.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@

import numpy as np
import voluptuous as vol
import zeroconf

from ledfx.color import parse_color
from ledfx.effects import DummyEffect
Expand Down Expand Up @@ -940,7 +939,6 @@ def cleanup_effects(e):

self._ledfx = ledfx
self._ledfx.events.add_listener(cleanup_effects, Event.LEDFX_SHUTDOWN)
self._zeroconf = zeroconf.Zeroconf()
self._virtuals = {}

def create_from_config(self, config):
Expand Down
145 changes: 145 additions & 0 deletions ledfx/zeroconf.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,145 @@
import logging

from zeroconf import ServiceStateChange, Zeroconf
from zeroconf.asyncio import (
AsyncServiceBrowser,
AsyncServiceInfo,
AsyncZeroconf,
)

from ledfx.utils import async_fire_and_forget

_LOGGER = logging.getLogger(__name__)


class ZeroConfRunner:
"""
Class responsible for handling zeroconf, WLED discovery and WLED device registration.
Attributes:
aiobrowser: The async service browser for zeroconf.
aiozc: The async zeroconf instance.
_ledfx: The ledfx instance.
Methods:
on_service_state_change: Callback function for service state change.
async_on_service_state_change: Asynchronous function for handling WLED service state change.
add_wled_device: Asynchronous function for adding discovered WLED devices to config.
discover_wled_devices: Asynchronous function for discovering WLED devices.
async_close: Asynchronous function for closing zeroconf listener.
"""

def __init__(self, ledfx):
self.aiobrowser = None
self.aiozc = None
self._ledfx = ledfx

def on_service_state_change(
self,
zeroconf: Zeroconf,
service_type: str,
name: str,
state_change: ServiceStateChange,
):
"""
Callback function for service state change.
Args:
zeroconf (Zeroconf): The zeroconf instance.
service_type (str): The service type.
name (str): The service name.
state_change (ServiceStateChange): The state change event.
"""
# Schedule the coroutine to be run on the event loop
async_fire_and_forget(
self.async_on_service_state_change(
zeroconf=zeroconf,
service_type=service_type,
name=name,
state_change=state_change,
),
self._ledfx.loop,
)

async def async_on_service_state_change(
self,
zeroconf: Zeroconf,
service_type: str,
name: str,
state_change: ServiceStateChange,
):
"""
Asynchronous function for handling service state change.
Args:
zeroconf (Zeroconf): The zeroconf instance.
service_type (str): The service type.
name (str): The service name.
state_change (ServiceStateChange): The state change event.
"""
_LOGGER.debug(
f"Service {name} of type {service_type} state changed: {state_change}"
)
if state_change is not ServiceStateChange.Added:
return

async_fire_and_forget(
self.add_wled_device(zeroconf, service_type, name),
self._ledfx.loop,
)

async def add_wled_device(
self, zeroconf: Zeroconf, service_type: str, name: str
) -> None:
"""
Asynchronous function for feeding discovered WLED devices to add_new_device.
Duplicate detection is handled within add_new_device.
Args:
zeroconf (Zeroconf): The zeroconf instance.
service_type (str): The service type.
name (str): The service name.
"""
info = AsyncServiceInfo(service_type, name)
await info.async_request(zeroconf, 3000)
if info:
hostname = str(info.server).rstrip(".")
_LOGGER.info(f"Found WLED device: {hostname}")

device_type = "wled"
device_config = {"ip_address": hostname}

def handle_exception(future):
# Ignore exceptions, these will be raised when a device is found that already exists
exc = future.exception()

async_fire_and_forget(
self._ledfx.devices.add_new_device(device_type, device_config),
loop=self._ledfx.loop,
exc_handler=handle_exception,
)

async def discover_wled_devices(self) -> None:
"""
Asynchronous function for discovering WLED devices.
"""
self.aiozc = AsyncZeroconf()
services = ["_wled._tcp.local."]
_LOGGER.info("Browsing for WLED devices...")
self.aiobrowser = AsyncServiceBrowser(
self.aiozc.zeroconf,
services,
handlers=[self.on_service_state_change],
)

async def async_close(self) -> None:
"""
Asynchronous function for closing zeroconf listener.
"""
# If aiobrowser exists, then aiozc must also exist.
if self.aiobrowser:
_LOGGER.info("Closing zeroconf listener.")
await self.aiobrowser.async_cancel()
await self.aiozc.async_close()
_LOGGER.info("Zeroconf closed.")

0 comments on commit cc8b9b2

Please sign in to comment.