Speed up bluetooth matching (#122626)

- use a defaultdict to avoid lots of setdefault
- move the intersection outside of the genexpr
  to avoid entering the genexpr if there is no
  intersection
This commit is contained in:
J. Nick Koston 2024-07-26 02:22:56 -05:00 committed by GitHub
parent 78a98afb8d
commit e262f759af
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass from dataclasses import dataclass
from fnmatch import translate from fnmatch import translate
from functools import lru_cache from functools import lru_cache
@ -173,10 +174,10 @@ class BluetoothMatcherIndexBase[
def __init__(self) -> None: def __init__(self) -> None:
"""Initialize the matcher index.""" """Initialize the matcher index."""
self.local_name: dict[str, list[_T]] = {} self.local_name: defaultdict[str, list[_T]] = defaultdict(list)
self.service_uuid: dict[str, list[_T]] = {} self.service_uuid: defaultdict[str, list[_T]] = defaultdict(list)
self.service_data_uuid: dict[str, list[_T]] = {} self.service_data_uuid: defaultdict[str, list[_T]] = defaultdict(list)
self.manufacturer_id: dict[int, list[_T]] = {} self.manufacturer_id: defaultdict[int, list[_T]] = defaultdict(list)
self.service_uuid_set: set[str] = set() self.service_uuid_set: set[str] = set()
self.service_data_uuid_set: set[str] = set() self.service_data_uuid_set: set[str] = set()
self.manufacturer_id_set: set[int] = set() self.manufacturer_id_set: set[int] = set()
@ -190,26 +191,22 @@ class BluetoothMatcherIndexBase[
""" """
# Local name is the cheapest to match since its just a dict lookup # Local name is the cheapest to match since its just a dict lookup
if LOCAL_NAME in matcher: if LOCAL_NAME in matcher:
self.local_name.setdefault( self.local_name[_local_name_to_index_key(matcher[LOCAL_NAME])].append(
_local_name_to_index_key(matcher[LOCAL_NAME]), [] matcher
).append(matcher) )
return True return True
# Manufacturer data is 2nd cheapest since its all ints # Manufacturer data is 2nd cheapest since its all ints
if MANUFACTURER_ID in matcher: if MANUFACTURER_ID in matcher:
self.manufacturer_id.setdefault(matcher[MANUFACTURER_ID], []).append( self.manufacturer_id[matcher[MANUFACTURER_ID]].append(matcher)
matcher
)
return True return True
if SERVICE_UUID in matcher: if SERVICE_UUID in matcher:
self.service_uuid.setdefault(matcher[SERVICE_UUID], []).append(matcher) self.service_uuid[matcher[SERVICE_UUID]].append(matcher)
return True return True
if SERVICE_DATA_UUID in matcher: if SERVICE_DATA_UUID in matcher:
self.service_data_uuid.setdefault(matcher[SERVICE_DATA_UUID], []).append( self.service_data_uuid[matcher[SERVICE_DATA_UUID]].append(matcher)
matcher
)
return True return True
return False return False
@ -260,32 +257,38 @@ class BluetoothMatcherIndexBase[
if ble_device_matches(matcher, service_info) if ble_device_matches(matcher, service_info)
) )
if self.service_data_uuid_set and service_info.service_data: if (
(service_data_uuid_set := self.service_data_uuid_set)
and (service_data := service_info.service_data)
and (matched_uuids := service_data_uuid_set.intersection(service_data))
):
matches.extend( matches.extend(
matcher matcher
for service_data_uuid in self.service_data_uuid_set.intersection( for service_data_uuid in matched_uuids
service_info.service_data
)
for matcher in self.service_data_uuid[service_data_uuid] for matcher in self.service_data_uuid[service_data_uuid]
if ble_device_matches(matcher, service_info) if ble_device_matches(matcher, service_info)
) )
if self.manufacturer_id_set and service_info.manufacturer_data: if (
(manufacturer_id_set := self.manufacturer_id_set)
and (manufacturer_data := service_info.manufacturer_data)
and (matched_ids := manufacturer_id_set.intersection(manufacturer_data))
):
matches.extend( matches.extend(
matcher matcher
for manufacturer_id in self.manufacturer_id_set.intersection( for manufacturer_id in matched_ids
service_info.manufacturer_data
)
for matcher in self.manufacturer_id[manufacturer_id] for matcher in self.manufacturer_id[manufacturer_id]
if ble_device_matches(matcher, service_info) if ble_device_matches(matcher, service_info)
) )
if self.service_uuid_set and service_info.service_uuids: if (
(service_uuid_set := self.service_uuid_set)
and (service_uuids := service_info.service_uuids)
and (matched_uuids := service_uuid_set.intersection(service_uuids))
):
matches.extend( matches.extend(
matcher matcher
for service_uuid in self.service_uuid_set.intersection( for service_uuid in matched_uuids
service_info.service_uuids
)
for matcher in self.service_uuid[service_uuid] for matcher in self.service_uuid[service_uuid]
if ble_device_matches(matcher, service_info) if ble_device_matches(matcher, service_info)
) )
@ -310,7 +313,9 @@ class BluetoothCallbackMatcherIndex(
def __init__(self) -> None: def __init__(self) -> None:
"""Initialize the matcher index.""" """Initialize the matcher index."""
super().__init__() super().__init__()
self.address: dict[str, list[BluetoothCallbackMatcherWithCallback]] = {} self.address: defaultdict[str, list[BluetoothCallbackMatcherWithCallback]] = (
defaultdict(list)
)
self.connectable: list[BluetoothCallbackMatcherWithCallback] = [] self.connectable: list[BluetoothCallbackMatcherWithCallback] = []
def add_callback_matcher( def add_callback_matcher(
@ -323,7 +328,7 @@ class BluetoothCallbackMatcherIndex(
We put them in the bucket that they are most likely to match. We put them in the bucket that they are most likely to match.
""" """
if ADDRESS in matcher: if ADDRESS in matcher:
self.address.setdefault(matcher[ADDRESS], []).append(matcher) self.address[matcher[ADDRESS]].append(matcher)
return return
if super().add(matcher): if super().add(matcher):