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 0998682dd8..de6197f7ff 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -175,48 +175,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) 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()) - 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 - - 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(): @@ -283,32 +300,32 @@ 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) 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) -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]): @@ -351,7 +368,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) @@ -364,7 +381,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 +463,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 +536,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 +576,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 +630,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 +790,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 +860,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 +892,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: @@ -907,23 +916,24 @@ 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, connections in list(USAGE_POOL.items()): + if sid in connections: + del connections[sid] + if not connections: + try: + del USAGE_POOL[model_id] + except KeyError: + pass + 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 09fa47de15..ad39893bba 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -112,6 +112,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.clear() @@ -140,8 +155,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) if self._signature_name: self.redis.set(self._signature_name, signature)