Private
Public Access
Expire presence so users can't get stuck online (#77)
Presence was a per-user refcount in Redis, decremented only when a connection's cleanup ran. Anything that ended a connection without cleanup (a crash, an unclean restart, a laptop asleep with the page frozen) leaked its count forever, so those users never showed offline. Production had 19 counts for 2 live connections. Presence is now one expiring entry per connection (90s TTL) for global, per-room and focus presence, in new presence:v2:* keys. Clients renew theirs by sending any frame, which the existing 30s ping/pong already does. The server closes a human connection that sends nothing for 90s (code 4408), which catches a frozen page whose browser still answers transport pings. Bot connections can't be expected to answer pings, so the server renews theirs itself while the handler is alive. An expiry has no disconnect event to announce it, so a sweeper on each worker (one per 15s overall, via a Redis lock) prunes expired entries and broadcasts the offline transition. The old presence:* hashes are deleted at startup. Tests cover the expiry store, the sweeper, and the websocket cases. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
7ff7f2ec0e
commit
dda85a2b78
11 files changed
+843
-97
No files matched your search
+19
-10
@@ -289,17 +289,26 @@ instance and running many behave identically:
|
||||
to `room:*`; each message it receives is handed to its own local
|
||||
`ConnectionManager.broadcast()`. One instance just talks to itself
|
||||
through Redis, so there's no separate single-instance code path.
|
||||
- `app/ws/presence.py` (`Presence`) — a Redis hash per room
|
||||
(`presence:{room_id}`, field = user ID, value = connection refcount) is
|
||||
the cross-instance answer to "is this member connected *anywhere* right
|
||||
now," which is what the Phase 4 offline-push check uses instead of the
|
||||
local `ConnectionManager`. Refcounted so a user connected from two tabs
|
||||
(or two instances) isn't marked offline until every connection closes.
|
||||
- `app/ws/presence.py` (`Presence`) — a Redis sorted set per room
|
||||
(`presence:v2:room:{id}`, one member per *connection*, `<user_id>:<conn_id>`,
|
||||
score = expiry time) is the cross-instance answer to "is this member
|
||||
connected *anywhere* right now," which is what the Phase 4 offline-push
|
||||
check uses instead of the local `ConnectionManager`. A user is connected
|
||||
while any of their entries is unexpired, so two tabs (or two instances)
|
||||
don't mark them offline until every connection is gone. `GlobalPresence`
|
||||
(`presence:v2:global`) and `FocusPresence` (`presence:v2:unfocused`) use the
|
||||
same store, `app/ws/expiring_presence.py` (`ExpiringPresenceSet`).
|
||||
|
||||
Known limitation: `Presence` has no heartbeat/TTL, so a hard process crash
|
||||
(not a clean disconnect) leaks that connection's increment forever — same
|
||||
category of simplification as the "no server-side session revocation" note
|
||||
below.
|
||||
Presence entries expire (#77) so nothing can be stuck online. A connection
|
||||
renews its entries (90s TTL) whenever the client sends a frame — the frontend
|
||||
answers the server's 30s `ping` with a `pong`, which is exactly that — and
|
||||
the server closes a human client's socket (code 4408) after 90s of silence.
|
||||
Bot (API-token) connections don't speak the heartbeat, so the server renews
|
||||
their entries itself while their handler is alive. An expiry has no
|
||||
disconnect event to announce it, so `app/ws/presence_sweeper.py` runs on every
|
||||
worker (a Redis lock makes it one sweep per 15s overall), prunes expired
|
||||
entries, and broadcasts the offline transition. The pre-#77 `presence:*`
|
||||
refcount hashes are deleted at startup.
|
||||
|
||||
## Push notifications (Phase 4)
|
||||
|
||||
|
||||
+21
-3
@@ -1,5 +1,6 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import pathlib
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
@@ -11,6 +12,7 @@ from redis.asyncio import Redis
|
||||
from starlette.middleware.sessions import SessionMiddleware
|
||||
|
||||
from app.config import settings
|
||||
from app.database import async_session_factory
|
||||
from app.routers import (
|
||||
admin,
|
||||
auth,
|
||||
@@ -31,6 +33,9 @@ from app.ws.connection_manager import ConnectionManager
|
||||
from app.ws.focus_presence import FocusPresence
|
||||
from app.ws.global_presence import GlobalPresence
|
||||
from app.ws.presence import Presence
|
||||
from app.ws.presence_sweeper import purge_legacy_presence, run_presence_sweeper
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# backend/app/main.py -> backend/ -> repo root -- matches both the local
|
||||
# monorepo layout and the production layout (/srv/ds-chat/backend,
|
||||
@@ -65,11 +70,23 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
# restarts on failure instead.
|
||||
listener_task = asyncio.create_task(app.state.broadcaster.run_forever())
|
||||
|
||||
# #77: presence entries now expire, and an expiry has no disconnect event
|
||||
# to announce it -- this loop is what notices. The purge clears the old
|
||||
# refcount hashes (and the counts that had leaked into them); it's
|
||||
# best-effort, since Redis being down at startup shouldn't stop the app.
|
||||
try:
|
||||
await purge_legacy_presence(app.state.redis)
|
||||
except Exception:
|
||||
logger.warning("Could not purge legacy presence keys", exc_info=True)
|
||||
sweeper_task = asyncio.create_task(run_presence_sweeper(app))
|
||||
|
||||
yield
|
||||
|
||||
listener_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await listener_task
|
||||
for task in (sweeper_task, listener_task):
|
||||
task.cancel()
|
||||
for task in (sweeper_task, listener_task):
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
await app.state.redis.aclose()
|
||||
|
||||
|
||||
@@ -85,6 +102,7 @@ def create_app() -> FastAPI:
|
||||
)
|
||||
|
||||
app.state.connection_manager = ConnectionManager()
|
||||
app.state.session_factory = async_session_factory
|
||||
app.state.redis = Redis.from_url(settings.redis_url, decode_responses=True)
|
||||
app.state.presence = Presence(app.state.redis)
|
||||
app.state.global_presence = GlobalPresence(app.state.redis)
|
||||
|
||||
+87
-10
@@ -1,4 +1,7 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
|
||||
@@ -28,6 +31,8 @@ from app.services.message_service import (
|
||||
from app.services.room_service import mark_room_read
|
||||
from app.services.session_service import resolve_session
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(tags=["ws"])
|
||||
|
||||
WS_UNAUTHENTICATED = 4401
|
||||
@@ -44,6 +49,22 @@ WS_UNAUTHENTICATED = 4401
|
||||
# itself actually recovering).
|
||||
WS_PING_INTERVAL_SECONDS = 30
|
||||
|
||||
# #77: a human client that has sent nothing -- not even the "pong" to the ping
|
||||
# above -- for this long is treated as gone and disconnected, so its presence
|
||||
# stops being renewed and cleanup runs. Needed on top of the transport's own
|
||||
# keepalive: a browser answers *protocol-level* pings from its network layer
|
||||
# even while the page's JavaScript is frozen (a suspended or discarded tab,
|
||||
# a sleeping laptop that keeps its network up), so the connection looks
|
||||
# healthy forever while nobody is actually there. Only the page itself can
|
||||
# send the app-level "pong". The frontend reconnects on any non-4401 close.
|
||||
WS_DEAD_AFTER_SECONDS = 90
|
||||
WS_HEARTBEAT_TIMEOUT = 4408
|
||||
|
||||
# Presence entries (see ExpiringPresenceSet) are renewed on incoming frames,
|
||||
# at most this often -- a chatty client shouldn't turn every message into
|
||||
# several extra Redis writes.
|
||||
PRESENCE_REFRESH_MIN_INTERVAL_SECONDS = 10
|
||||
|
||||
|
||||
class ClientEnvelope(BaseModel):
|
||||
type: str
|
||||
@@ -74,6 +95,21 @@ def _missing_scope(api_token: ApiToken | None, scope: str) -> bool:
|
||||
return api_token is not None and scope not in api_token.scopes
|
||||
|
||||
|
||||
async def _broadcast_offline(db: AsyncSession, broadcaster, user_id: uuid.UUID) -> None:
|
||||
await broadcast_member_updated(db, broadcaster, user_id)
|
||||
await broadcast_dm_presence_update(db, broadcaster, user_id, online=False)
|
||||
|
||||
|
||||
async def _best_effort(label: str, awaitable):
|
||||
# Disconnect cleanup is several independent Redis/DB steps; one failing
|
||||
# must not skip the ones after it.
|
||||
try:
|
||||
return await awaitable
|
||||
except Exception:
|
||||
logger.warning("Presence cleanup step failed: %s", label, exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
@router.websocket("/ws/chat")
|
||||
async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)) -> None:
|
||||
api_token: ApiToken | None = None
|
||||
@@ -114,13 +150,19 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
||||
# refcount if the socket closes while still blurred -- mirroring how
|
||||
# joined_rooms tracks per-connection room membership for its own cleanup.
|
||||
is_blurred = False
|
||||
# Identifies this connection's own presence entries (#77) -- see
|
||||
# ExpiringPresenceSet -- so one connection ending can't affect another
|
||||
# belonging to the same user, and a dead one can simply stop renewing.
|
||||
conn_id = uuid.uuid4().hex
|
||||
last_frame_at = time.monotonic()
|
||||
last_presence_refresh = last_frame_at
|
||||
manager.register_user(user.id, websocket)
|
||||
# Only broadcast on a genuine offline->online transition (this user's
|
||||
# first open connection), not for every extra tab -- broadcast_member_
|
||||
# updated tells every room this user's in to refresh, which would be
|
||||
# wasted churn on a transition that didn't actually change anything
|
||||
# visible.
|
||||
if await global_presence.connect(user.id):
|
||||
if await global_presence.connect(user.id, conn_id):
|
||||
await broadcast_member_updated(db, broadcaster, user.id)
|
||||
await broadcast_dm_presence_update(db, broadcaster, user.id, online=True)
|
||||
# This session is shared for the connection's entire lifetime (which can
|
||||
@@ -135,6 +177,18 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
||||
# waiting for the next one.
|
||||
await db.commit()
|
||||
|
||||
async def refresh_presence() -> None:
|
||||
# Renews every presence entry this connection holds. If the global
|
||||
# one had already lapsed (a long stall) this is also the moment the
|
||||
# user comes back online, and has to be announced like a fresh connect.
|
||||
if await global_presence.refresh(user.id, conn_id):
|
||||
await broadcast_member_updated(db, broadcaster, user.id)
|
||||
await broadcast_dm_presence_update(db, broadcaster, user.id, online=True)
|
||||
for room_id in joined_rooms:
|
||||
await presence.refresh(room_id, user.id, conn_id)
|
||||
if is_blurred:
|
||||
await focus_presence.refresh(user.id, conn_id)
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
@@ -147,9 +201,30 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
||||
# send (raising, caught below same as any other disconnect);
|
||||
# a merely quiet one just gets proded and the loop continues
|
||||
# waiting.
|
||||
if api_token is None and time.monotonic() - last_frame_at >= WS_DEAD_AFTER_SECONDS:
|
||||
# Gone (see WS_DEAD_AFTER_SECONDS). The socket may be too
|
||||
# dead to even take a close frame; either way, leaving the
|
||||
# loop runs the normal cleanup below.
|
||||
with contextlib.suppress(Exception):
|
||||
await websocket.close(code=WS_HEARTBEAT_TIMEOUT)
|
||||
break
|
||||
if api_token is not None:
|
||||
# A bot isn't expected to speak our app-level heartbeat
|
||||
# (it may just ignore the ping), so it can't be dropped
|
||||
# for not answering and its presence can't depend on
|
||||
# answers. While the handler is alive the transport's own
|
||||
# keepalive is what vouches for it -- a dead bot socket
|
||||
# ends this handler, and cleanup does the rest.
|
||||
await refresh_presence()
|
||||
last_presence_refresh = time.monotonic()
|
||||
await db.commit()
|
||||
await websocket.send_json({"type": "ping"})
|
||||
continue
|
||||
last_frame_at = time.monotonic()
|
||||
try:
|
||||
if last_frame_at - last_presence_refresh >= PRESENCE_REFRESH_MIN_INTERVAL_SECONDS:
|
||||
await refresh_presence()
|
||||
last_presence_refresh = last_frame_at
|
||||
try:
|
||||
envelope = ClientEnvelope.model_validate(raw)
|
||||
except ValidationError:
|
||||
@@ -166,7 +241,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
||||
)
|
||||
continue
|
||||
manager.join(envelope.room_id, websocket)
|
||||
await presence.join(envelope.room_id, user.id)
|
||||
await presence.join(envelope.room_id, user.id, conn_id)
|
||||
joined_rooms.add(envelope.room_id)
|
||||
await websocket.send_json({"type": "joined", "room_id": str(envelope.room_id)})
|
||||
|
||||
@@ -183,7 +258,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
||||
await websocket.send_json({"type": "error", "detail": "room_id required"})
|
||||
continue
|
||||
manager.leave(envelope.room_id, websocket)
|
||||
await presence.leave(envelope.room_id, user.id)
|
||||
await presence.leave(envelope.room_id, user.id, conn_id)
|
||||
joined_rooms.discard(envelope.room_id)
|
||||
|
||||
elif envelope.type == "focus":
|
||||
@@ -198,11 +273,11 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
||||
continue
|
||||
if envelope.focused:
|
||||
if is_blurred:
|
||||
await focus_presence.mark_focused(user.id)
|
||||
await focus_presence.mark_focused(user.id, conn_id)
|
||||
is_blurred = False
|
||||
else:
|
||||
if not is_blurred:
|
||||
await focus_presence.mark_blurred(user.id)
|
||||
await focus_presence.mark_blurred(user.id, conn_id)
|
||||
is_blurred = True
|
||||
|
||||
elif envelope.type == "message":
|
||||
@@ -396,9 +471,11 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
||||
manager.leave_all(websocket)
|
||||
manager.unregister_user(user.id, websocket)
|
||||
for room_id in joined_rooms:
|
||||
await presence.leave(room_id, user.id)
|
||||
await _best_effort("room presence", presence.leave(room_id, user.id, conn_id))
|
||||
if is_blurred:
|
||||
await focus_presence.mark_focused(user.id)
|
||||
if await global_presence.disconnect(user.id):
|
||||
await broadcast_member_updated(db, broadcaster, user.id)
|
||||
await broadcast_dm_presence_update(db, broadcaster, user.id, online=False)
|
||||
await _best_effort("focus presence", focus_presence.mark_focused(user.id, conn_id))
|
||||
if await _best_effort("global presence", global_presence.disconnect(user.id, conn_id)):
|
||||
await _best_effort(
|
||||
"offline broadcast",
|
||||
_broadcast_offline(db, broadcaster, user.id),
|
||||
)
|
||||
@@ -0,0 +1,103 @@
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
# #77: an entry that isn't renewed for this long stops counting. Chosen as 3x
|
||||
# the app-level ping interval (chat.py's WS_PING_INTERVAL_SECONDS): a healthy
|
||||
# connection renews roughly every ping, so two missed beats are tolerated
|
||||
# before someone flips offline.
|
||||
PRESENCE_TTL_SECONDS = 90
|
||||
|
||||
|
||||
def _text(value: str | bytes) -> str:
|
||||
return value.decode() if isinstance(value, bytes) else value
|
||||
|
||||
|
||||
def _user_part(member: str) -> str:
|
||||
return member.split(":", 1)[0]
|
||||
|
||||
|
||||
class ExpiringPresenceSet:
|
||||
""""Which users have at least one live connection" for one Redis key,
|
||||
where every connection is its own entry with its own expiry.
|
||||
|
||||
This replaced per-user refcounts (#77): a refcount only goes down when a
|
||||
connection's cleanup code actually runs, so any connection that ended
|
||||
without it -- a server crash or reboot, a process killed mid-restart, a
|
||||
client that vanished without the server ever noticing -- left a +1 that
|
||||
nothing could ever subtract, showing that person online forever (and in
|
||||
the per-room case, suppressing their push notifications). Here a dead
|
||||
connection simply stops being renewed and ages out on its own, so there
|
||||
is no state that can leak.
|
||||
|
||||
One sorted set per key: member `<user_id>:<connection_id>`, score = the
|
||||
moment it expires. "Live" means score > now. Scans are O(connections) --
|
||||
fine for a self-hosted team chat; only the global key grows with the
|
||||
whole server rather than one room.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis: Redis,
|
||||
ttl_seconds: float = PRESENCE_TTL_SECONDS,
|
||||
clock: Callable[[], float] = time.time,
|
||||
) -> None:
|
||||
self._redis = redis
|
||||
self._ttl = ttl_seconds
|
||||
self._clock = clock
|
||||
|
||||
@staticmethod
|
||||
def _member(user_id: uuid.UUID, conn_id: str) -> str:
|
||||
return f"{user_id}:{conn_id}"
|
||||
|
||||
async def _live_members(self, key: str) -> list[str]:
|
||||
members = await self._redis.zrangebyscore(key, f"({self._clock()}", "+inf")
|
||||
return [_text(m) for m in members]
|
||||
|
||||
async def add(self, key: str, user_id: uuid.UUID, conn_id: str) -> bool:
|
||||
"""Creates or renews this connection's entry. Returns True iff the
|
||||
user had no live entry beforehand -- a genuine offline->online
|
||||
transition worth telling anyone about. Renewing an entry that's still
|
||||
live (the common case, every heartbeat) is a single cheap write."""
|
||||
member = self._member(user_id, conn_id)
|
||||
now = self._clock()
|
||||
score = await self._redis.zscore(key, member)
|
||||
came_online = False
|
||||
if score is None or score <= now:
|
||||
came_online = not await self.is_live(key, user_id)
|
||||
await self._redis.zadd(key, {member: now + self._ttl})
|
||||
# Lets a key nobody renews any more (an idle room) disappear on its
|
||||
# own instead of sitting there holding only expired entries.
|
||||
await self._redis.expire(key, int(self._ttl * 2))
|
||||
return came_online
|
||||
|
||||
async def remove(self, key: str, user_id: uuid.UUID, conn_id: str) -> bool:
|
||||
"""Drops this connection's entry. Returns True iff that left the user
|
||||
with no live entry -- a genuine online->offline transition. False if
|
||||
the entry was already gone: whoever removed it (see prune) already
|
||||
reported the transition, and reporting it twice would be noise."""
|
||||
removed = await self._redis.zrem(key, self._member(user_id, conn_id))
|
||||
if not removed:
|
||||
return False
|
||||
return not await self.is_live(key, user_id)
|
||||
|
||||
async def is_live(self, key: str, user_id: uuid.UUID) -> bool:
|
||||
prefix = f"{user_id}:"
|
||||
return any(m.startswith(prefix) for m in await self._live_members(key))
|
||||
|
||||
async def live_user_ids(self, key: str) -> set[uuid.UUID]:
|
||||
return {uuid.UUID(_user_part(m)) for m in await self._live_members(key)}
|
||||
|
||||
async def prune(self, key: str) -> set[uuid.UUID]:
|
||||
"""Deletes expired entries. Returns the users whose last live entry
|
||||
was among them -- the offline transitions nobody was told about,
|
||||
since an expiry has no disconnect event to hang a broadcast on."""
|
||||
now = self._clock()
|
||||
expired = [_text(m) for m in await self._redis.zrangebyscore(key, "-inf", now)]
|
||||
if not expired:
|
||||
return set()
|
||||
await self._redis.zremrangebyscore(key, "-inf", now)
|
||||
gone = {_user_part(m) for m in expired} - {_user_part(m) for m in await self._live_members(key)}
|
||||
return {uuid.UUID(u) for u in gone}
|
||||
@@ -2,44 +2,45 @@ import uuid
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from app.ws.expiring_presence import ExpiringPresenceSet
|
||||
|
||||
|
||||
class FocusPresence:
|
||||
"""Cross-instance "is this desktop-mode user's window currently
|
||||
unfocused," used only to widen desktop_notification/push eligibility
|
||||
beyond plain room-connection state (#59). A Redis hash (field =
|
||||
user_id, value = refcount of that user's currently-blurred desktop
|
||||
connections), parallel to Presence/GlobalPresence.
|
||||
beyond plain room-connection state (#59). One expiring entry per
|
||||
currently-blurred desktop connection (see ExpiringPresenceSet),
|
||||
parallel to Presence/GlobalPresence.
|
||||
|
||||
Absence from this hash is the default and means "focused." That default
|
||||
Absence from the set is the default and means "focused." That default
|
||||
is also exactly right for every browser-tab connection: only the
|
||||
desktop client ever sends focus/blur frames at all (see chat.py's
|
||||
"focus" envelope handling), so a browser user never appears here --
|
||||
their attention is already fully captured by Presence's room-connection
|
||||
state, which stays visibility-gated with no separate focus signal.
|
||||
|
||||
Refcounted for the same multi-connection reason as Presence/
|
||||
GlobalPresence, with the same known simplification: two desktop windows
|
||||
for one user, one focused and one blurred, count as "unfocused" here
|
||||
(refcount > 0) even though the user does have attention somewhere. That
|
||||
errs toward notifying rather than silently missing one, which is the
|
||||
safer failure mode for a notification.
|
||||
Same known simplification as before: two desktop windows for one user,
|
||||
one focused and one blurred, count as "unfocused" (any live blurred
|
||||
entry) even though the user does have attention somewhere. That errs
|
||||
toward notifying rather than silently missing one, which is the safer
|
||||
failure mode for a notification. Entries expire (#77) -- a desktop
|
||||
window that died while blurred used to leave its user "unfocused"
|
||||
forever.
|
||||
"""
|
||||
|
||||
def __init__(self, redis: Redis) -> None:
|
||||
self._redis = redis
|
||||
KEY = "presence:v2:unfocused"
|
||||
|
||||
def _key(self) -> str:
|
||||
return "presence:unfocused"
|
||||
def __init__(self, redis: Redis, store: ExpiringPresenceSet | None = None) -> None:
|
||||
self._store = store or ExpiringPresenceSet(redis)
|
||||
|
||||
async def mark_blurred(self, user_id: uuid.UUID) -> None:
|
||||
await self._redis.hincrby(self._key(), str(user_id), 1)
|
||||
async def mark_blurred(self, user_id: uuid.UUID, conn_id: str) -> None:
|
||||
await self._store.add(self.KEY, user_id, conn_id)
|
||||
|
||||
async def mark_focused(self, user_id: uuid.UUID) -> None:
|
||||
key = self._key()
|
||||
field = str(user_id)
|
||||
remaining = await self._redis.hincrby(key, field, -1)
|
||||
if remaining <= 0:
|
||||
await self._redis.hdel(key, field)
|
||||
async def mark_focused(self, user_id: uuid.UUID, conn_id: str) -> None:
|
||||
await self._store.remove(self.KEY, user_id, conn_id)
|
||||
|
||||
async def refresh(self, user_id: uuid.UUID, conn_id: str) -> None:
|
||||
await self._store.add(self.KEY, user_id, conn_id)
|
||||
|
||||
async def is_unfocused(self, user_id: uuid.UUID) -> bool:
|
||||
return await self._redis.hexists(self._key(), str(user_id))
|
||||
return await self._store.is_live(self.KEY, user_id)
|
||||
@@ -2,50 +2,55 @@ import uuid
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from app.ws.expiring_presence import ExpiringPresenceSet
|
||||
|
||||
|
||||
class GlobalPresence:
|
||||
"""Cross-instance "does this user have the app open at all right now,"
|
||||
independent of which (if any) room they currently have open -- backing
|
||||
the online/offline presence dot shown wherever a user's avatar renders.
|
||||
A single Redis hash (field = user_id, value = connection refcount),
|
||||
parallel to but separate from Presence's per-room hashes.
|
||||
|
||||
Refcounted for the same reason as Presence: multiple tabs/instances for
|
||||
one user shouldn't flip them offline until the last connection closes.
|
||||
Each connection has its own expiring entry (see ExpiringPresenceSet), so
|
||||
a user with several tabs/devices stays online until the last one goes,
|
||||
and a connection that dies without cleaning up ages out instead of
|
||||
leaving them online forever (#77). The key is versioned (`v2`) because
|
||||
the old refcount hash lives at `presence:global` and its leaked counts
|
||||
must not be read as live entries.
|
||||
"""
|
||||
|
||||
def __init__(self, redis: Redis) -> None:
|
||||
self._redis = redis
|
||||
KEY = "presence:v2:global"
|
||||
|
||||
def _key(self) -> str:
|
||||
return "presence:global"
|
||||
def __init__(self, redis: Redis, store: ExpiringPresenceSet | None = None) -> None:
|
||||
self._store = store or ExpiringPresenceSet(redis)
|
||||
|
||||
async def connect(self, user_id: uuid.UUID) -> bool:
|
||||
"""Returns True iff this was the user's first open connection --
|
||||
a genuine offline->online transition worth telling anyone about."""
|
||||
count = await self._redis.hincrby(self._key(), str(user_id), 1)
|
||||
return count == 1
|
||||
async def connect(self, user_id: uuid.UUID, conn_id: str) -> bool:
|
||||
"""Returns True iff this made the user newly online -- a genuine
|
||||
offline->online transition worth telling anyone about."""
|
||||
return await self._store.add(self.KEY, user_id, conn_id)
|
||||
|
||||
async def disconnect(self, user_id: uuid.UUID) -> bool:
|
||||
"""Returns True iff this was the user's last open connection -- a
|
||||
async def refresh(self, user_id: uuid.UUID, conn_id: str) -> bool:
|
||||
"""Heartbeat. Same as connect() on purpose: a connection whose entry
|
||||
had already expired (a long stall) comes back online here, and the
|
||||
caller needs to announce that exactly like a fresh connect."""
|
||||
return await self._store.add(self.KEY, user_id, conn_id)
|
||||
|
||||
async def disconnect(self, user_id: uuid.UUID, conn_id: str) -> bool:
|
||||
"""Returns True iff that was the user's last live connection -- a
|
||||
genuine online->offline transition."""
|
||||
key = self._key()
|
||||
field = str(user_id)
|
||||
remaining = await self._redis.hincrby(key, field, -1)
|
||||
if remaining <= 0:
|
||||
await self._redis.hdel(key, field)
|
||||
return True
|
||||
return False
|
||||
return await self._store.remove(self.KEY, user_id, conn_id)
|
||||
|
||||
async def sweep(self) -> set[uuid.UUID]:
|
||||
"""Drops expired entries and returns the users that went offline as a
|
||||
result -- the transitions no disconnect ever reported."""
|
||||
return await self._store.prune(self.KEY)
|
||||
|
||||
async def is_online(self, user_id: uuid.UUID) -> bool:
|
||||
return await self._redis.hexists(self._key(), str(user_id))
|
||||
return await self._store.is_live(self.KEY, user_id)
|
||||
|
||||
async def online_user_ids(self, user_ids: list[uuid.UUID]) -> set[uuid.UUID]:
|
||||
if not user_ids:
|
||||
return set()
|
||||
values = await self._redis.hmget(self._key(), [str(u) for u in user_ids])
|
||||
return {uid for uid, v in zip(user_ids, values) if v is not None}
|
||||
return set(user_ids) & await self._store.live_user_ids(self.KEY)
|
||||
|
||||
async def all_online_user_ids(self) -> set[uuid.UUID]:
|
||||
fields = await self._redis.hkeys(self._key())
|
||||
return {uuid.UUID(f) for f in fields}
|
||||
return await self._store.live_user_ids(self.KEY)
|
||||
+26
-24
@@ -2,38 +2,40 @@ import uuid
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from app.ws.expiring_presence import ExpiringPresenceSet
|
||||
|
||||
|
||||
class Presence:
|
||||
"""Cross-instance "who's connected to this room," backed by a Redis hash
|
||||
per room (field = user_id, value = connection refcount).
|
||||
"""Cross-instance "who's connected to this room," one expiring set per
|
||||
room (see ExpiringPresenceSet).
|
||||
|
||||
Refcounted rather than a plain set so a user with two connections to the
|
||||
same room -- two tabs, or one per app instance -- doesn't get marked
|
||||
offline when only one of those connections closes.
|
||||
|
||||
Known limitation: a hard crash (not a clean disconnect) leaks that
|
||||
connection's increment forever, since there's no heartbeat/TTL here to
|
||||
reclaim it -- out of scope for this phase, same category of
|
||||
simplification as the "no server-side session revocation" note in the
|
||||
README.
|
||||
This is what the offline-push check reads, so a leaked entry used to mean
|
||||
a user who'd gone away was still treated as "watching the room" and never
|
||||
got a push notification (#77). Entries now expire unless the connection
|
||||
keeps renewing them, so that can't outlive a dead connection by more than
|
||||
the TTL. Keys are versioned (`v2`) so the old refcount hashes
|
||||
(`presence:{room_id}`), with whatever they leaked, are never read.
|
||||
"""
|
||||
|
||||
def __init__(self, redis: Redis) -> None:
|
||||
self._redis = redis
|
||||
def __init__(self, redis: Redis, store: ExpiringPresenceSet | None = None) -> None:
|
||||
self._store = store or ExpiringPresenceSet(redis)
|
||||
|
||||
def _key(self, room_id: uuid.UUID) -> str:
|
||||
return f"presence:{room_id}"
|
||||
return f"presence:v2:room:{room_id}"
|
||||
|
||||
async def join(self, room_id: uuid.UUID, user_id: uuid.UUID) -> None:
|
||||
await self._redis.hincrby(self._key(room_id), str(user_id), 1)
|
||||
|
||||
async def leave(self, room_id: uuid.UUID, user_id: uuid.UUID) -> None:
|
||||
async def join(self, room_id: uuid.UUID, user_id: uuid.UUID, conn_id: str) -> None:
|
||||
key = self._key(room_id)
|
||||
field = str(user_id)
|
||||
remaining = await self._redis.hincrby(key, field, -1)
|
||||
if remaining <= 0:
|
||||
await self._redis.hdel(key, field)
|
||||
# Per-room offline transitions aren't broadcast anywhere, so unlike
|
||||
# the global key there's nothing to lose by pruning here -- it just
|
||||
# keeps a busy room's set from accumulating expired entries.
|
||||
await self._store.prune(key)
|
||||
await self._store.add(key, user_id, conn_id)
|
||||
|
||||
async def refresh(self, room_id: uuid.UUID, user_id: uuid.UUID, conn_id: str) -> None:
|
||||
await self._store.add(self._key(room_id), user_id, conn_id)
|
||||
|
||||
async def leave(self, room_id: uuid.UUID, user_id: uuid.UUID, conn_id: str) -> None:
|
||||
await self._store.remove(self._key(room_id), user_id, conn_id)
|
||||
|
||||
async def connected_user_ids(self, room_id: uuid.UUID) -> set[uuid.UUID]:
|
||||
fields = await self._redis.hkeys(self._key(room_id))
|
||||
return {uuid.UUID(f) for f in fields}
|
||||
return await self._store.live_user_ids(self._key(room_id))
|
||||
@@ -0,0 +1,83 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
|
||||
from redis.asyncio import Redis
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.services.message_events import broadcast_dm_presence_update, broadcast_member_updated
|
||||
from app.ws.broadcaster import Broadcaster
|
||||
from app.ws.global_presence import GlobalPresence
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PRESENCE_SWEEP_INTERVAL_SECONDS = 15
|
||||
SWEEP_LOCK_KEY = "presence:v2:sweeper-lock"
|
||||
|
||||
|
||||
async def sweep_once(
|
||||
redis: Redis,
|
||||
global_presence: GlobalPresence,
|
||||
broadcaster: Broadcaster,
|
||||
session_factory: Callable[[], AsyncSession],
|
||||
lock_seconds: int | None = max(1, PRESENCE_SWEEP_INTERVAL_SECONDS - 1),
|
||||
) -> set[uuid.UUID]:
|
||||
"""Expires dead connections' presence entries and tells everyone about
|
||||
the users that went offline as a result.
|
||||
|
||||
A clean disconnect broadcasts its own offline transition (chat.py), but a
|
||||
connection that just stops being renewed has no such event -- this is the
|
||||
only thing that notices. Every app worker runs this loop; the Redis lock
|
||||
makes it one sweep per interval across the whole deployment rather than
|
||||
N workers each doing (and broadcasting) the same thing. `lock_seconds=None`
|
||||
skips the lock (tests).
|
||||
"""
|
||||
if lock_seconds is not None and not await redis.set(SWEEP_LOCK_KEY, "1", nx=True, ex=lock_seconds):
|
||||
return set()
|
||||
went_offline = await global_presence.sweep()
|
||||
for user_id in went_offline:
|
||||
# One user's failed broadcast must not cost everyone after them
|
||||
# theirs -- their entry is already gone, so this is the only chance.
|
||||
try:
|
||||
async with session_factory() as db:
|
||||
await broadcast_member_updated(db, broadcaster, user_id)
|
||||
await broadcast_dm_presence_update(db, broadcaster, user_id, online=False)
|
||||
except Exception:
|
||||
logger.error("Failed to broadcast offline for %s", user_id, exc_info=True)
|
||||
return went_offline
|
||||
|
||||
|
||||
async def run_presence_sweeper(app, interval: float = PRESENCE_SWEEP_INTERVAL_SECONDS) -> None:
|
||||
# Same lesson as Broadcaster.run_forever (#76): a background task that
|
||||
# dies on its first error does so silently, and here that would quietly
|
||||
# bring the stuck-online bug back.
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
try:
|
||||
await sweep_once(
|
||||
app.state.redis,
|
||||
app.state.global_presence,
|
||||
app.state.broadcaster,
|
||||
app.state.session_factory,
|
||||
lock_seconds=max(1, int(interval) - 1),
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.error("Presence sweep failed", exc_info=True)
|
||||
|
||||
|
||||
async def purge_legacy_presence(redis: Redis) -> int:
|
||||
"""Deletes the old refcount hashes (`presence:global`, `presence:unfocused`,
|
||||
`presence:{room_id}`). Nothing reads them any more -- the live data is
|
||||
under `presence:v2:` -- and they hold whatever counts leaked before #77,
|
||||
so there's nothing worth keeping. Safe to run on every startup."""
|
||||
deleted = 0
|
||||
async for key in redis.scan_iter("presence:*"):
|
||||
name = key.decode() if isinstance(key, bytes) else key
|
||||
if name.startswith("presence:v2:"):
|
||||
continue
|
||||
await redis.delete(name)
|
||||
deleted += 1
|
||||
return deleted
|
||||
@@ -0,0 +1,155 @@
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import pytest_asyncio
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from app.ws.expiring_presence import ExpiringPresenceSet
|
||||
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = time.time()
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.now += seconds
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def redis_client():
|
||||
client = Redis.from_url(os.environ["REDIS_URL"], decode_responses=True)
|
||||
yield client
|
||||
await client.aclose()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def setup(redis_client):
|
||||
clock = Clock()
|
||||
store = ExpiringPresenceSet(redis_client, ttl_seconds=90, clock=clock)
|
||||
key = f"presence:v2:test:{uuid.uuid4()}"
|
||||
yield store, clock, key
|
||||
await redis_client.delete(key)
|
||||
|
||||
|
||||
async def test_first_connection_is_a_transition_and_a_second_is_not(setup):
|
||||
store, _, key = setup
|
||||
user = uuid.uuid4()
|
||||
assert await store.add(key, user, "tab1") is True
|
||||
assert await store.add(key, user, "tab2") is False
|
||||
assert await store.live_user_ids(key) == {user}
|
||||
|
||||
|
||||
async def test_renewing_a_live_entry_is_not_a_transition(setup):
|
||||
store, clock, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
clock.advance(30)
|
||||
assert await store.add(key, user, "tab1") is False
|
||||
|
||||
|
||||
async def test_entry_stops_counting_after_the_ttl(setup):
|
||||
store, clock, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
clock.advance(89)
|
||||
assert await store.is_live(key, user) is True
|
||||
clock.advance(2)
|
||||
assert await store.is_live(key, user) is False
|
||||
assert await store.live_user_ids(key) == set()
|
||||
|
||||
|
||||
async def test_renewal_keeps_a_connection_alive_past_the_original_expiry(setup):
|
||||
store, clock, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
for _ in range(5):
|
||||
clock.advance(60)
|
||||
await store.add(key, user, "tab1")
|
||||
assert await store.is_live(key, user) is True
|
||||
|
||||
|
||||
async def test_coming_back_after_expiry_is_a_transition_again(setup):
|
||||
store, clock, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
clock.advance(200)
|
||||
assert await store.add(key, user, "tab1") is True
|
||||
|
||||
|
||||
async def test_one_live_connection_keeps_the_user_online_when_another_expires(setup):
|
||||
store, clock, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "old-tab")
|
||||
clock.advance(60)
|
||||
await store.add(key, user, "new-tab")
|
||||
clock.advance(60) # old-tab is now 120s stale, new-tab 60s
|
||||
assert await store.is_live(key, user) is True
|
||||
|
||||
|
||||
async def test_remove_reports_only_the_last_connection_leaving(setup):
|
||||
store, _, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
await store.add(key, user, "tab2")
|
||||
assert await store.remove(key, user, "tab1") is False
|
||||
assert await store.remove(key, user, "tab2") is True
|
||||
|
||||
|
||||
async def test_remove_of_an_entry_that_is_already_gone_reports_nothing(setup):
|
||||
# Whoever removed it first (e.g. prune) already reported the transition.
|
||||
store, _, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
await store.remove(key, user, "tab1")
|
||||
assert await store.remove(key, user, "tab1") is False
|
||||
|
||||
|
||||
async def test_remove_of_an_expired_but_unpruned_entry_still_reports_offline(setup):
|
||||
# Nobody has been told this user went offline yet (no prune has run), so
|
||||
# their eventual clean close is the one place left to report it.
|
||||
store, clock, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
clock.advance(200)
|
||||
assert await store.remove(key, user, "tab1") is True
|
||||
|
||||
|
||||
async def test_users_are_independent(setup):
|
||||
store, _, key = setup
|
||||
alice, bob = uuid.uuid4(), uuid.uuid4()
|
||||
await store.add(key, alice, "a")
|
||||
await store.add(key, bob, "b")
|
||||
await store.remove(key, alice, "a")
|
||||
assert await store.live_user_ids(key) == {bob}
|
||||
|
||||
|
||||
async def test_prune_reports_users_who_lost_their_last_live_entry(setup):
|
||||
store, clock, key = setup
|
||||
gone, still_here = uuid.uuid4(), uuid.uuid4()
|
||||
await store.add(key, gone, "tab1")
|
||||
await store.add(key, still_here, "old")
|
||||
clock.advance(60)
|
||||
await store.add(key, still_here, "new")
|
||||
clock.advance(40) # gone/tab1 and still_here/old are expired; still_here/new is not
|
||||
assert await store.prune(key) == {gone}
|
||||
|
||||
|
||||
async def test_prune_is_idempotent_and_removes_what_it_reports(setup):
|
||||
store, clock, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
clock.advance(200)
|
||||
assert await store.prune(key) == {user}
|
||||
assert await store.prune(key) == set()
|
||||
|
||||
|
||||
async def test_prune_with_nothing_expired_does_nothing(setup):
|
||||
store, _, key = setup
|
||||
user = uuid.uuid4()
|
||||
await store.add(key, user, "tab1")
|
||||
assert await store.prune(key) == set()
|
||||
assert await store.is_live(key, user) is True
|
||||
@@ -130,3 +130,108 @@ def test_online_users_endpoint_respects_appear_offline(ws_client):
|
||||
|
||||
online = ws_client.get("/api/users/online").json()
|
||||
assert alice["id"] not in online
|
||||
|
||||
|
||||
# --- #77: liveness -----------------------------------------------------------
|
||||
# A connection that stops answering the app-level heartbeat is dropped by the
|
||||
# server, and a connection that does answer keeps its presence entry renewed.
|
||||
# Timings are shrunk so nothing here waits on the real 30s/90s values.
|
||||
|
||||
|
||||
def _fast_heartbeat(monkeypatch):
|
||||
monkeypatch.setattr("app.ws.chat.WS_PING_INTERVAL_SECONDS", 0.05)
|
||||
monkeypatch.setattr("app.ws.chat.WS_DEAD_AFTER_SECONDS", 0.2)
|
||||
monkeypatch.setattr("app.ws.chat.PRESENCE_REFRESH_MIN_INTERVAL_SECONDS", 0)
|
||||
|
||||
|
||||
def test_unresponsive_client_is_disconnected_by_the_server(ws_client, monkeypatch):
|
||||
# The frozen-page / sleeping-laptop case: the transport stays open (the
|
||||
# browser answers protocol pings itself) but the page never replies to
|
||||
# the app-level ping.
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
_fast_heartbeat(monkeypatch)
|
||||
_register_ws(ws_client, _unique("alice"))
|
||||
|
||||
with ws_client.websocket_connect("/ws/chat") as ws:
|
||||
pings = 0
|
||||
try:
|
||||
while True:
|
||||
assert ws.receive_json() == {"type": "ping"}
|
||||
pings += 1
|
||||
assert pings < 100, "server never gave up on a silent client"
|
||||
except WebSocketDisconnect as exc:
|
||||
assert exc.code == 4408
|
||||
assert pings >= 2
|
||||
|
||||
|
||||
def test_client_that_answers_pings_is_not_disconnected(ws_client, monkeypatch):
|
||||
_fast_heartbeat(monkeypatch)
|
||||
_register_ws(ws_client, _unique("alice"))
|
||||
room = ws_client.post("/api/rooms", json={"name": _unique("general")}).json()
|
||||
|
||||
with ws_client.websocket_connect("/ws/chat") as ws:
|
||||
# Far longer than WS_DEAD_AFTER_SECONDS in total.
|
||||
for _ in range(12):
|
||||
assert ws.receive_json() == {"type": "ping"}
|
||||
ws.send_json({"type": "pong"})
|
||||
ws.send_json({"type": "join", "room_id": room["id"]})
|
||||
while (msg := ws.receive_json())["type"] == "ping":
|
||||
pass
|
||||
assert msg == {"type": "joined", "room_id": room["id"]}
|
||||
|
||||
|
||||
def test_bot_connection_is_not_dropped_for_ignoring_pings(ws_client_factory, monkeypatch):
|
||||
# API-token clients aren't expected to speak the app heartbeat.
|
||||
_fast_heartbeat(monkeypatch)
|
||||
ws_client = ws_client_factory()
|
||||
admin = _register_ws(ws_client, _unique("admin"))
|
||||
|
||||
async def _promote():
|
||||
from app.models import User
|
||||
|
||||
async with ws_client.session_factory() as session:
|
||||
user = await session.get(User, uuid.UUID(admin["id"]))
|
||||
user.is_site_admin = True
|
||||
await session.commit()
|
||||
|
||||
ws_client.portal.call(_promote)
|
||||
room = ws_client.post("/api/rooms", json={"name": _unique("general")}).json()
|
||||
bot = ws_client.post("/api/admin/bots", json={"username": _unique("bot")}).json()
|
||||
token = ws_client.post(
|
||||
f"/api/admin/bots/{bot['id']}/tokens", json={"scopes": ["read:messages", "write:messages"]}
|
||||
).json()["token"]
|
||||
ws_client.post(f"/api/rooms/{room['id']}/join", headers={"Authorization": f"Bearer {token}"})
|
||||
|
||||
with ws_client.websocket_connect("/ws/chat", headers={"Authorization": f"Bearer {token}"}) as ws:
|
||||
for _ in range(8): # ~0.4s of silence, twice WS_DEAD_AFTER_SECONDS
|
||||
assert ws.receive_json() == {"type": "ping"}
|
||||
ws.send_json({"type": "join", "room_id": room["id"]})
|
||||
while (msg := ws.receive_json())["type"] == "ping":
|
||||
pass
|
||||
assert msg["type"] == "joined"
|
||||
|
||||
|
||||
def test_presence_entry_is_renewed_while_the_client_answers(ws_client, monkeypatch):
|
||||
import time
|
||||
|
||||
_fast_heartbeat(monkeypatch)
|
||||
alice = _register_ws(ws_client, _unique("alice"))
|
||||
|
||||
async def _scores():
|
||||
redis = ws_client.app.state.redis
|
||||
return await redis.zrange("presence:v2:global", 0, -1, withscores=True)
|
||||
|
||||
with ws_client.websocket_connect("/ws/chat") as ws:
|
||||
# Wait for the first ping so the connect-time entry exists.
|
||||
assert ws.receive_json() == {"type": "ping"}
|
||||
first = [s for m, s in ws_client.portal.call(_scores) if m.startswith(alice["id"])]
|
||||
assert len(first) == 1
|
||||
|
||||
time.sleep(0.1)
|
||||
ws.send_json({"type": "pong"})
|
||||
# Any later frame arriving back proves the pong was processed.
|
||||
while ws.receive_json()["type"] != "ping":
|
||||
pass
|
||||
second = [s for m, s in ws_client.portal.call(_scores) if m.startswith(alice["id"])]
|
||||
assert second and second[0] > first[0]
|
||||
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import pytest_asyncio
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from app.schemas.room import RoomCreate
|
||||
from app.schemas.user import UserCreate
|
||||
from app.services.auth_service import register_user
|
||||
from app.services.room_service import create_room, find_or_create_dm
|
||||
from app.ws.expiring_presence import ExpiringPresenceSet
|
||||
from app.ws.global_presence import GlobalPresence
|
||||
from app.ws.presence_sweeper import SWEEP_LOCK_KEY, purge_legacy_presence, sweep_once
|
||||
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = time.time()
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.now += seconds
|
||||
|
||||
|
||||
class FakeBroadcaster:
|
||||
def __init__(self) -> None:
|
||||
self.room_events: list[tuple[uuid.UUID, dict]] = []
|
||||
self.user_events: list[tuple[uuid.UUID, dict]] = []
|
||||
|
||||
async def publish(self, room_id, payload):
|
||||
self.room_events.append((room_id, payload))
|
||||
|
||||
async def publish_to_user(self, user_id, payload):
|
||||
self.user_events.append((user_id, payload))
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def redis_client():
|
||||
client = Redis.from_url(os.environ["REDIS_URL"], decode_responses=True)
|
||||
await client.delete(GlobalPresence.KEY, SWEEP_LOCK_KEY)
|
||||
yield client
|
||||
await client.delete(GlobalPresence.KEY, SWEEP_LOCK_KEY)
|
||||
await client.aclose()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def world(redis_client, db_session):
|
||||
clock = Clock()
|
||||
presence = GlobalPresence(redis_client, ExpiringPresenceSet(redis_client, ttl_seconds=90, clock=clock))
|
||||
broadcaster = FakeBroadcaster()
|
||||
|
||||
class _Factory:
|
||||
def __call__(self):
|
||||
return self
|
||||
|
||||
async def __aenter__(self):
|
||||
return db_session
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
async def make_user(prefix: str):
|
||||
name = f"{prefix}-{uuid.uuid4().hex[:8]}"
|
||||
return await register_user(
|
||||
db_session, UserCreate(username=name, email=f"{name}@example.com", password="password123")
|
||||
)
|
||||
|
||||
async def sweep(**kwargs):
|
||||
return await sweep_once(redis_client, presence, broadcaster, _Factory(), **kwargs)
|
||||
|
||||
return clock, presence, broadcaster, make_user, sweep
|
||||
|
||||
|
||||
async def test_sweep_announces_a_user_whose_connection_expired(world, db_session):
|
||||
clock, presence, broadcaster, make_user, sweep = world
|
||||
alice, bob = await make_user("alice"), await make_user("bob")
|
||||
room = await create_room(db_session, alice.id, RoomCreate(name=f"r-{uuid.uuid4().hex[:8]}"))
|
||||
dm = await find_or_create_dm(db_session, alice.id, bob.id)
|
||||
|
||||
await presence.connect(alice.id, "tab1")
|
||||
clock.advance(200)
|
||||
|
||||
assert await sweep(lock_seconds=None) == {alice.id}
|
||||
|
||||
room_ids = {rid for rid, p in broadcaster.room_events if p["type"] == "member_updated"}
|
||||
assert {room.id, dm.id} <= room_ids
|
||||
assert (bob.id, {"type": "dm_presence_update", "user_id": str(alice.id), "status": "offline"}) in (
|
||||
broadcaster.user_events
|
||||
)
|
||||
assert await presence.is_online(alice.id) is False
|
||||
|
||||
|
||||
async def test_sweep_announces_each_expiry_only_once(world, db_session):
|
||||
clock, presence, broadcaster, make_user, sweep = world
|
||||
alice = await make_user("alice")
|
||||
await create_room(db_session, alice.id, RoomCreate(name=f"r-{uuid.uuid4().hex[:8]}"))
|
||||
await presence.connect(alice.id, "tab1")
|
||||
clock.advance(200)
|
||||
|
||||
await sweep(lock_seconds=None)
|
||||
sent = len(broadcaster.room_events)
|
||||
assert sent > 0
|
||||
assert await sweep(lock_seconds=None) == set()
|
||||
assert len(broadcaster.room_events) == sent
|
||||
|
||||
|
||||
async def test_sweep_stays_quiet_while_the_user_has_another_live_connection(world, db_session):
|
||||
clock, presence, broadcaster, make_user, sweep = world
|
||||
alice = await make_user("alice")
|
||||
await create_room(db_session, alice.id, RoomCreate(name=f"r-{uuid.uuid4().hex[:8]}"))
|
||||
await presence.connect(alice.id, "old-tab")
|
||||
clock.advance(60)
|
||||
await presence.connect(alice.id, "new-tab")
|
||||
clock.advance(60)
|
||||
|
||||
assert await sweep(lock_seconds=None) == set()
|
||||
assert broadcaster.room_events == []
|
||||
assert await presence.is_online(alice.id) is True
|
||||
|
||||
|
||||
async def test_sweep_leaves_healthy_connections_alone(world, db_session):
|
||||
clock, presence, broadcaster, make_user, sweep = world
|
||||
alice = await make_user("alice")
|
||||
await presence.connect(alice.id, "tab1")
|
||||
clock.advance(30)
|
||||
assert await sweep(lock_seconds=None) == set()
|
||||
assert broadcaster.room_events == []
|
||||
|
||||
|
||||
async def test_only_one_sweeper_runs_per_interval(world, db_session, redis_client):
|
||||
clock, presence, broadcaster, make_user, sweep = world
|
||||
alice = await make_user("alice")
|
||||
await presence.connect(alice.id, "tab1")
|
||||
clock.advance(200)
|
||||
|
||||
# Another worker holds the lock: this one must not sweep (or broadcast).
|
||||
await redis_client.set(SWEEP_LOCK_KEY, "1", ex=14)
|
||||
assert await sweep(lock_seconds=14) == set()
|
||||
assert await redis_client.zcard(GlobalPresence.KEY) == 1
|
||||
|
||||
await redis_client.delete(SWEEP_LOCK_KEY)
|
||||
assert await sweep(lock_seconds=14) == {alice.id}
|
||||
# And the lock is now held, so an immediate second attempt does nothing.
|
||||
assert await sweep(lock_seconds=14) == set()
|
||||
|
||||
|
||||
async def test_one_failed_broadcast_does_not_stop_the_rest(world, db_session):
|
||||
clock, presence, broadcaster, make_user, sweep = world
|
||||
alice, bob = await make_user("alice"), await make_user("bob")
|
||||
room = await create_room(db_session, alice.id, RoomCreate(name=f"r-{uuid.uuid4().hex[:8]}"))
|
||||
other = await create_room(db_session, bob.id, RoomCreate(name=f"r-{uuid.uuid4().hex[:8]}"))
|
||||
await presence.connect(alice.id, "a")
|
||||
await presence.connect(bob.id, "b")
|
||||
clock.advance(200)
|
||||
|
||||
real_publish = broadcaster.publish
|
||||
failed = []
|
||||
|
||||
async def flaky_publish(room_id, payload):
|
||||
if room_id == room.id:
|
||||
failed.append(room_id)
|
||||
raise RuntimeError("redis hiccup")
|
||||
await real_publish(room_id, payload)
|
||||
|
||||
broadcaster.publish = flaky_publish
|
||||
|
||||
assert await sweep(lock_seconds=None) == {alice.id, bob.id}
|
||||
assert failed == [room.id]
|
||||
assert other.id in {rid for rid, _ in broadcaster.room_events}
|
||||
|
||||
|
||||
async def test_purge_removes_legacy_keys_but_keeps_v2(redis_client):
|
||||
legacy = ["presence:global", "presence:unfocused", f"presence:{uuid.uuid4()}"]
|
||||
current = f"presence:v2:test:{uuid.uuid4()}"
|
||||
for key in legacy:
|
||||
await redis_client.hset(key, "someone", 3)
|
||||
await redis_client.zadd(current, {"x:y": time.time() + 90})
|
||||
|
||||
try:
|
||||
assert await purge_legacy_presence(redis_client) >= len(legacy)
|
||||
for key in legacy:
|
||||
assert await redis_client.exists(key) == 0
|
||||
assert await redis_client.exists(current) == 1
|
||||
finally:
|
||||
await redis_client.delete(current, *legacy)
|
||||
Reference in new issue
Block a user