151 lines
6.0 KiB
Python
151 lines
6.0 KiB
Python
"""Redis-backed presence store.
|
|
|
|
Presence is determined by a heartbeat TTL key in Redis:
|
|
- heartbeat(user_id) sets a TTL key per user (default 60 s).
|
|
- is_online(user_id) / is_online_batch(user_ids) check whether the key(s)
|
|
exist.
|
|
|
|
status_message and last_seen_at are persisted in the ``users`` table by the
|
|
router (app/routers/presence.py); this store only owns the Redis side.
|
|
|
|
The Redis client is created lazily, on first actual use, rather than
|
|
requiring an explicit ``connect()`` call during app startup. This matters
|
|
because FastAPI's lifespan handler does not run under the plain
|
|
``TestClient(app)`` instantiation this codebase's tests use (no
|
|
``with`` block), so anything that only connected during lifespan startup
|
|
would never connect in tests. Lazy connection means the store works
|
|
identically whether the app was started for real or is being driven by
|
|
pytest.
|
|
|
|
The client is also cached per-*event loop* rather than globally: the same
|
|
plain ``TestClient(app)`` pattern can dispatch different requests onto
|
|
different event loops (see the NullPool comment on ``_override_db`` in
|
|
tests/conftest.py for the analogous asyncpg problem), and an
|
|
``aioredis.Redis`` connection, like an asyncpg connection, breaks with
|
|
"Event loop is closed" if reused across loops. In the real running server
|
|
there is exactly one event loop for the process lifetime, so this
|
|
degrades to a single cached, pooled connection there.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
import redis.asyncio as aioredis
|
|
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Key prefix for heartbeat entries.
|
|
_PRESENCE_KEY_PREFIX = "presence:heartbeat:"
|
|
|
|
# Default heartbeat TTL in seconds (1 minute).
|
|
# Clients should heartbeat at least once per interval.
|
|
_DEFAULT_HEARTBEAT_TTL = 60
|
|
|
|
|
|
class PresenceStore:
|
|
"""Redis-backed presence store with async, lazily-connected operations."""
|
|
|
|
def __init__(self, redis_url: str | None = None, heartbeat_ttl: int = _DEFAULT_HEARTBEAT_TTL):
|
|
# redis_url is resolved from settings at first use (not at
|
|
# construction time) so tests that override settings.REDIS_URL
|
|
# before the first call still take effect -- this module is
|
|
# imported (and the module-level singleton below constructed)
|
|
# well before any such override runs.
|
|
self._redis_url = redis_url
|
|
self._heartbeat_ttl = heartbeat_ttl
|
|
# Keyed by id(running loop) -- see the module docstring.
|
|
self._clients: dict[int, aioredis.Redis] = {}
|
|
self._connect_lock = asyncio.Lock()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Lifecycle
|
|
# ------------------------------------------------------------------
|
|
|
|
async def _ensure_client(self) -> aioredis.Redis:
|
|
"""Return a connected client for the current event loop, connecting
|
|
on first use per loop."""
|
|
loop_key = id(asyncio.get_running_loop())
|
|
client = self._clients.get(loop_key)
|
|
if client is None:
|
|
async with self._connect_lock:
|
|
client = self._clients.get(loop_key)
|
|
if client is None:
|
|
url = self._redis_url or settings.REDIS_URL
|
|
client = aioredis.from_url(
|
|
url,
|
|
encoding="utf-8",
|
|
decode_responses=True,
|
|
)
|
|
try:
|
|
await client.ping()
|
|
except Exception:
|
|
logger.warning("Redis ping failed for %s", url)
|
|
raise
|
|
self._clients[loop_key] = client
|
|
return client
|
|
|
|
async def close(self) -> None:
|
|
"""Close the Redis connection for the current event loop, if one was
|
|
opened."""
|
|
loop_key = id(asyncio.get_running_loop())
|
|
client = self._clients.pop(loop_key, None)
|
|
if client is not None:
|
|
await client.aclose()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Heartbeat
|
|
# ------------------------------------------------------------------
|
|
|
|
async def heartbeat(self, user_id: int) -> None:
|
|
"""Record a heartbeat for *user_id*.
|
|
|
|
Sets a TTL key so that the user is considered online as long as
|
|
heartbeats arrive at least once per heartbeat_ttl seconds.
|
|
"""
|
|
client = await self._ensure_client()
|
|
await client.set(self._heartbeat_key(user_id), "1", ex=self._heartbeat_ttl)
|
|
|
|
async def clear_heartbeat(self, user_id: int) -> None:
|
|
"""Clear heartbeat for *user_id* to mark them offline immediately."""
|
|
client = await self._ensure_client()
|
|
await client.delete(self._heartbeat_key(user_id))
|
|
|
|
# ------------------------------------------------------------------
|
|
# Online checks
|
|
# ------------------------------------------------------------------
|
|
|
|
async def is_online(self, user_id: int) -> bool:
|
|
"""Return True if the user has a valid heartbeat key."""
|
|
client = await self._ensure_client()
|
|
exists = await client.exists(self._heartbeat_key(user_id))
|
|
return bool(exists)
|
|
|
|
async def is_online_batch(self, user_ids: list[int]) -> set[int]:
|
|
"""Return the subset of user_ids that currently have a valid
|
|
heartbeat key. Uses MGET (one round trip) rather than one EXISTS
|
|
per user.
|
|
"""
|
|
if not user_ids:
|
|
return set()
|
|
client = await self._ensure_client()
|
|
keys = [self._heartbeat_key(uid) for uid in user_ids]
|
|
values = await client.mget(keys)
|
|
return {uid for uid, v in zip(user_ids, values) if v is not None}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _heartbeat_key(user_id: int) -> str:
|
|
return f"{_PRESENCE_KEY_PREFIX}{user_id}"
|
|
|
|
|
|
# Singleton instance shared by the router, so all requests share the same
|
|
# connection pool.
|
|
presence_store = PresenceStore()
|