Improve scalability of DHCP matchers (#109406)

pull/109628/head^2
J. Nick Koston 2024-02-04 16:50:08 -06:00 committed by GitHub
parent 52d27230bc
commit 73589015c3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 202 additions and 62 deletions

View File

@ -9,6 +9,7 @@ from dataclasses import dataclass
from datetime import timedelta
from fnmatch import translate
from functools import lru_cache
import itertools
import logging
import os
import re
@ -89,11 +90,55 @@ class DhcpServiceInfo(BaseServiceInfo):
macaddress: str
@dataclass(slots=True)
class DhcpMatchers:
"""Prepared info from dhcp entries."""
registered_devices_domains: set[str]
no_oui_matchers: dict[str, list[DHCPMatcher]]
oui_matchers: dict[str, list[DHCPMatcher]]
def async_index_integration_matchers(
integration_matchers: list[DHCPMatcher],
) -> DhcpMatchers:
"""Index the integration matchers.
We have three types of matchers:
1. Registered devices
2. Devices with no OUI - index by first char of lower() hostname
3. Devices with OUI - index by OUI
"""
registered_devices_domains: set[str] = set()
no_oui_matchers: dict[str, list[DHCPMatcher]] = {}
oui_matchers: dict[str, list[DHCPMatcher]] = {}
for matcher in integration_matchers:
domain = matcher["domain"]
if REGISTERED_DEVICES in matcher:
registered_devices_domains.add(domain)
continue
if mac_address := matcher.get(MAC_ADDRESS):
oui_matchers.setdefault(mac_address[:6], []).append(matcher)
continue
if hostname := matcher.get(HOSTNAME):
first_char = hostname[0].lower()
no_oui_matchers.setdefault(first_char, []).append(matcher)
return DhcpMatchers(
registered_devices_domains=registered_devices_domains,
no_oui_matchers=no_oui_matchers,
oui_matchers=oui_matchers,
)
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up the dhcp component."""
watchers: list[WatcherBase] = []
address_data: dict[str, dict[str, str]] = {}
integration_matchers = await async_get_dhcp(hass)
integration_matchers = async_index_integration_matchers(await async_get_dhcp(hass))
# For the passive classes we need to start listening
# for state changes and connect the dispatchers before
# everything else starts up or we will miss events
@ -125,7 +170,7 @@ class WatcherBase(ABC):
self,
hass: HomeAssistant,
address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher],
integration_matchers: DhcpMatchers,
) -> None:
"""Initialize class."""
super().__init__()
@ -189,28 +234,29 @@ class WatcherBase(ABC):
lowercase_hostname,
)
matched_domains = set()
device_domains = set()
matched_domains: set[str] = set()
matchers = self._integration_matchers
registered_devices_domains = matchers.registered_devices_domains
dev_reg: DeviceRegistry = async_get(self.hass)
if device := dev_reg.async_get_device(
connections={(CONNECTION_NETWORK_MAC, uppercase_mac)}
):
for entry_id in device.config_entries:
if entry := self.hass.config_entries.async_get_entry(entry_id):
device_domains.add(entry.domain)
if (
entry := self.hass.config_entries.async_get_entry(entry_id)
) and entry.domain in registered_devices_domains:
matched_domains.add(entry.domain)
for matcher in self._integration_matchers:
oui = uppercase_mac[:6]
lowercase_hostname_first_char = (
lowercase_hostname[0] if len(lowercase_hostname) else ""
)
for matcher in itertools.chain(
matchers.no_oui_matchers.get(lowercase_hostname_first_char, ()),
matchers.oui_matchers.get(oui, ()),
):
domain = matcher["domain"]
if matcher.get(REGISTERED_DEVICES) and domain not in device_domains:
continue
if (
matcher_mac := matcher.get(MAC_ADDRESS)
) is not None and not _memorized_fnmatch(uppercase_mac, matcher_mac):
continue
if (
matcher_hostname := matcher.get(HOSTNAME)
) is not None and not _memorized_fnmatch(
@ -241,7 +287,7 @@ class NetworkWatcher(WatcherBase):
self,
hass: HomeAssistant,
address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher],
integration_matchers: DhcpMatchers,
) -> None:
"""Initialize class."""
super().__init__(hass, address_data, integration_matchers)
@ -294,7 +340,7 @@ class DeviceTrackerWatcher(WatcherBase):
self,
hass: HomeAssistant,
address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher],
integration_matchers: DhcpMatchers,
) -> None:
"""Initialize class."""
super().__init__(hass, address_data, integration_matchers)
@ -349,7 +395,7 @@ class DeviceTrackerRegisteredWatcher(WatcherBase):
self,
hass: HomeAssistant,
address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher],
integration_matchers: DhcpMatchers,
) -> None:
"""Initialize class."""
super().__init__(hass, address_data, integration_matchers)
@ -387,7 +433,7 @@ class DHCPWatcher(WatcherBase):
self,
hass: HomeAssistant,
address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher],
integration_matchers: DhcpMatchers,
) -> None:
"""Initialize class."""
super().__init__(hass, address_data, integration_matchers)

View File

@ -1,6 +1,8 @@
"""Test the DHCP discovery integration."""
from collections.abc import Awaitable, Callable
import datetime
import threading
from typing import Any, cast
from unittest.mock import MagicMock, patch
import pytest
@ -130,13 +132,15 @@ RAW_DHCP_REQUEST_WITHOUT_HOSTNAME = (
)
async def _async_get_handle_dhcp_packet(hass, integration_matchers):
async def _async_get_handle_dhcp_packet(
hass: HomeAssistant, integration_matchers: dhcp.DhcpMatchers
) -> Callable[[Any], Awaitable[None]]:
dhcp_watcher = dhcp.DHCPWatcher(
hass,
{},
integration_matchers,
)
async_handle_dhcp_packet = None
async_handle_dhcp_packet: Callable[[Any], Awaitable[None]] | None = None
def _mock_sniffer(*args, **kwargs):
nonlocal async_handle_dhcp_packet
@ -158,14 +162,14 @@ async def _async_get_handle_dhcp_packet(hass, integration_matchers):
):
await dhcp_watcher.async_start()
return async_handle_dhcp_packet
return cast("Callable[[Any], Awaitable[None]]", async_handle_dhcp_packet)
async def test_dhcp_match_hostname_and_macaddress(hass: HomeAssistant) -> None:
"""Test matching based on hostname and macaddress."""
integration_matchers = [
{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}
]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
async_handle_dhcp_packet = await _async_get_handle_dhcp_packet(
@ -190,9 +194,9 @@ async def test_dhcp_match_hostname_and_macaddress(hass: HomeAssistant) -> None:
async def test_dhcp_renewal_match_hostname_and_macaddress(hass: HomeAssistant) -> None:
"""Test renewal matching based on hostname and macaddress."""
integration_matchers = [
{"domain": "mock-domain", "hostname": "irobot-*", "macaddress": "501479*"}
]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "irobot-*", "macaddress": "501479*"}]
)
packet = Ether(RAW_DHCP_RENEWAL)
@ -220,10 +224,12 @@ async def test_registered_devices(
hass: HomeAssistant, device_registry: dr.DeviceRegistry
) -> None:
"""Test discovery flows are created for registered devices."""
integration_matchers = [
{"domain": "not-matching", "registered_devices": True},
{"domain": "mock-domain", "registered_devices": True},
]
integration_matchers = dhcp.async_index_integration_matchers(
[
{"domain": "not-matching", "registered_devices": True},
{"domain": "mock-domain", "registered_devices": True},
]
)
packet = Ether(RAW_DHCP_RENEWAL)
@ -265,7 +271,9 @@ async def test_registered_devices(
async def test_dhcp_match_hostname(hass: HomeAssistant) -> None:
"""Test matching based on hostname only."""
integration_matchers = [{"domain": "mock-domain", "hostname": "connect"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "connect"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -289,7 +297,9 @@ async def test_dhcp_match_hostname(hass: HomeAssistant) -> None:
async def test_dhcp_match_macaddress(hass: HomeAssistant) -> None:
"""Test matching based on macaddress only."""
integration_matchers = [{"domain": "mock-domain", "macaddress": "B8B7F1*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "macaddress": "B8B7F1*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -313,10 +323,12 @@ async def test_dhcp_match_macaddress(hass: HomeAssistant) -> None:
async def test_dhcp_multiple_match_only_one_flow(hass: HomeAssistant) -> None:
"""Test matching the domain multiple times only generates one flow."""
integration_matchers = [
{"domain": "mock-domain", "macaddress": "B8B7F1*"},
{"domain": "mock-domain", "hostname": "connect"},
]
integration_matchers = dhcp.async_index_integration_matchers(
[
{"domain": "mock-domain", "macaddress": "B8B7F1*"},
{"domain": "mock-domain", "hostname": "connect"},
]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -340,7 +352,9 @@ async def test_dhcp_multiple_match_only_one_flow(hass: HomeAssistant) -> None:
async def test_dhcp_match_macaddress_without_hostname(hass: HomeAssistant) -> None:
"""Test matching based on macaddress only."""
integration_matchers = [{"domain": "mock-domain", "macaddress": "606BBD*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "macaddress": "606BBD*"}]
)
packet = Ether(RAW_DHCP_REQUEST_WITHOUT_HOSTNAME)
@ -364,7 +378,9 @@ async def test_dhcp_match_macaddress_without_hostname(hass: HomeAssistant) -> No
async def test_dhcp_nomatch(hass: HomeAssistant) -> None:
"""Test not matching based on macaddress only."""
integration_matchers = [{"domain": "mock-domain", "macaddress": "ABC123*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "macaddress": "ABC123*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -379,7 +395,9 @@ async def test_dhcp_nomatch(hass: HomeAssistant) -> None:
async def test_dhcp_nomatch_hostname(hass: HomeAssistant) -> None:
"""Test not matching based on hostname only."""
integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "nomatch*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -394,7 +412,9 @@ async def test_dhcp_nomatch_hostname(hass: HomeAssistant) -> None:
async def test_dhcp_nomatch_non_dhcp_packet(hass: HomeAssistant) -> None:
"""Test matching does not throw on a non-dhcp packet."""
integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "nomatch*"}]
)
packet = Ether(b"")
@ -409,7 +429,9 @@ async def test_dhcp_nomatch_non_dhcp_packet(hass: HomeAssistant) -> None:
async def test_dhcp_nomatch_non_dhcp_request_packet(hass: HomeAssistant) -> None:
"""Test nothing happens with the wrong message-type."""
integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "nomatch*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -433,7 +455,9 @@ async def test_dhcp_nomatch_non_dhcp_request_packet(hass: HomeAssistant) -> None
async def test_dhcp_invalid_hostname(hass: HomeAssistant) -> None:
"""Test we ignore invalid hostnames."""
integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "nomatch*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -457,7 +481,9 @@ async def test_dhcp_invalid_hostname(hass: HomeAssistant) -> None:
async def test_dhcp_missing_hostname(hass: HomeAssistant) -> None:
"""Test we ignore missing hostnames."""
integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "nomatch*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -481,7 +507,9 @@ async def test_dhcp_missing_hostname(hass: HomeAssistant) -> None:
async def test_dhcp_invalid_option(hass: HomeAssistant) -> None:
"""Test we ignore invalid hostname option."""
integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}]
integration_matchers = dhcp.async_index_integration_matchers(
[{"domain": "mock-domain", "hostname": "nomatch*"}]
)
packet = Ether(RAW_DHCP_REQUEST)
@ -628,7 +656,15 @@ async def test_device_tracker_hostname_and_macaddress_exists_before_start(
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -653,7 +689,15 @@ async def test_device_tracker_registered(hass: HomeAssistant) -> None:
device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -684,7 +728,15 @@ async def test_device_tracker_registered_hostname_none(hass: HomeAssistant) -> N
device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -709,7 +761,15 @@ async def test_device_tracker_hostname_and_macaddress_after_start(
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -748,7 +808,15 @@ async def test_device_tracker_hostname_and_macaddress_after_start_not_home(
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -877,7 +945,15 @@ async def test_device_tracker_ignore_self_assigned_ips_before_start(
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -902,7 +978,15 @@ async def test_aiodiscover_finds_new_hosts(hass: HomeAssistant) -> None:
device_tracker_watcher = dhcp.NetworkWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -953,13 +1037,15 @@ async def test_aiodiscover_does_not_call_again_on_shorter_hostname(
device_tracker_watcher = dhcp.NetworkWatcher(
hass,
{},
[
{
"domain": "mock-domain",
"hostname": "irobot-*",
"macaddress": "B8B7F1*",
}
],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "irobot-*",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -996,7 +1082,15 @@ async def test_aiodiscover_finds_new_hosts_after_interval(hass: HomeAssistant) -
device_tracker_watcher = dhcp.NetworkWatcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
)
await device_tracker_watcher.async_start()
await hass.async_block_till_done()