perf: cache broadcast audience counts

backend/bot/app/web/admin_api_impl/broadcast.py

@@ -5,6 +5,8 @@ from .common import _panel_user_connection_activity
 import asyncio
 from collections import defaultdict
 
+from bot.utils.ttl_cache import AsyncTTLCache
+
 
 BROADCAST_TARGET_ACTIVE_NEVER_CONNECTED = "active_never_connected"
 BROADCAST_TARGETS = {
@@ -16,6 +18,7 @@ BROADCAST_TARGETS = {
     BROADCAST_TARGET_ACTIVE_NEVER_CONNECTED,
 }
 PANEL_ACTIVITY_LOOKUP_CONCURRENCY = 10
+_ADMIN_BROADCAST_AUDIENCE_COUNT_CACHES: Dict[tuple[int, int], AsyncTTLCache] = {}
 
 
 def _resolve_panel_service(request: web.Request) -> Any:
@@ -73,14 +76,92 @@ async def _user_ids_with_active_subscription_never_connected(
         async with semaphore:
             return await _panel_connection_status(panel_service, panel_uuid)
 
+    panel_uuids = list(
+        dict.fromkeys(
+            panel_uuid
+            for user_panel_uuids in panel_uuids_by_user.values()
+            for panel_uuid in user_panel_uuids
+        )
+    )
+    statuses_by_uuid = dict(
+        zip(
+            panel_uuids,
+            await asyncio.gather(*(lookup(uuid) for uuid in panel_uuids)),
+        )
+    )
+
     user_ids: List[int] = []
     for user_id, panel_uuids in panel_uuids_by_user.items():
-        statuses = await asyncio.gather(*(lookup(panel_uuid) for panel_uuid in panel_uuids))
+        statuses = [statuses_by_uuid.get(panel_uuid, "unknown") for panel_uuid in panel_uuids]
         if statuses and all(status == "never" for status in statuses):
             user_ids.append(user_id)
     return user_ids
 
 
+def _admin_broadcast_audience_counts_cache(settings: Settings) -> Optional[AsyncTTLCache]:
+    ttl_seconds = int(
+        getattr(settings, "ADMIN_BROADCAST_AUDIENCE_COUNTS_CACHE_TTL_SECONDS", 30) or 0
+    )
+    if ttl_seconds <= 0:
+        return None
+    cache_key = (id(settings), ttl_seconds)
+    cache = _ADMIN_BROADCAST_AUDIENCE_COUNT_CACHES.get(cache_key)
+    if cache is None:
+        cache = AsyncTTLCache(
+            ttl_seconds=ttl_seconds,
+            settings=settings,
+            namespace="admin:broadcast_audience_counts",
+        )
+        _ADMIN_BROADCAST_AUDIENCE_COUNT_CACHES[cache_key] = cache
+    return cache
+
+
+async def _load_broadcast_audience_counts(
+    settings: Settings,
+    async_session_factory: sessionmaker,
+    panel_service: Any,
+) -> Dict[str, Optional[int]]:
+    cache = _admin_broadcast_audience_counts_cache(settings)
+    if cache is None:
+        return await _load_broadcast_audience_counts_uncached(
+            async_session_factory,
+            panel_service,
+        )
+    cache_key = "with-panel" if panel_service is not None else "without-panel"
+    return await cache.get_or_load(
+        cache_key,
+        lambda: _load_broadcast_audience_counts_uncached(
+            async_session_factory,
+            panel_service,
+        ),
+    )
+
+
+async def _load_broadcast_audience_counts_uncached(
+    async_session_factory: sessionmaker,
+    panel_service: Any,
+) -> Dict[str, Optional[int]]:
+    async with async_session_factory() as session:
+        counts: Dict[str, Optional[int]] = {
+            "all": await user_dal.count_all_active_users_for_broadcast(session),
+            "active": await user_dal.count_users_with_active_subscription_for_broadcast(session),
+            "inactive": await user_dal.count_users_without_active_subscription_for_broadcast(
+                session
+            ),
+            "expired": await user_dal.count_users_with_expired_subscription_for_broadcast(session),
+            "never": await user_dal.count_users_without_any_subscription_for_broadcast(session),
+            BROADCAST_TARGET_ACTIVE_NEVER_CONNECTED: None,
+        }
+        if panel_service is not None:
+            counts[BROADCAST_TARGET_ACTIVE_NEVER_CONNECTED] = len(
+                await _user_ids_with_active_subscription_never_connected(
+                    session,
+                    panel_service,
+                )
+            )
+    return counts
+
+
 async def admin_broadcast_route(request: web.Request) -> web.Response:
     actor_id = _require_admin_user_id(request)
     payload = await _read_json(request)
@@ -149,23 +230,13 @@ async def admin_broadcast_audience_counts_route(request: web.Request) -> web.Res
     """Return how many users each broadcast audience currently resolves to."""
     _require_admin_user_id(request)
 
+    settings: Settings = request.app["settings"]
     async_session_factory: sessionmaker = request.app["async_session_factory"]
-    async with async_session_factory() as session:
-        panel_service = _resolve_panel_service(request)
-        counts = {
-            "all": len(await user_dal.get_all_active_user_ids_for_broadcast(session)),
-            "active": len(await user_dal.get_user_ids_with_active_subscription(session)),
-            "inactive": len(await user_dal.get_user_ids_without_active_subscription(session)),
-            "expired": len(await user_dal.get_user_ids_with_expired_subscription(session)),
-            "never": len(await user_dal.get_user_ids_without_any_subscription(session)),
-            BROADCAST_TARGET_ACTIVE_NEVER_CONNECTED: None,
-        }
-        if panel_service is not None:
-            counts[BROADCAST_TARGET_ACTIVE_NEVER_CONNECTED] = len(
-                await _user_ids_with_active_subscription_never_connected(
-                    session,
-                    panel_service,
-                )
-            )
+    panel_service = _resolve_panel_service(request)
+    counts = await _load_broadcast_audience_counts(
+        settings,
+        async_session_factory,
+        panel_service,
+    )
 
     return _ok({"counts": counts})