From cce8250d52a8cffa7cec890a10009738ba197ef0 Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:35:36 +0200 Subject: [PATCH 1/2] 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 --- backend/open_webui/routers/channels.py | 2 +- backend/open_webui/socket/main.py | 105 ++++++++++++------------- backend/open_webui/socket/utils.py | 18 ++++- 3 files changed, 69 insertions(+), 56 deletions(-) diff --git a/backend/open_webui/routers/channels.py b/backend/open_webui/routers/channels.py index 652cbdf07f..5251d472e8 100644 --- a/backend/open_webui/routers/channels.py +++ b/backend/open_webui/routers/channels.py @@ -1224,7 +1224,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 diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 7cb7904fed..6896c664ec 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -176,6 +176,7 @@ YDOC_MANAGER = YdocManager( 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) + is_redis = WEBSOCKET_MANAGER == 'redis' while True: if not session_aquire_func(): log.debug('Session cleanup lock held by another node. Retrying.') @@ -189,15 +190,24 @@ 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 - await asyncio.sleep(SESSION_POOL_TIMEOUT) + batches = SESSION_POOL.scan_batches() if is_redis else [list(SESSION_POOL.items())] + for batch in batches: + expired = { + sid: entry.get('id') + for sid, entry in batch + if now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT + } + if expired: + log.warning('Reaping %d orphaned session(s) (sid: user): %s', len(expired), expired) + if is_redis: + SESSION_POOL.discard(*expired) + else: + for sid in expired: + SESSION_POOL.pop(sid, None) + await asyncio.sleep(0) # don't hold the loop for the whole sweep + + # Never sleep past half the lock TTL; renewals only happen between sweeps. + await asyncio.sleep(min(SESSION_POOL_TIMEOUT, WEBSOCKET_REDIS_LOCK_TIMEOUT / 2)) finally: session_release_func() @@ -265,6 +275,14 @@ def get_user_id_from_session_pool(sid): return None +async def get_socket_session_user(sid): + """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) @@ -278,19 +296,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): + 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]): @@ -346,7 +354,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()) @@ -428,7 +436,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']) @@ -501,7 +509,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 @@ -541,15 +549,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 @@ -603,7 +603,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 @@ -763,7 +763,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 @@ -828,7 +828,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: @@ -860,7 +860,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: @@ -884,23 +884,22 @@ async def yjs_awareness_update(sid, data): @sio.event async def disconnect(sid, reason=None): - if sid in SESSION_POOL: + try: 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: - del connections[sid] - if not connections: - del USAGE_POOL[model_id] - else: - USAGE_POOL[model_id] = connections - - await YDOC_MANAGER.remove_user_from_all_documents(sid) - else: + except KeyError: pass - # print(f"Unknown session ID {sid} disconnected") + + # 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: + del connections[sid] + if not connections: + del USAGE_POOL[model_id] + else: + USAGE_POOL[model_id] = connections + + await YDOC_MANAGER.remove_user_from_all_documents(sid) async def _make_channel_emitter(request_info): diff --git a/backend/open_webui/socket/utils.py b/backend/open_webui/socket/utils.py index 00f8424aae..c8c1f2868f 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -104,6 +104,21 @@ 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=200) + if batch: + yield [(k, JSONCodec.loads(v)) for k, v in batch.items()] + if cursor == 0: + break + + def discard(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) + def set(self, mapping: dict): if not mapping: self.redis.delete(self.name) @@ -137,8 +152,7 @@ class RedisDict: # We never DELETE the whole hash — this eliminates the race window # where concurrent readers would see an empty models dict. self.redis.hset(self.name, mapping=serialized) - if keys_to_remove: - self.redis.hdel(self.name, *keys_to_remove) + self.discard(*keys_to_remove) self._last_signature = signature From 75fe65c2f85825535419a74e6d7fba17e252f83a Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:36:45 +0200 Subject: [PATCH 2/2] 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. --- backend/open_webui/socket/main.py | 103 +++++++++++++++++------------- 1 file changed, 58 insertions(+), 45 deletions(-) diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index d09e32d3dc..e2efb783dd 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -173,56 +173,65 @@ YDOC_MANAGER = YdocManager( ) +def 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) - is_redis = WEBSOCKET_MANAGER == 'redis' renew_interval = max(WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, 0.5) while True: - if not session_aquire_func(): - log.debug('Session cleanup lock held by another node. Retrying.') - await asyncio.sleep(retry_delay) - continue - try: - while True: - if not session_renew_func(): - log.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.') - break + if not session_aquire_func(): + log.debug('Session cleanup lock held by another node. Retrying.') + await asyncio.sleep(retry_delay) + continue - now = int(time.time()) - batches = SESSION_POOL.scan_batches() if is_redis else [list(SESSION_POOL.items())] - for batch in batches: - expired = { - sid: entry.get('id') - for sid, entry in batch - if now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT - } - if expired: - log.warning('Reaping %d orphaned session(s) (sid: user): %s', len(expired), expired) - if is_redis: - SESSION_POOL.discard(*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 + try: while True: - sleep_for = min(renew_interval, next_cleanup_at - time.monotonic()) - if sleep_for <= 0: - break - await asyncio.sleep(sleep_for) if not session_renew_func(): log.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.') - lock_lost = True break - if lock_lost: - break - finally: - session_release_func() + now = int(time.time()) + for batch in session_pool_batches(): + expired = { + sid: entry.get('id') + for sid, entry in batch + if now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT + } + if expired: + log.warning('Reaping %d orphaned session(s) (sid: user): %s', len(expired), expired) + if WEBSOCKET_MANAGER == 'redis': + SESSION_POOL.discard(*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 + while True: + sleep_for = min(renew_interval, next_cleanup_at - time.monotonic()) + if sleep_for <= 0: + break + await asyncio.sleep(sleep_for) + if not session_renew_func(): + log.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.') + lock_lost = True + break + + if lock_lost: + break + finally: + session_release_func() + except Exception: + log.exception('Session pool cleanup failed. Retrying.') + await asyncio.sleep(retry_delay) async def periodic_usage_pool_cleanup(): @@ -302,10 +311,12 @@ def get_session_ids_from_room(room): return list(members) if members else [] -def get_session_ids_by_user_id(user_id: str) -> list[str]: +async def get_session_ids_by_user_id(user_id: str) -> list[str]: """Get known session IDs for a user across the local rooms and shared session pool.""" session_ids = set(get_session_ids_from_room(f'user:{user_id}')) - session_ids.update(sid for sid, entry in SESSION_POOL.items() if entry and entry.get('id') == user_id) + for batch in session_pool_batches(): + session_ids.update(sid for sid, entry in batch if entry.get('id') == user_id) + await asyncio.sleep(0) # don't hold the loop for the whole pool return list(session_ids) @@ -354,7 +365,7 @@ async def disconnect_user_sessions(user_id: str): The client will automatically reconnect and re-authenticate with fresh data from the database. """ - session_ids = get_session_ids_by_user_id(user_id) + session_ids = await get_session_ids_by_user_id(user_id) for sid in session_ids: try: await sio.disconnect(sid) @@ -903,12 +914,14 @@ async def disconnect(sid, reason=None): pass # 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] + try: + del USAGE_POOL[model_id] + except KeyError: + pass else: USAGE_POOL[model_id] = connections