From b780d5c5561c72a65d668359a788970eae0b25ce Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 15 Feb 2026 18:41:16 -0600 Subject: [PATCH] refac --- backend/open_webui/main.py | 2 + backend/open_webui/socket/main.py | 92 ++++++++++++++++++++++++------ backend/open_webui/socket/utils.py | 34 +++++++++++ 3 files changed, 110 insertions(+), 18 deletions(-) diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index cdd6a65173..3a14ade8bc 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -64,6 +64,7 @@ from open_webui.socket.main import ( MODELS, app as socket_app, periodic_usage_pool_cleanup, + periodic_session_pool_cleanup, get_event_emitter, get_models_in_use, ) @@ -635,6 +636,7 @@ async def lifespan(app: FastAPI): limiter.total_tokens = THREAD_POOL_SIZE asyncio.create_task(periodic_usage_pool_cleanup()) + asyncio.create_task(periodic_session_pool_cleanup()) if app.state.config.ENABLE_BASE_MODELS_CACHE: await get_all_models( diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index b43c56b4e6..78df66b8dc 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -99,6 +99,7 @@ else: # Timeout duration in seconds TIMEOUT_DURATION = 3 +SESSION_POOL_TIMEOUT = 120 # seconds without heartbeat before session is reaped # Dictionary to maintain the user pool @@ -147,6 +148,17 @@ if WEBSOCKET_MANAGER == "redis": aquire_func = clean_up_lock.aquire_lock renew_func = clean_up_lock.renew_lock release_func = clean_up_lock.release_lock + + session_cleanup_lock = RedisLock( + redis_url=WEBSOCKET_REDIS_URL, + lock_name=f"{REDIS_KEY_PREFIX}:session_cleanup_lock", + timeout_secs=WEBSOCKET_REDIS_LOCK_TIMEOUT, + redis_sentinels=redis_sentinels, + redis_cluster=WEBSOCKET_REDIS_CLUSTER, + ) + session_aquire_func = session_cleanup_lock.aquire_lock + session_renew_func = session_cleanup_lock.renew_lock + session_release_func = session_cleanup_lock.release_lock else: MODELS = {} @@ -154,6 +166,7 @@ else: USAGE_POOL = {} aquire_func = release_func = renew_func = lambda: True + session_aquire_func = session_release_func = session_renew_func = lambda: True YDOC_MANAGER = YdocManager( @@ -162,6 +175,31 @@ YDOC_MANAGER = YdocManager( ) +async def periodic_session_pool_cleanup(): + """Reap orphaned SESSION_POOL entries that missed heartbeats (e.g. crashed instance).""" + if not session_aquire_func(): + log.debug("Session cleanup lock held by another node. Skipping.") + return + + try: + while True: + if not session_renew_func(): + log.error("Unable to renew session cleanup lock. Exiting.") + return + + 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')})" + ) + del SESSION_POOL[sid] + await asyncio.sleep(SESSION_POOL_TIMEOUT) + finally: + session_release_func() + + async def periodic_usage_pool_cleanup(): max_retries = 2 retry_delay = random.uniform( @@ -313,15 +351,18 @@ async def connect(sid, environ, auth): user = Users.get_user_by_id(data["id"]) if user: - SESSION_POOL[sid] = user.model_dump( - exclude=[ - "profile_image_url", - "profile_banner_image_url", - "date_of_birth", - "bio", - "gender", - ] - ) + SESSION_POOL[sid] = { + **user.model_dump( + exclude=[ + "profile_image_url", + "profile_banner_image_url", + "date_of_birth", + "bio", + "gender", + ] + ), + "last_seen_at": int(time.time()), + } await sio.enter_room(sid, f"user:{user.id}") @@ -340,15 +381,18 @@ async def user_join(sid, data): if not user: return - SESSION_POOL[sid] = user.model_dump( - exclude=[ - "profile_image_url", - "profile_banner_image_url", - "date_of_birth", - "bio", - "gender", - ] - ) + SESSION_POOL[sid] = { + **user.model_dump( + exclude=[ + "profile_image_url", + "profile_banner_image_url", + "date_of_birth", + "bio", + "gender", + ] + ), + "last_seen_at": int(time.time()), + } await sio.enter_room(sid, f"user:{user.id}") @@ -366,6 +410,7 @@ async def user_join(sid, data): async def heartbeat(sid, data): user = SESSION_POOL.get(sid) if user: + SESSION_POOL[sid] = {**user, "last_seen_at": int(time.time())} Users.update_last_active_by_id(user["id"]) @@ -709,6 +754,17 @@ async def disconnect(sid): if sid in SESSION_POOL: user = SESSION_POOL[sid] 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: pass diff --git a/backend/open_webui/socket/utils.py b/backend/open_webui/socket/utils.py index 327348626a..c33af2e71d 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -118,6 +118,8 @@ class RedisDict: class YdocManager: + COMPACTION_THRESHOLD = 500 + def __init__( self, redis=None, @@ -133,10 +135,42 @@ class YdocManager: if self._redis: redis_key = f"{self._redis_key_prefix}:{document_id}:updates" await self._redis.rpush(redis_key, json.dumps(list(update))) + list_len = await self._redis.llen(redis_key) + if list_len >= self.COMPACTION_THRESHOLD: + await self._compact_updates_redis(document_id) else: if document_id not in self._updates: self._updates[document_id] = [] self._updates[document_id].append(update) + if len(self._updates[document_id]) >= self.COMPACTION_THRESHOLD: + self._compact_updates_memory(document_id) + + async def _compact_updates_redis(self, document_id: str): + """Rolling compaction: squash oldest half into one snapshot.""" + redis_key = f"{self._redis_key_prefix}:{document_id}:updates" + all_updates = await self._redis.lrange(redis_key, 0, -1) + if len(all_updates) <= 1: + return + mid = len(all_updates) // 2 + ydoc = Y.Doc() + for raw in all_updates[:mid]: + ydoc.apply_update(bytes(json.loads(raw))) + snapshot = json.dumps(list(ydoc.get_update())) + pipe = self._redis.pipeline() + pipe.delete(redis_key) + pipe.rpush(redis_key, snapshot, *all_updates[mid:]) + await pipe.execute() + + def _compact_updates_memory(self, document_id: str): + """Rolling compaction: squash oldest half into one snapshot.""" + updates = self._updates.get(document_id, []) + if len(updates) <= 1: + return + mid = len(updates) // 2 + ydoc = Y.Doc() + for update in updates[:mid]: + ydoc.apply_update(bytes(update)) + self._updates[document_id] = [ydoc.get_update()] + updates[mid:] async def get_updates(self, document_id: str) -> List[bytes]: document_id = document_id.replace(":", "_")