diff --git a/backend/open_webui/routers/channels.py b/backend/open_webui/routers/channels.py index 0b20bc6a17..3c165f16ff 100644 --- a/backend/open_webui/routers/channels.py +++ b/backend/open_webui/routers/channels.py @@ -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 diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index d95158e0bd..00af651d6a 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -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] diff --git a/backend/open_webui/socket/utils.py b/backend/open_webui/socket/utils.py index 09fa47de15..97df5d4c38 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -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()