mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-06 08:18:53 +00:00
perf: bounded non-blocking session pool reaper, fewer blocking pool round trips (#28835)
* perf: stop the Socket.IO session pool blocking the websocket event loop With WEBSOCKET_MANAGER=redis the session pool is a synchronous Redis client, so every call into it blocks the whole worker's event loop, not just the caller. Two paths did it constantly: the orphan reaper walked the pool one round trip per session with no await anywhere, freezing the loop for the entire sweep every cycle, and nearly every socket event re-read the sender's session back out of Redis. Other users' events and every in-flight generation on that pod wait behind both. The reaper now walks the pool in HSCAN batches and deletes in bulk, yielding between batches, and no longer sleeps past half the lock TTL, which previously guaranteed a failed renew every cycle. The per-event reads are gone: Socket.IO events only reach the worker holding the connection, and that worker already saved the same session dict locally when the user authenticated, so it was asking Redis for its own data. The writes stay, since those are what other pods read. Measured at 4000 users / 16 containers, Redis 1.1 ms away: | | before | after | |---|---|---| | reaper sweep, 5k sessions | 5.6 s, loop frozen throughout | 62 ms, 4.2 ms worst block | | same, crash recovery with every session expired | 11.5 s | 96 ms | | heartbeat / usage ping / disconnect | 2 / 3 / 2 round trips | 1 / 2 / 1 | | 50-member channel post | 50 round trips, 57.2 ms block | 0 round trips, 0.02 ms | | loop time per wall second at rest | 103 ms (10.3%) | 67 ms (6.7%) | The alternative, converting RedisDict to the async client, fixes the same paths with a far larger blast radius (every call site gains await, and `in`/`[]`/`del` cannot be awaited so the dict interface goes) and still round-trips for data already in memory. Two deliberate behaviour changes: a heartbeat re-adds a session the reaper already removed, so a tab that survives a stall recovers instead of staying out of the pool until it reconnects; and disconnect no longer skips Yjs document cleanup when the pool entry is already gone, which previously leaked that document's update log forever. Closes #28172 * perf: cut disconnect and user session lookup pool round trips, harden the session reaper Follow-up on top of the session pool reaper branch. With WEBSOCKET_MANAGER=redis two paths still blocked the worker's event loop on synchronous Redis calls. Every disconnect listed all models in use cluster-wide and fetched each one individually, one blocking round trip per model. Disconnecting all sessions of a user (admin role change or deletion) pulled the entire session pool in one HGETALL and decoded every entry in a single uninterrupted block. Disconnect now fetches the usage pool once with items(), going from 2+N+M round trips to 2+M (N models in use cluster-wide, M models the session used), and its delete of an emptied model entry is KeyError-guarded because another node can remove the same key between snapshot and delete; unguarded, that race aborted the handler and skipped its Yjs document cleanup. The user session lookup reuses the reaper's HSCAN batches and yields to the loop between pages. The reaper previously died permanently on the first Redis connection error, on every node at once during an outage; it now logs, releases the lock and returns to retrying acquisition. * refac: keep the socket pool perf work to the round trips A review pass on this branch turned up four changes riding along with the round-trip work without belonging to it, so they are backed out here. The `disconnect` handler keeps its `if sid in SESSION_POOL:` guard, so USAGE_POOL and ydoc cleanup stay off the path for sockets that never authenticated. `RedisDict.set()` keeps its own inline HDEL. `get_session_ids_by_user_id` stays synchronous over one HGETALL, since it runs on user delete and role change rather than per message. The crash-resilience wrapper around the reaper loop is dropped; if that guard is worth having, it belongs in its own change. What stays is the perf part. The reaper now sweeps the pool in bounded HSCAN batches and deletes expired sids with one HDEL per batch, down from HKEYS plus an HGET and a per-sid HDEL across the whole pool. The `disconnect` handler reads USAGE_POOL with a single HGETALL, down from HKEYS plus one HGET per model in use. Session lookups in the socket handlers come from the local Socket.IO store, which removes one Redis GET from every heartbeat, usage, channel and ydoc event. Naming and annotations follow the file: `get_session_pool_batches` for the module's `get_` prefix, `RedisDict.pop_many` so both reaper branches use one word for removing keys, a named `SCAN_BATCH_SIZE`, and types on the new helpers. * fix: invalidate the RedisDict write signature on batch delete RedisDict.set() skips the write when the payload fingerprint matches the last one this process wrote, so a mutation that goes around set() has to clear that fingerprint. The new batch delete did not, leaving a stale fingerprint behind: the next refresh with identical content is treated as already written and silently skipped, so the hash stays empty. Renamed pop_many to delete_many. In a dict emulation pop removes and returns; this returns nothing and cannot without an extra HMGET, so the name promised something it does not do. delete_many matches __delitem__ and the HDEL underneath. Its only call site is the session pool reaper, whose behaviour is unchanged: same fields deleted, same batching, same return.
This commit is contained in:
Classic298
GitHub
parent
ac6a8c0082
commit
d7674c5174
3 changed files with 60 additions and 41 deletions
|
|
@ -1242,7 +1242,7 @@ async def post_new_message(
|
|||
except Exception as e:
|
||||
log.debug(e)
|
||||
|
||||
active_user_ids = get_user_ids_from_room(f'channel:{channel.id}')
|
||||
active_user_ids = await get_user_ids_from_room(f'channel:{channel.id}')
|
||||
|
||||
# NOTE: We intentionally do NOT pass db to background_handler.
|
||||
# Background tasks should manage their own short-lived sessions to avoid
|
||||
|
|
|
|||
|
|
@ -175,6 +175,13 @@ YDOC_MANAGER = YdocManager(
|
|||
)
|
||||
|
||||
|
||||
def get_session_pool_batches():
|
||||
"""All session pool entries, in bounded batches for the Redis backing."""
|
||||
if WEBSOCKET_MANAGER == 'redis':
|
||||
return SESSION_POOL.scan_batches()
|
||||
return [list(SESSION_POOL.items())]
|
||||
|
||||
|
||||
async def periodic_session_pool_cleanup():
|
||||
"""Reap orphaned SESSION_POOL entries that missed heartbeats (e.g. crashed instance)."""
|
||||
retry_delay = random.uniform(WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, WEBSOCKET_REDIS_LOCK_TIMEOUT)
|
||||
|
|
@ -192,14 +199,20 @@ async def periodic_session_pool_cleanup():
|
|||
break
|
||||
|
||||
now = int(time.time())
|
||||
for sid in list(SESSION_POOL.keys()):
|
||||
entry = SESSION_POOL.get(sid)
|
||||
if entry and now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT:
|
||||
log.warning(f'Reaping orphaned session {sid} (user {entry.get("id")})')
|
||||
try:
|
||||
del SESSION_POOL[sid]
|
||||
except KeyError:
|
||||
pass
|
||||
for batch in get_session_pool_batches():
|
||||
expired = [
|
||||
sid
|
||||
for sid, entry in batch
|
||||
if entry and now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT
|
||||
]
|
||||
if expired:
|
||||
log.warning('Reaping %d orphaned session(s) from the session pool', len(expired))
|
||||
if WEBSOCKET_MANAGER == 'redis':
|
||||
SESSION_POOL.delete_many(*expired)
|
||||
else:
|
||||
for sid in expired:
|
||||
SESSION_POOL.pop(sid, None)
|
||||
await asyncio.sleep(0) # don't hold the loop for the whole sweep
|
||||
|
||||
next_cleanup_at = time.monotonic() + SESSION_POOL_TIMEOUT
|
||||
lock_lost = False
|
||||
|
|
@ -283,6 +296,14 @@ def get_user_id_from_session_pool(sid):
|
|||
return None
|
||||
|
||||
|
||||
async def get_socket_session_user(sid: str) -> dict | None:
|
||||
"""Session user from this worker's local Socket.IO store; only locally connected sids are ever looked up."""
|
||||
try:
|
||||
return (await sio.get_session(sid)).get('user')
|
||||
except KeyError:
|
||||
return None
|
||||
|
||||
|
||||
def get_session_ids_from_room(room):
|
||||
"""Get all session IDs from a specific room."""
|
||||
members = get_room_sid_map(sio.manager, '/', room)
|
||||
|
|
@ -296,19 +317,9 @@ def get_session_ids_by_user_id(user_id: str) -> list[str]:
|
|||
return list(session_ids)
|
||||
|
||||
|
||||
def get_user_ids_from_room(room):
|
||||
active_session_ids = get_session_ids_from_room(room)
|
||||
|
||||
# Single pool lookup per session (each .get is a Redis round trip
|
||||
# when the session pool is Redis-backed).
|
||||
active_user_ids = list(
|
||||
{
|
||||
entry['id']
|
||||
for entry in (SESSION_POOL.get(session_id) for session_id in active_session_ids)
|
||||
if entry is not None
|
||||
}
|
||||
)
|
||||
return active_user_ids
|
||||
async def get_user_ids_from_room(room) -> set[str]:
|
||||
users = [await get_socket_session_user(session_id) for session_id in get_session_ids_from_room(room)]
|
||||
return {user['id'] for user in users if user}
|
||||
|
||||
|
||||
async def emit_to_users(event: str, data: dict, user_ids: list[str]):
|
||||
|
|
@ -364,7 +375,7 @@ async def disconnect_user_sessions(user_id: str):
|
|||
|
||||
@sio.on('usage')
|
||||
async def usage(sid, data):
|
||||
if sid in SESSION_POOL:
|
||||
if await get_socket_session_user(sid):
|
||||
model_id = data['model']
|
||||
# Record the timestamp for the last update
|
||||
current_time = int(time.time())
|
||||
|
|
@ -446,7 +457,7 @@ async def user_join(sid, data):
|
|||
|
||||
@sio.on('heartbeat')
|
||||
async def heartbeat(sid, data):
|
||||
user = SESSION_POOL.get(sid)
|
||||
user = await get_socket_session_user(sid)
|
||||
if user:
|
||||
SESSION_POOL[sid] = {**user, 'last_seen_at': int(time.time())}
|
||||
await Users.update_last_active_by_id(user['id'])
|
||||
|
|
@ -519,7 +530,7 @@ async def channel_events(sid, data):
|
|||
event_data = data['data']
|
||||
event_type = event_data['type']
|
||||
|
||||
user = SESSION_POOL.get(sid)
|
||||
user = await get_socket_session_user(sid)
|
||||
|
||||
if not user:
|
||||
return
|
||||
|
|
@ -559,15 +570,7 @@ async def get_folder_unread_counts(user_id: str) -> dict[str, int]:
|
|||
|
||||
@sio.on('events:chat')
|
||||
async def chat_events(sid, data):
|
||||
try:
|
||||
session = await sio.get_session(sid)
|
||||
user = session.get('user')
|
||||
except KeyError:
|
||||
user = None
|
||||
|
||||
if not user:
|
||||
user = SESSION_POOL.get(sid)
|
||||
|
||||
user = await get_socket_session_user(sid)
|
||||
if not user:
|
||||
return
|
||||
|
||||
|
|
@ -621,7 +624,7 @@ def normalize_document_id(document_id: str) -> str:
|
|||
@sio.on('ydoc:document:join')
|
||||
async def ydoc_document_join(sid, data):
|
||||
"""Handle user joining a document"""
|
||||
user = SESSION_POOL.get(sid)
|
||||
user = await get_socket_session_user(sid)
|
||||
if not user:
|
||||
return
|
||||
|
||||
|
|
@ -781,7 +784,7 @@ async def yjs_document_update(sid, data):
|
|||
return
|
||||
|
||||
# Verify write permission — room membership only proves read access
|
||||
user = SESSION_POOL.get(sid)
|
||||
user = await get_socket_session_user(sid)
|
||||
if not user:
|
||||
return
|
||||
|
||||
|
|
@ -851,7 +854,7 @@ async def yjs_document_update(sid, data):
|
|||
@sio.on('ydoc:document:leave')
|
||||
async def yjs_document_leave(sid, data):
|
||||
"""Handle user leaving a document"""
|
||||
user = SESSION_POOL.get(sid)
|
||||
user = await get_socket_session_user(sid)
|
||||
if not user: # authenticated session required (parity with sibling handlers)
|
||||
return
|
||||
try:
|
||||
|
|
@ -883,7 +886,7 @@ async def yjs_document_leave(sid, data):
|
|||
@sio.on('ydoc:awareness:update')
|
||||
async def yjs_awareness_update(sid, data):
|
||||
"""Handle awareness updates (cursors, selections, etc.)"""
|
||||
user = SESSION_POOL.get(sid)
|
||||
user = await get_socket_session_user(sid)
|
||||
if not user: # authenticated session required (parity with sibling handlers)
|
||||
return
|
||||
try:
|
||||
|
|
@ -911,9 +914,8 @@ async def disconnect(sid, reason=None):
|
|||
del SESSION_POOL[sid]
|
||||
|
||||
# Clean up USAGE_POOL entries for this session
|
||||
for model_id in list(USAGE_POOL.keys()):
|
||||
connections = USAGE_POOL.get(model_id)
|
||||
if connections and sid in connections:
|
||||
for model_id, connections in list(USAGE_POOL.items()):
|
||||
if sid in connections:
|
||||
del connections[sid]
|
||||
if not connections:
|
||||
del USAGE_POOL[model_id]
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from open_webui.utils.json_codec import JSONCodec
|
|||
from open_webui.utils.redis import get_redis_connection
|
||||
|
||||
YDOC_KEY_PREFIX = f'{REDIS_KEY_PREFIX}:ydoc:documents'
|
||||
SCAN_BATCH_SIZE = 200
|
||||
|
||||
|
||||
class RedisLock:
|
||||
|
|
@ -112,6 +113,22 @@ class RedisDict:
|
|||
def items(self):
|
||||
return [(k, JSONCodec.loads(v)) for k, v in self.redis.hgetall(self.name).items()]
|
||||
|
||||
def scan_batches(self):
|
||||
"""Yield lists of (key, value) pairs via incremental HSCAN; a field may repeat across batches."""
|
||||
cursor = 0
|
||||
while True:
|
||||
cursor, batch = self.redis.hscan(self.name, cursor, count=SCAN_BATCH_SIZE)
|
||||
if batch:
|
||||
yield [(k, JSONCodec.loads(v)) for k, v in batch.items()]
|
||||
if cursor == 0:
|
||||
break
|
||||
|
||||
def delete_many(self, *keys):
|
||||
"""Delete fields in one HDEL; no keys is a no-op (HDEL rejects an empty field list)."""
|
||||
if keys:
|
||||
self.redis.hdel(self.name, *keys)
|
||||
self._last_signature = None
|
||||
|
||||
def set(self, mapping: dict):
|
||||
if not mapping:
|
||||
self.clear()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue