From 740209912c5a7086fb702f448d249a23869d91c8 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 21 Jan 2024 17:49:06 -1000 Subject: [PATCH] Small performance improvements to handing revoke token callbacks (#108625) - Use a set to avoid linear search for remove - Avoid recreating the unregister function each time --- homeassistant/auth/__init__.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/homeassistant/auth/__init__.py b/homeassistant/auth/__init__.py index ac9bbaaf593..0e9a2429fe4 100644 --- a/homeassistant/auth/__init__.py +++ b/homeassistant/auth/__init__.py @@ -5,6 +5,7 @@ import asyncio from collections import OrderedDict from collections.abc import Mapping from datetime import timedelta +from functools import partial import time from typing import Any, cast @@ -157,7 +158,7 @@ class AuthManager: self._providers = providers self._mfa_modules = mfa_modules self.login_flow = AuthManagerFlowManager(hass, self) - self._revoke_callbacks: dict[str, list[CALLBACK_TYPE]] = {} + self._revoke_callbacks: dict[str, set[CALLBACK_TYPE]] = {} @property def auth_providers(self) -> list[AuthProvider]: @@ -475,27 +476,28 @@ class AuthManager: """Delete a refresh token.""" await self._store.async_remove_refresh_token(refresh_token) - callbacks = self._revoke_callbacks.pop(refresh_token.id, []) + callbacks = self._revoke_callbacks.pop(refresh_token.id, ()) for revoke_callback in callbacks: revoke_callback() + @callback + def _async_unregister( + self, callbacks: set[CALLBACK_TYPE], callback_: CALLBACK_TYPE + ) -> None: + """Unregister a callback.""" + callbacks.remove(callback_) + @callback def async_register_revoke_token_callback( self, refresh_token_id: str, revoke_callback: CALLBACK_TYPE ) -> CALLBACK_TYPE: """Register a callback to be called when the refresh token id is revoked.""" if refresh_token_id not in self._revoke_callbacks: - self._revoke_callbacks[refresh_token_id] = [] + self._revoke_callbacks[refresh_token_id] = set() callbacks = self._revoke_callbacks[refresh_token_id] - callbacks.append(revoke_callback) - - @callback - def unregister() -> None: - if revoke_callback in callbacks: - callbacks.remove(revoke_callback) - - return unregister + callbacks.add(revoke_callback) + return partial(self._async_unregister, callbacks, revoke_callback) @callback def async_create_access_token(