Improve scalability of DHCP matchers (#109406)

This commit is contained in:
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 datetime import timedelta
from fnmatch import translate from fnmatch import translate
from functools import lru_cache from functools import lru_cache
import itertools
import logging import logging
import os import os
import re import re
@ -89,11 +90,55 @@ class DhcpServiceInfo(BaseServiceInfo):
macaddress: str 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: async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up the dhcp component.""" """Set up the dhcp component."""
watchers: list[WatcherBase] = [] watchers: list[WatcherBase] = []
address_data: dict[str, dict[str, str]] = {} 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 the passive classes we need to start listening
# for state changes and connect the dispatchers before # for state changes and connect the dispatchers before
# everything else starts up or we will miss events # everything else starts up or we will miss events
@ -125,7 +170,7 @@ class WatcherBase(ABC):
self, self,
hass: HomeAssistant, hass: HomeAssistant,
address_data: dict[str, dict[str, str]], address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher], integration_matchers: DhcpMatchers,
) -> None: ) -> None:
"""Initialize class.""" """Initialize class."""
super().__init__() super().__init__()
@ -189,28 +234,29 @@ class WatcherBase(ABC):
lowercase_hostname, lowercase_hostname,
) )
matched_domains = set() matched_domains: set[str] = set()
device_domains = set() matchers = self._integration_matchers
registered_devices_domains = matchers.registered_devices_domains
dev_reg: DeviceRegistry = async_get(self.hass) dev_reg: DeviceRegistry = async_get(self.hass)
if device := dev_reg.async_get_device( if device := dev_reg.async_get_device(
connections={(CONNECTION_NETWORK_MAC, uppercase_mac)} connections={(CONNECTION_NETWORK_MAC, uppercase_mac)}
): ):
for entry_id in device.config_entries: for entry_id in device.config_entries:
if entry := self.hass.config_entries.async_get_entry(entry_id): if (
device_domains.add(entry.domain) 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"] 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 ( if (
matcher_hostname := matcher.get(HOSTNAME) matcher_hostname := matcher.get(HOSTNAME)
) is not None and not _memorized_fnmatch( ) is not None and not _memorized_fnmatch(
@ -241,7 +287,7 @@ class NetworkWatcher(WatcherBase):
self, self,
hass: HomeAssistant, hass: HomeAssistant,
address_data: dict[str, dict[str, str]], address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher], integration_matchers: DhcpMatchers,
) -> None: ) -> None:
"""Initialize class.""" """Initialize class."""
super().__init__(hass, address_data, integration_matchers) super().__init__(hass, address_data, integration_matchers)
@ -294,7 +340,7 @@ class DeviceTrackerWatcher(WatcherBase):
self, self,
hass: HomeAssistant, hass: HomeAssistant,
address_data: dict[str, dict[str, str]], address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher], integration_matchers: DhcpMatchers,
) -> None: ) -> None:
"""Initialize class.""" """Initialize class."""
super().__init__(hass, address_data, integration_matchers) super().__init__(hass, address_data, integration_matchers)
@ -349,7 +395,7 @@ class DeviceTrackerRegisteredWatcher(WatcherBase):
self, self,
hass: HomeAssistant, hass: HomeAssistant,
address_data: dict[str, dict[str, str]], address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher], integration_matchers: DhcpMatchers,
) -> None: ) -> None:
"""Initialize class.""" """Initialize class."""
super().__init__(hass, address_data, integration_matchers) super().__init__(hass, address_data, integration_matchers)
@ -387,7 +433,7 @@ class DHCPWatcher(WatcherBase):
self, self,
hass: HomeAssistant, hass: HomeAssistant,
address_data: dict[str, dict[str, str]], address_data: dict[str, dict[str, str]],
integration_matchers: list[DHCPMatcher], integration_matchers: DhcpMatchers,
) -> None: ) -> None:
"""Initialize class.""" """Initialize class."""
super().__init__(hass, address_data, integration_matchers) super().__init__(hass, address_data, integration_matchers)

View File

@ -1,6 +1,8 @@
"""Test the DHCP discovery integration.""" """Test the DHCP discovery integration."""
from collections.abc import Awaitable, Callable
import datetime import datetime
import threading import threading
from typing import Any, cast
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
import pytest 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( dhcp_watcher = dhcp.DHCPWatcher(
hass, hass,
{}, {},
integration_matchers, integration_matchers,
) )
async_handle_dhcp_packet = None async_handle_dhcp_packet: Callable[[Any], Awaitable[None]] | None = None
def _mock_sniffer(*args, **kwargs): def _mock_sniffer(*args, **kwargs):
nonlocal async_handle_dhcp_packet nonlocal async_handle_dhcp_packet
@ -158,14 +162,14 @@ async def _async_get_handle_dhcp_packet(hass, integration_matchers):
): ):
await dhcp_watcher.async_start() 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: async def test_dhcp_match_hostname_and_macaddress(hass: HomeAssistant) -> None:
"""Test matching based on hostname and macaddress.""" """Test matching based on hostname and macaddress."""
integration_matchers = [ integration_matchers = dhcp.async_index_integration_matchers(
{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"} [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}]
] )
packet = Ether(RAW_DHCP_REQUEST) packet = Ether(RAW_DHCP_REQUEST)
async_handle_dhcp_packet = await _async_get_handle_dhcp_packet( 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: async def test_dhcp_renewal_match_hostname_and_macaddress(hass: HomeAssistant) -> None:
"""Test renewal matching based on hostname and macaddress.""" """Test renewal matching based on hostname and macaddress."""
integration_matchers = [ integration_matchers = dhcp.async_index_integration_matchers(
{"domain": "mock-domain", "hostname": "irobot-*", "macaddress": "501479*"} [{"domain": "mock-domain", "hostname": "irobot-*", "macaddress": "501479*"}]
] )
packet = Ether(RAW_DHCP_RENEWAL) packet = Ether(RAW_DHCP_RENEWAL)
@ -220,10 +224,12 @@ async def test_registered_devices(
hass: HomeAssistant, device_registry: dr.DeviceRegistry hass: HomeAssistant, device_registry: dr.DeviceRegistry
) -> None: ) -> None:
"""Test discovery flows are created for registered devices.""" """Test discovery flows are created for registered devices."""
integration_matchers = [ integration_matchers = dhcp.async_index_integration_matchers(
{"domain": "not-matching", "registered_devices": True}, [
{"domain": "mock-domain", "registered_devices": True}, {"domain": "not-matching", "registered_devices": True},
] {"domain": "mock-domain", "registered_devices": True},
]
)
packet = Ether(RAW_DHCP_RENEWAL) packet = Ether(RAW_DHCP_RENEWAL)
@ -265,7 +271,9 @@ async def test_registered_devices(
async def test_dhcp_match_hostname(hass: HomeAssistant) -> None: async def test_dhcp_match_hostname(hass: HomeAssistant) -> None:
"""Test matching based on hostname only.""" """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) 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: async def test_dhcp_match_macaddress(hass: HomeAssistant) -> None:
"""Test matching based on macaddress only.""" """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) 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: async def test_dhcp_multiple_match_only_one_flow(hass: HomeAssistant) -> None:
"""Test matching the domain multiple times only generates one flow.""" """Test matching the domain multiple times only generates one flow."""
integration_matchers = [ integration_matchers = dhcp.async_index_integration_matchers(
{"domain": "mock-domain", "macaddress": "B8B7F1*"}, [
{"domain": "mock-domain", "hostname": "connect"}, {"domain": "mock-domain", "macaddress": "B8B7F1*"},
] {"domain": "mock-domain", "hostname": "connect"},
]
)
packet = Ether(RAW_DHCP_REQUEST) 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: async def test_dhcp_match_macaddress_without_hostname(hass: HomeAssistant) -> None:
"""Test matching based on macaddress only.""" """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) 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: async def test_dhcp_nomatch(hass: HomeAssistant) -> None:
"""Test not matching based on macaddress only.""" """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) 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: async def test_dhcp_nomatch_hostname(hass: HomeAssistant) -> None:
"""Test not matching based on hostname only.""" """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) 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: async def test_dhcp_nomatch_non_dhcp_packet(hass: HomeAssistant) -> None:
"""Test matching does not throw on a non-dhcp packet.""" """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"") 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: async def test_dhcp_nomatch_non_dhcp_request_packet(hass: HomeAssistant) -> None:
"""Test nothing happens with the wrong message-type.""" """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) 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: async def test_dhcp_invalid_hostname(hass: HomeAssistant) -> None:
"""Test we ignore invalid hostnames.""" """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) 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: async def test_dhcp_missing_hostname(hass: HomeAssistant) -> None:
"""Test we ignore missing hostnames.""" """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) 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: async def test_dhcp_invalid_option(hass: HomeAssistant) -> None:
"""Test we ignore invalid hostname option.""" """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) 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( device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() await hass.async_block_till_done()
@ -653,7 +689,15 @@ async def test_device_tracker_registered(hass: HomeAssistant) -> None:
device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher( device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() 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( device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() 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( device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() 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( device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() 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( device_tracker_watcher = dhcp.DeviceTrackerWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() 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( device_tracker_watcher = dhcp.NetworkWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() 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( device_tracker_watcher = dhcp.NetworkWatcher(
hass, hass,
{}, {},
[ dhcp.async_index_integration_matchers(
{ [
"domain": "mock-domain", {
"hostname": "irobot-*", "domain": "mock-domain",
"macaddress": "B8B7F1*", "hostname": "irobot-*",
} "macaddress": "B8B7F1*",
], }
]
),
) )
await device_tracker_watcher.async_start() await device_tracker_watcher.async_start()
await hass.async_block_till_done() 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( device_tracker_watcher = dhcp.NetworkWatcher(
hass, 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 device_tracker_watcher.async_start()
await hass.async_block_till_done() await hass.async_block_till_done()