mirror of
https://github.com/fastapi-practices/fastapi-best-architecture.git
synced 2026-10-05 01:20:10 +00:00
Wu Clan
110a9c2e0c
* Optimize Redis connection pool and batch operations * Refactor token session index and online status
119 lines
4.2 KiB
Python
119 lines
4.2 KiB
Python
import urllib.parse
|
|
import uuid
|
|
|
|
from typing import Any
|
|
|
|
import socketio
|
|
|
|
from starlette_context import request_cycle_context
|
|
|
|
from backend.common.log import log
|
|
from backend.common.security.jwt import jwt_authentication, jwt_decode
|
|
from backend.core.conf import settings
|
|
from backend.database.redis import redis_client
|
|
from backend.utils.timezone import timezone
|
|
|
|
# 创建 Socket.IO 服务器实例
|
|
sio = socketio.AsyncServer(
|
|
client_manager=socketio.AsyncRedisManager(
|
|
f'redis://:{urllib.parse.quote(settings.REDIS_PASSWORD)}@{settings.REDIS_HOST}:{settings.REDIS_PORT}/{settings.REDIS_DATABASE}',
|
|
redis_options={
|
|
'socket_timeout': None,
|
|
'socket_connect_timeout': settings.REDIS_TIMEOUT,
|
|
},
|
|
),
|
|
async_mode='asgi',
|
|
cors_allowed_origins=settings.CORS_ALLOWED_ORIGINS,
|
|
cors_credentials=True,
|
|
namespaces=['/', '/ws'],
|
|
)
|
|
|
|
|
|
@sio.event(namespace='*')
|
|
async def connect(namespace: str, sid: str, _environ: dict[str, Any], auth: dict[str, Any] | None) -> bool:
|
|
"""
|
|
Socket 连接事件
|
|
|
|
:param namespace: 命名空间
|
|
:param sid: 连接 ID
|
|
:param _environ: 连接环境
|
|
:param auth: 授权信息
|
|
:return:
|
|
"""
|
|
if namespace not in sio.namespaces:
|
|
return False
|
|
if not isinstance(auth, dict):
|
|
log.error('WebSocket 连接失败：无授权')
|
|
return False
|
|
session_uuid = auth.get('session_uuid')
|
|
token = auth.get('token')
|
|
if not isinstance(token, str) or not token or not isinstance(session_uuid, str) or not session_uuid:
|
|
log.error('WebSocket 连接失败：授权失败，请检查')
|
|
return False
|
|
|
|
# 免授权直连
|
|
if token == settings.WS_NO_AUTH_MARKER:
|
|
if settings.ENVIRONMENT == 'prod':
|
|
log.error('WebSocket 连接失败：生产环境禁止免授权直连')
|
|
return False
|
|
expire = settings.TOKEN_EXPIRE_SECONDS
|
|
else:
|
|
try:
|
|
with request_cycle_context({settings.TRACE_ID_REQUEST_HEADER_KEY: uuid.uuid4().hex}):
|
|
await jwt_authentication(token)
|
|
token_payload = jwt_decode(token)
|
|
except Exception as e:
|
|
log.info(f'WebSocket 连接失败：{e!s}')
|
|
return False
|
|
session_uuid = token_payload.session_uuid
|
|
expire = int((token_payload.expire_time - timezone.now()).total_seconds())
|
|
if expire <= 0:
|
|
log.info('WebSocket 连接失败：Token 已过期')
|
|
return False
|
|
|
|
await sio.save_session(sid, {'session_uuid': session_uuid}, namespace=namespace)
|
|
sid_key = f'{settings.TOKEN_ONLINE_REDIS_PREFIX}:sid:{sid}'
|
|
session_key = f'{settings.TOKEN_ONLINE_REDIS_PREFIX}:session:{session_uuid}'
|
|
await redis_client.set(sid_key, session_uuid, ex=expire)
|
|
await redis_client.sadd(session_key, sid)
|
|
session_ttl = await redis_client.ttl(session_key)
|
|
# 新集合尚无 TTL 时按本次连接设置；多连接取最长剩余寿命
|
|
new_ttl = expire if session_ttl < 0 else max(session_ttl, expire)
|
|
await redis_client.expire(session_key, new_ttl)
|
|
return True
|
|
|
|
|
|
@sio.event(namespace='*')
|
|
async def disconnect(namespace: str, sid: str, _reason: str | None = None) -> None:
|
|
"""
|
|
Socket 断开连接事件
|
|
|
|
:param namespace: 命名空间
|
|
:param sid: 连接 ID
|
|
:param _reason: 断开原因
|
|
:return:
|
|
"""
|
|
sid_key = f'{settings.TOKEN_ONLINE_REDIS_PREFIX}:sid:{sid}'
|
|
session_uuid = await redis_client.get(sid_key)
|
|
if not session_uuid:
|
|
try:
|
|
session_data = await sio.get_session(sid, namespace=namespace)
|
|
except KeyError:
|
|
return
|
|
session_uuid = session_data.get('session_uuid')
|
|
if not session_uuid:
|
|
return
|
|
|
|
session_key = f'{settings.TOKEN_ONLINE_REDIS_PREFIX}:session:{session_uuid}'
|
|
await redis_client.delete(sid_key)
|
|
await redis_client.srem(session_key, sid)
|
|
remaining = list(await redis_client.smembers(session_key))
|
|
if not remaining:
|
|
return
|
|
mappings = await redis_client.mget_batched([
|
|
f'{settings.TOKEN_ONLINE_REDIS_PREFIX}:sid:{other_sid}' for other_sid in remaining
|
|
])
|
|
stale_sids = [other_sid for other_sid, mapping in zip(remaining, mappings, strict=True) if mapping != session_uuid]
|
|
if stale_sids:
|
|
await redis_client.srem(session_key, *stale_sids)
|