Merge pull request #7964 from jk-f5/fix/cleanup

Fix Redis Lock Expiration in periodic_usage_pool_cleanup
This commit is contained in:
Timothy Jaeryang Baek 2024-12-19 14:45:18 -08:00 committed by GitHub
commit f133353734
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -49,7 +49,6 @@ TIMEOUT_DURATION = 3
# Dictionary to maintain the user pool # Dictionary to maintain the user pool
run_cleanup = True
if WEBSOCKET_MANAGER == "redis": if WEBSOCKET_MANAGER == "redis":
log.debug("Using Redis to manage websockets.") log.debug("Using Redis to manage websockets.")
SESSION_POOL = RedisDict("open-webui:session_pool", redis_url=WEBSOCKET_REDIS_URL) SESSION_POOL = RedisDict("open-webui:session_pool", redis_url=WEBSOCKET_REDIS_URL)
@ -61,18 +60,18 @@ if WEBSOCKET_MANAGER == "redis":
lock_name="usage_cleanup_lock", lock_name="usage_cleanup_lock",
timeout_secs=TIMEOUT_DURATION * 2, timeout_secs=TIMEOUT_DURATION * 2,
) )
run_cleanup = clean_up_lock.aquire_lock() aquire_func = clean_up_lock.aquire_lock
renew_func = clean_up_lock.renew_lock renew_func = clean_up_lock.renew_lock
release_func = clean_up_lock.release_lock release_func = clean_up_lock.release_lock
else: else:
SESSION_POOL = {} SESSION_POOL = {}
USER_POOL = {} USER_POOL = {}
USAGE_POOL = {} USAGE_POOL = {}
release_func = renew_func = lambda: True aquire_func = release_func = renew_func = lambda: True
async def periodic_usage_pool_cleanup(): async def periodic_usage_pool_cleanup():
if not run_cleanup: if not aquire_func():
log.debug("Usage pool cleanup lock already exists. Not running it.") log.debug("Usage pool cleanup lock already exists. Not running it.")
return return
log.debug("Running periodic_usage_pool_cleanup") log.debug("Running periodic_usage_pool_cleanup")