mirror of
https://github.com/open-webui/open-webui.git
synced 2026-11-14 09:21:22 +00:00
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.
This commit is contained in:
Classic298
parent
75fe65c2f8
commit
fabe845337
2 changed files with 60 additions and 66 deletions
|
|
@ -173,7 +173,7 @@ YDOC_MANAGER = YdocManager(
|
|||
)
|
||||
|
||||
|
||||
def session_pool_batches():
|
||||
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()
|
||||
|
|
@ -185,53 +185,49 @@ async def periodic_session_pool_cleanup():
|
|||
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:
|
||||
try:
|
||||
if not session_aquire_func():
|
||||
log.debug('Session cleanup lock held by another node. Retrying.')
|
||||
await asyncio.sleep(retry_delay)
|
||||
continue
|
||||
if not session_aquire_func():
|
||||
log.debug('Session cleanup lock held by another node. Retrying.')
|
||||
await asyncio.sleep(retry_delay)
|
||||
continue
|
||||
|
||||
try:
|
||||
try:
|
||||
while True:
|
||||
if not session_renew_func():
|
||||
log.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.')
|
||||
break
|
||||
|
||||
now = int(time.time())
|
||||
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.pop_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
|
||||
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
|
||||
|
||||
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)
|
||||
if lock_lost:
|
||||
break
|
||||
finally:
|
||||
session_release_func()
|
||||
|
||||
|
||||
async def periodic_usage_pool_cleanup():
|
||||
|
|
@ -297,7 +293,7 @@ def get_user_id_from_session_pool(sid):
|
|||
return None
|
||||
|
||||
|
||||
async def get_socket_session_user(sid):
|
||||
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')
|
||||
|
|
@ -311,16 +307,14 @@ def get_session_ids_from_room(room):
|
|||
return list(members) if members else []
|
||||
|
||||
|
||||
async def get_session_ids_by_user_id(user_id: str) -> list[str]:
|
||||
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}'))
|
||||
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
|
||||
session_ids.update(sid for sid, entry in SESSION_POOL.items() if entry and entry.get('id') == user_id)
|
||||
return list(session_ids)
|
||||
|
||||
|
||||
async def get_user_ids_from_room(room):
|
||||
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}
|
||||
|
||||
|
|
@ -365,7 +359,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 = await get_session_ids_by_user_id(user_id)
|
||||
session_ids = get_session_ids_by_user_id(user_id)
|
||||
for sid in session_ids:
|
||||
try:
|
||||
await sio.disconnect(sid)
|
||||
|
|
@ -908,24 +902,22 @@ async def yjs_awareness_update(sid, data):
|
|||
|
||||
@sio.event
|
||||
async def disconnect(sid, reason=None):
|
||||
try:
|
||||
if sid in SESSION_POOL:
|
||||
del SESSION_POOL[sid]
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
# 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:
|
||||
# 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:
|
||||
del USAGE_POOL[model_id]
|
||||
except KeyError:
|
||||
pass
|
||||
else:
|
||||
USAGE_POOL[model_id] = connections
|
||||
else:
|
||||
USAGE_POOL[model_id] = connections
|
||||
|
||||
await YDOC_MANAGER.remove_user_from_all_documents(sid)
|
||||
await YDOC_MANAGER.remove_user_from_all_documents(sid)
|
||||
else:
|
||||
pass
|
||||
# print(f"Unknown session ID {sid} disconnected")
|
||||
|
||||
|
||||
async def _make_channel_emitter(request_info):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -108,13 +109,13 @@ class RedisDict:
|
|||
"""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)
|
||||
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 discard(self, *keys):
|
||||
def pop_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)
|
||||
|
|
@ -152,7 +153,8 @@ 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)
|
||||
self.discard(*keys_to_remove)
|
||||
if keys_to_remove:
|
||||
self.redis.hdel(self.name, *keys_to_remove)
|
||||
|
||||
self._last_signature = signature
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue