Refactor rainmachine to increase chance of reusing the connection (#111573)

This commit is contained in:
J. Nick Koston 2024-02-26 16:31:18 -10:00 committed by GitHub
parent 26079a6eaf
commit 4bdd8dbd40
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 20 deletions

View File

@ -1,7 +1,6 @@
"""Support for RainMachine devices.""" """Support for RainMachine devices."""
from __future__ import annotations from __future__ import annotations
import asyncio
from collections.abc import Callable, Coroutine from collections.abc import Callable, Coroutine
from dataclasses import dataclass from dataclasses import dataclass
from datetime import timedelta from datetime import timedelta
@ -206,13 +205,10 @@ async def async_update_programs_and_zones(
programs affect zones and certain combinations of zones affect programs. programs affect zones and certain combinations of zones affect programs.
""" """
data: RainMachineData = hass.data[DOMAIN][entry.entry_id] data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
# No gather here to allow http keep-alive to reuse
await asyncio.gather( # the connection for each coordinator.
*[ await data.coordinators[DATA_PROGRAMS].async_refresh()
data.coordinators[DATA_PROGRAMS].async_refresh(), await data.coordinators[DATA_ZONES].async_refresh()
data.coordinators[DATA_ZONES].async_refresh(),
]
)
async def async_setup_entry( # noqa: C901 async def async_setup_entry( # noqa: C901
@ -302,14 +298,6 @@ async def async_setup_entry( # noqa: C901
return data return data
async def async_init_coordinator(
coordinator: RainMachineDataUpdateCoordinator,
) -> None:
"""Initialize a RainMachineDataUpdateCoordinator."""
await coordinator.async_initialize()
await coordinator.async_config_entry_first_refresh()
controller_init_tasks = []
coordinators = {} coordinators = {}
for api_category, update_interval in COORDINATOR_UPDATE_INTERVAL_MAP.items(): for api_category, update_interval in COORDINATOR_UPDATE_INTERVAL_MAP.items():
coordinator = coordinators[api_category] = RainMachineDataUpdateCoordinator( coordinator = coordinators[api_category] = RainMachineDataUpdateCoordinator(
@ -320,9 +308,11 @@ async def async_setup_entry( # noqa: C901
update_interval=update_interval, update_interval=update_interval,
update_method=partial(async_update, api_category), update_method=partial(async_update, api_category),
) )
controller_init_tasks.append(async_init_coordinator(coordinator)) coordinator.async_initialize()
# Its generally faster not to gather here so we can
await asyncio.gather(*controller_init_tasks) # reuse the connection instead of creating a new
# connection for each coordinator.
await coordinator.async_config_entry_first_refresh()
hass.data.setdefault(DOMAIN, {}) hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = RainMachineData( hass.data[DOMAIN][entry.entry_id] = RainMachineData(

View File

@ -120,7 +120,8 @@ class RainMachineDataUpdateCoordinator(DataUpdateCoordinator[dict]): # pylint:
self.config_entry.entry_id self.config_entry.entry_id
) )
async def async_initialize(self) -> None: @callback
def async_initialize(self) -> None:
"""Initialize the coordinator.""" """Initialize the coordinator."""
@callback @callback