Implement real-time WebSocket gateway for presence, DM, and world chat (US-4, US-2)
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""WebSocket real-time gateway package."""
|
||||
@@ -0,0 +1,538 @@
|
||||
"""Real-time WebSocket gateway.
|
||||
|
||||
Multiplexes presence, DM, and world-chat delivery over a single WebSocket
|
||||
endpoint (docs/planning/social.md US-2, US-4).
|
||||
|
||||
Key features:
|
||||
- Authentication via JWT token (query param or initial auth frame).
|
||||
- Scoped subscriptions to DM conversations and world channels.
|
||||
- Server-derived presence heartbeat (Redis TTL refreshed periodically).
|
||||
- Broadcast presence changes (online/offline) to mutual friends.
|
||||
- Fan-out via Redis pub/sub.
|
||||
- Thread-safe & event-loop-safe dispatch for multi-instance and test environments.
|
||||
- Graceful reconnect without duplicate delivery or orphaned state.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status
|
||||
from jose import JWTError, jwt
|
||||
from sqlalchemy import or_, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
from app.config import settings
|
||||
from app.database import async_session_factory, get_db
|
||||
from app.models.chat import ConversationParticipant
|
||||
from app.models.social import Friendship, FriendshipStatus
|
||||
from app.models.user import User
|
||||
from app.models.world import World, WorldVisibility
|
||||
from app.services.presence import presence_store
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_db_session():
|
||||
"""Helper to acquire an async DB session respecting test overrides if present."""
|
||||
from app.main import app
|
||||
|
||||
override = app.dependency_overrides.get(get_db)
|
||||
if override:
|
||||
res = override()
|
||||
if hasattr(res, "__anext__"):
|
||||
session = await res.__anext__()
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
try:
|
||||
await res.__anext__()
|
||||
except (StopAsyncIteration, Exception):
|
||||
pass
|
||||
return
|
||||
elif hasattr(res, "__next__"):
|
||||
session = next(res)
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
try:
|
||||
next(res)
|
||||
except (StopIteration, Exception):
|
||||
pass
|
||||
return
|
||||
|
||||
async with async_session_factory() as session:
|
||||
yield session
|
||||
|
||||
|
||||
async def authenticate_token(token: str) -> User | None:
|
||||
"""Validate JWT token and return authenticated User, or None."""
|
||||
try:
|
||||
payload = jwt.decode(token, settings.public_pem, algorithms=[settings.ALGORITHM])
|
||||
user_id_raw = payload.get("sub")
|
||||
if user_id_raw is None:
|
||||
return None
|
||||
user_id = int(user_id_raw)
|
||||
except (JWTError, ValueError, TypeError):
|
||||
return None
|
||||
|
||||
async with get_db_session() as db:
|
||||
res = await db.execute(select(User).where(User.id == user_id))
|
||||
user = res.scalar_one_or_none()
|
||||
if user and user.is_active:
|
||||
return user
|
||||
return None
|
||||
|
||||
|
||||
async def get_friend_ids(user_id: int) -> set[int]:
|
||||
"""Return set of accepted friend user IDs for user_id."""
|
||||
try:
|
||||
async with get_db_session() as db:
|
||||
res = await db.execute(
|
||||
select(Friendship).where(
|
||||
or_(Friendship.requester_id == user_id, Friendship.addressee_id == user_id),
|
||||
Friendship.status == FriendshipStatus.ACCEPTED,
|
||||
)
|
||||
)
|
||||
friendships = res.scalars().all()
|
||||
return {
|
||||
f.addressee_id if f.requester_id == user_id else f.requester_id
|
||||
for f in friendships
|
||||
}
|
||||
except Exception as exc:
|
||||
logger.error("Failed to get friend IDs for user %s: %s", user_id, exc, exc_info=True)
|
||||
return set()
|
||||
|
||||
|
||||
async def check_world_access(world_id: int, user_id: int) -> bool:
|
||||
"""Verify if user has permission to join/view a world."""
|
||||
async with get_db_session() as db:
|
||||
res = await db.execute(select(World).where(World.id == world_id))
|
||||
world = res.scalar_one_or_none()
|
||||
if not world:
|
||||
return False
|
||||
if world.visibility == WorldVisibility.PUBLIC:
|
||||
return True
|
||||
if world.owner_id == user_id:
|
||||
return True
|
||||
if world.visibility == WorldVisibility.FRIENDS_ONLY:
|
||||
res_fr = await db.execute(
|
||||
select(Friendship.id).where(
|
||||
or_(
|
||||
and_(Friendship.requester_id == world.owner_id, Friendship.addressee_id == user_id),
|
||||
and_(Friendship.requester_id == user_id, Friendship.addressee_id == world.owner_id),
|
||||
),
|
||||
Friendship.status == FriendshipStatus.ACCEPTED,
|
||||
).limit(1)
|
||||
)
|
||||
return res_fr.scalar_one_or_none() is not None
|
||||
return False
|
||||
|
||||
|
||||
async def check_conversation_access(conversation_id: int, user_id: int) -> bool:
|
||||
"""Verify if user is a participant in the given conversation."""
|
||||
async with get_db_session() as db:
|
||||
res = await db.execute(
|
||||
select(ConversationParticipant.id).where(
|
||||
ConversationParticipant.conversation_id == conversation_id,
|
||||
ConversationParticipant.user_id == user_id,
|
||||
).limit(1)
|
||||
)
|
||||
return res.scalar_one_or_none() is not None
|
||||
|
||||
|
||||
async def _async_update_last_seen(user_id: int, now: datetime):
|
||||
try:
|
||||
async with async_session_factory() as session:
|
||||
res = await session.execute(select(User).where(User.id == user_id))
|
||||
db_user = res.scalar_one_or_none()
|
||||
if db_user:
|
||||
db_user.last_seen_at = now
|
||||
await session.commit()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class ConnectionManager:
|
||||
"""Manages active WebSocket connections, subscriptions, presence, and fan-out."""
|
||||
|
||||
def __init__(self):
|
||||
# ws -> User
|
||||
self._connections: dict[WebSocket, User] = {}
|
||||
# ws -> event loop id
|
||||
self._ws_loops: dict[WebSocket, int] = {}
|
||||
# user_id -> set of active WebSockets
|
||||
self._user_connections: dict[int, set[WebSocket]] = {}
|
||||
# user_id -> cached friend user IDs
|
||||
self._user_friend_ids: dict[int, set[int]] = {}
|
||||
# world_id -> set of WebSockets
|
||||
self._world_subs: dict[int, set[WebSocket]] = {}
|
||||
# conversation_id -> set of WebSockets
|
||||
self._dm_subs: dict[int, set[WebSocket]] = {}
|
||||
# ws -> asyncio heartbeat task
|
||||
self._heartbeat_tasks: dict[WebSocket, asyncio.Task] = {}
|
||||
# Per-event-loop Redis pub/sub listener tasks and events
|
||||
self._redis_tasks: dict[int, asyncio.Task] = {}
|
||||
self._subscribed_events: dict[int, asyncio.Event] = {}
|
||||
|
||||
async def ensure_redis_listener(self):
|
||||
"""Ensure background Redis pub/sub listener is connected and subscribed for current loop."""
|
||||
loop = asyncio.get_running_loop()
|
||||
loop_key = id(loop)
|
||||
task = self._redis_tasks.get(loop_key)
|
||||
if task is None or task.done():
|
||||
event = asyncio.Event()
|
||||
self._subscribed_events[loop_key] = event
|
||||
task = asyncio.create_task(self._redis_listener(loop_key, event))
|
||||
self._redis_tasks[loop_key] = task
|
||||
else:
|
||||
event = self._subscribed_events.get(loop_key)
|
||||
|
||||
if event:
|
||||
try:
|
||||
await asyncio.wait_for(event.wait(), timeout=2.0)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
async def connect(self, websocket: WebSocket, user: User):
|
||||
"""Register an authenticated connection."""
|
||||
await self.ensure_redis_listener()
|
||||
loop_key = id(asyncio.get_running_loop())
|
||||
|
||||
self._connections[websocket] = user
|
||||
self._ws_loops[websocket] = loop_key
|
||||
self._user_connections.setdefault(user.id, set()).add(websocket)
|
||||
|
||||
# Cache friend IDs for instant presence dispatch
|
||||
friend_ids = await get_friend_ids(user.id)
|
||||
self._user_friend_ids[user.id] = friend_ids
|
||||
|
||||
# Mark user online in Redis and DB
|
||||
await presence_store.heartbeat(user.id)
|
||||
now = datetime.now(timezone.utc)
|
||||
asyncio.create_task(_async_update_last_seen(user.id, now))
|
||||
|
||||
# Start periodic heartbeat task (every 20 seconds)
|
||||
self._heartbeat_tasks[websocket] = asyncio.create_task(
|
||||
self._run_periodic_heartbeat(user.id)
|
||||
)
|
||||
|
||||
# Broadcast online presence to friends via Redis
|
||||
await self.broadcast_presence(user.id, user.username, is_online=True, friend_ids=friend_ids)
|
||||
|
||||
async def disconnect(self, websocket: WebSocket):
|
||||
"""Unregister a connection and clean up subscriptions and presence."""
|
||||
user = self._connections.pop(websocket, None)
|
||||
self._ws_loops.pop(websocket, None)
|
||||
task = self._heartbeat_tasks.pop(websocket, None)
|
||||
if task:
|
||||
task.cancel()
|
||||
|
||||
# Clean subscriptions
|
||||
for world_id in list(self._world_subs.keys()):
|
||||
self._world_subs[world_id].discard(websocket)
|
||||
if not self._world_subs[world_id]:
|
||||
self._world_subs.pop(world_id, None)
|
||||
|
||||
for conv_id in list(self._dm_subs.keys()):
|
||||
self._dm_subs[conv_id].discard(websocket)
|
||||
if not self._dm_subs[conv_id]:
|
||||
self._dm_subs.pop(conv_id, None)
|
||||
|
||||
if user is None:
|
||||
return
|
||||
|
||||
user_id = user.id
|
||||
username = user.username
|
||||
user_set = self._user_connections.get(user_id)
|
||||
if user_set:
|
||||
user_set.discard(websocket)
|
||||
if not user_set:
|
||||
self._user_connections.pop(user_id, None)
|
||||
|
||||
# If user has no more open connections, mark offline immediately
|
||||
if user_id not in self._user_connections:
|
||||
await presence_store.clear_heartbeat(user_id)
|
||||
now = datetime.now(timezone.utc)
|
||||
asyncio.create_task(_async_update_last_seen(user_id, now))
|
||||
friend_ids = self._user_friend_ids.pop(user_id, set())
|
||||
await self.broadcast_presence(
|
||||
user_id, username, is_online=False, last_seen_at=now, friend_ids=friend_ids
|
||||
)
|
||||
|
||||
async def _run_periodic_heartbeat(self, user_id: int):
|
||||
"""Keep Redis TTL alive while connection is open."""
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(20)
|
||||
await presence_store.heartbeat(user_id)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def subscribe_world(self, websocket: WebSocket, world_id: int) -> bool:
|
||||
user = self._connections.get(websocket)
|
||||
if not user:
|
||||
return False
|
||||
allowed = await check_world_access(world_id, user.id)
|
||||
if not allowed:
|
||||
return False
|
||||
self._world_subs.setdefault(world_id, set()).add(websocket)
|
||||
return True
|
||||
|
||||
def unsubscribe_world(self, websocket: WebSocket, world_id: int):
|
||||
if world_id in self._world_subs:
|
||||
self._world_subs[world_id].discard(websocket)
|
||||
|
||||
async def subscribe_dm(self, websocket: WebSocket, conversation_id: int) -> bool:
|
||||
user = self._connections.get(websocket)
|
||||
if not user:
|
||||
return False
|
||||
allowed = await check_conversation_access(conversation_id, user.id)
|
||||
if not allowed:
|
||||
return False
|
||||
self._dm_subs.setdefault(conversation_id, set()).add(websocket)
|
||||
return True
|
||||
|
||||
def unsubscribe_dm(self, websocket: WebSocket, conversation_id: int):
|
||||
if conversation_id in self._dm_subs:
|
||||
self._dm_subs[conversation_id].discard(websocket)
|
||||
|
||||
async def send_to_ws(self, websocket: WebSocket, message: dict):
|
||||
"""Deliver frame to websocket if it belongs to current event loop."""
|
||||
cur_loop = id(asyncio.get_running_loop())
|
||||
if self._ws_loops.get(websocket) != cur_loop:
|
||||
return
|
||||
try:
|
||||
await websocket.send_text(json.dumps(message, default=str))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def broadcast_world_chat(
|
||||
self, world_id: int, message_payload: dict, room_id: int | None = None
|
||||
):
|
||||
"""Send world chat message to all connected clients on this loop subscribed to world_id."""
|
||||
targets = list(self._world_subs.get(world_id, set()))
|
||||
if not targets:
|
||||
return
|
||||
frame = {
|
||||
"type": "world_chat",
|
||||
"world_id": world_id,
|
||||
"room_id": room_id,
|
||||
"message": message_payload,
|
||||
}
|
||||
for ws in targets:
|
||||
await self.send_to_ws(ws, frame)
|
||||
|
||||
async def broadcast_dm_chat(self, conversation_id: int, message_payload: dict):
|
||||
"""Send DM message to all connected clients on this loop subscribed to conversation_id."""
|
||||
targets = list(self._dm_subs.get(conversation_id, set()))
|
||||
frame = {
|
||||
"type": "dm_chat",
|
||||
"conversation_id": conversation_id,
|
||||
"message": message_payload,
|
||||
}
|
||||
for ws in targets:
|
||||
await self.send_to_ws(ws, frame)
|
||||
|
||||
async def broadcast_presence(
|
||||
self,
|
||||
user_id: int,
|
||||
username: str,
|
||||
is_online: bool,
|
||||
last_seen_at: datetime | None = None,
|
||||
friend_ids: set[int] | None = None,
|
||||
):
|
||||
"""Broadcast presence update to mutual friends via Redis."""
|
||||
if friend_ids is None:
|
||||
friend_ids = await get_friend_ids(user_id)
|
||||
if not friend_ids:
|
||||
return
|
||||
|
||||
frame = {
|
||||
"type": "presence",
|
||||
"user_id": user_id,
|
||||
"username": username,
|
||||
"is_online": is_online,
|
||||
"last_seen_at": last_seen_at.isoformat() if last_seen_at else None,
|
||||
"friend_ids": list(friend_ids),
|
||||
}
|
||||
|
||||
# Redis broadcast for all loops/instances
|
||||
try:
|
||||
client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True)
|
||||
async with client:
|
||||
await client.publish("presence:events", json.dumps(frame, default=str))
|
||||
except Exception as exc:
|
||||
logger.error("Failed to publish presence to Redis: %s", exc, exc_info=True)
|
||||
|
||||
async def _redis_listener(self, loop_key: int, event: asyncio.Event):
|
||||
"""Listen to Redis pub/sub channels for cross-instance message dispatch."""
|
||||
while True:
|
||||
try:
|
||||
client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True)
|
||||
pubsub = client.pubsub()
|
||||
await pubsub.psubscribe("world:*:chat", "dm:*:chat", "presence:events")
|
||||
event.set()
|
||||
logger.info("Gateway Redis pub/sub listener active for loop %s", loop_key)
|
||||
async for message in pubsub.listen():
|
||||
if message["type"] != "pmessage":
|
||||
continue
|
||||
channel = message["channel"]
|
||||
raw_data = message["data"]
|
||||
try:
|
||||
data = json.loads(raw_data)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if channel.startswith("world:") and channel.endswith(":chat"):
|
||||
parts = channel.split(":")
|
||||
world_id = int(parts[1])
|
||||
await self.broadcast_world_chat(world_id, data, room_id=data.get("room_id"))
|
||||
elif channel.startswith("dm:") and channel.endswith(":chat"):
|
||||
parts = channel.split(":")
|
||||
conv_id = int(parts[1])
|
||||
await self.broadcast_dm_chat(conv_id, data)
|
||||
elif channel == "presence:events":
|
||||
logger.info("Redis presence event received: %s. Active users: %s", data, list(self._user_connections.keys()))
|
||||
for friend_id in data.get("friend_ids", []):
|
||||
targets = list(self._user_connections.get(friend_id, set()))
|
||||
logger.info("Friend %s has %s targets", friend_id, len(targets))
|
||||
for ws in targets:
|
||||
await self.send_to_ws(ws, data)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as exc:
|
||||
logger.debug("Redis listener encountered error, retrying in 2s: %s", exc)
|
||||
event.clear()
|
||||
await asyncio.sleep(2)
|
||||
|
||||
|
||||
connection_manager = ConnectionManager()
|
||||
|
||||
|
||||
# ──────────────────────────────── WebSocket Route ────────────────────────────
|
||||
|
||||
|
||||
@router.websocket("/ws")
|
||||
async def websocket_gateway(websocket: WebSocket, token: str | None = None):
|
||||
"""Main WebSocket gateway endpoint for real-time presence and chat."""
|
||||
user: User | None = None
|
||||
# 1. Check query parameter token
|
||||
if token:
|
||||
user = await authenticate_token(token)
|
||||
if user is None:
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
||||
return
|
||||
|
||||
await websocket.accept()
|
||||
|
||||
# 2. If not authenticated via query param, wait for first auth frame
|
||||
if user is None:
|
||||
try:
|
||||
raw = await asyncio.wait_for(websocket.receive_text(), timeout=3.0)
|
||||
data = json.loads(raw)
|
||||
if data.get("type") == "auth" and data.get("token"):
|
||||
user = await authenticate_token(data["token"])
|
||||
except (asyncio.TimeoutError, json.JSONDecodeError, WebSocketDisconnect):
|
||||
pass
|
||||
|
||||
if user is None:
|
||||
try:
|
||||
await websocket.send_text(json.dumps({"type": "error", "message": "Authentication failed"}))
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
# Mark user online in Redis immediately upon successful auth
|
||||
await presence_store.heartbeat(user.id)
|
||||
|
||||
# Send auth_ok handshake immediately upon successful authentication
|
||||
await websocket.send_text(
|
||||
json.dumps({
|
||||
"type": "auth_ok",
|
||||
"user_id": user.id,
|
||||
"username": user.username,
|
||||
})
|
||||
)
|
||||
|
||||
# Register with connection manager for presence and real-time event delivery
|
||||
await connection_manager.connect(websocket, user)
|
||||
|
||||
try:
|
||||
while True:
|
||||
text = await websocket.receive_text()
|
||||
try:
|
||||
msg = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
await websocket.send_text(json.dumps({"type": "error", "message": "Invalid JSON"}))
|
||||
continue
|
||||
|
||||
msg_type = msg.get("type")
|
||||
if msg_type == "ping":
|
||||
await websocket.send_text(json.dumps({"type": "pong"}))
|
||||
elif msg_type == "heartbeat":
|
||||
await presence_store.heartbeat(user.id)
|
||||
await websocket.send_text(json.dumps({"type": "heartbeat_ack"}))
|
||||
elif msg_type in ("subscribe_world", "subscribe") and (
|
||||
msg.get("world_id") or msg.get("channel") == "world"
|
||||
):
|
||||
world_id = msg.get("world_id") or msg.get("id")
|
||||
if isinstance(world_id, int):
|
||||
ok = await connection_manager.subscribe_world(websocket, world_id)
|
||||
if ok:
|
||||
await websocket.send_text(
|
||||
json.dumps({"type": "subscribed", "channel": "world", "id": world_id})
|
||||
)
|
||||
else:
|
||||
await websocket.send_text(
|
||||
json.dumps({"type": "error", "message": f"Cannot subscribe to world {world_id}"})
|
||||
)
|
||||
elif msg_type in ("unsubscribe_world", "unsubscribe") and (
|
||||
msg.get("world_id") or msg.get("channel") == "world"
|
||||
):
|
||||
world_id = msg.get("world_id") or msg.get("id")
|
||||
if isinstance(world_id, int):
|
||||
connection_manager.unsubscribe_world(websocket, world_id)
|
||||
await websocket.send_text(
|
||||
json.dumps({"type": "unsubscribed", "channel": "world", "id": world_id})
|
||||
)
|
||||
elif msg_type in ("subscribe_dm", "subscribe") and (
|
||||
msg.get("conversation_id") or msg.get("channel") == "dm"
|
||||
):
|
||||
conv_id = msg.get("conversation_id") or msg.get("id")
|
||||
if isinstance(conv_id, int):
|
||||
ok = await connection_manager.subscribe_dm(websocket, conv_id)
|
||||
if ok:
|
||||
await websocket.send_text(
|
||||
json.dumps({"type": "subscribed", "channel": "dm", "id": conv_id})
|
||||
)
|
||||
else:
|
||||
await websocket.send_text(
|
||||
json.dumps({"type": "error", "message": f"Cannot subscribe to conversation {conv_id}"})
|
||||
)
|
||||
elif msg_type in ("unsubscribe_dm", "unsubscribe") and (
|
||||
msg.get("conversation_id") or msg.get("channel") == "dm"
|
||||
):
|
||||
conv_id = msg.get("conversation_id") or msg.get("id")
|
||||
if isinstance(conv_id, int):
|
||||
connection_manager.unsubscribe_dm(websocket, conv_id)
|
||||
await websocket.send_text(
|
||||
json.dumps({"type": "unsubscribed", "channel": "dm", "id": conv_id})
|
||||
)
|
||||
else:
|
||||
await websocket.send_text(json.dumps({"type": "error", "message": f"Unknown action: {msg_type}"}))
|
||||
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
finally:
|
||||
await connection_manager.disconnect(websocket)
|
||||
Reference in New Issue
Block a user