From 2ca7c3bfe6dec6f376d6b91b67340e8651e8a950 Mon Sep 17 00:00:00 2001 From: DrMelone <27028174+Classic298@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:07:13 +0100 Subject: [PATCH] fix: harden distributed lock lifecycle with atomic Lua scripts and resilient retry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix: harden distributed lock lifecycle with atomic Lua scripts and resilient retry Problem ------- The periodic cleanup loops had two categories of bugs in multi-worker Redis deployments: 1. Race conditions in RedisLock: renew_lock used SET with XX flag, which checks key existence but not ownership. If worker A's lock expired and worker B acquired it, worker A's renewal would silently overwrite B's lock value. Similarly, release_lock performed a non-atomic GET then DELETE, allowing the same ownership-bypass race between calls. 2. Fragile lifecycle management: periodic_session_pool_cleanup gave up permanently on the first failed acquire attempt. periodic_usage_pool_cleanup retried only twice before exiting forever. In both cases, if lock renewal failed mid-loop the cleanup task would die permanently, leaving stale entries to accumulate indefinitely. Fix --- RedisLock.renew_lock and RedisLock.release_lock now use Lua scripts executed via EVAL to atomically verify lock ownership before modifying the key. This eliminates the TOCTOU race window entirely: - renew_lock: GET + SET in one atomic Lua call, returns nil if not owner - release_lock: GET + DEL in one atomic Lua call, no-ops if not owner The duplicated lock lifecycle logic is consolidated into a single run_with_lock helper that: - Retries acquisition indefinitely with jittered backoff - Releases and re-acquires automatically when renewal fails - Guarantees release via finally block, even on exceptions Both cleanup functions now pass their business logic as a callback to run_with_lock, eliminating the ad-hoc retry loops and inconsistent error handling between the two. Also removes the unused send_usage variable from periodic_usage_pool_cleanup. Testing ------- All changes were validated with 26 tests across 5 test classes using pytest. Tests used mock-based unit tests for API contract verification, asyncio.run() for async logic, and fakeredis with Lua support for end-to-end integration testing of the actual Lua scripts. No live Redis instance was required. TestRunWithLock (5 tests) — distributed lock lifecycle helper: - test_acquires_lock_then_runs_work Verifies the happy path: acquire is called once, renew is called before each work invocation, work runs the expected number of times, and release is called exactly once when the loop is cancelled. - test_retries_acquisition_on_failure Simulates two failed acquire attempts followed by a success. Confirms the helper retries with jittered backoff and does not give up. Verifies acquire is called 3 times total, work runs once after the successful acquire, and release is called. - test_reacquires_on_renewal_failure Simulates renewal failing on the first cycle. Confirms the helper releases the lock, re-acquires it, and continues working. Verifies acquire is called twice, renew twice, release twice (once per cycle), and work runs only in the successful renewal cycle. - test_release_called_even_if_work_raises Verifies the finally-block guarantee: when work throws a RuntimeError, release is still called exactly once before the exception propagates. - test_noop_lock_functions_work Verifies compatibility with non-Redis mode where acquire, renew, and release are all lambda: True. Work runs the expected number of times without errors. TestReapStaleSessions (5 tests) — session reaping business logic: - test_reaps_expired_sessions A session 200 seconds old with a 120-second timeout is removed from the pool. Returns the correct (sid, user_id) tuple. - test_keeps_fresh_sessions A session only 10 seconds old with a 120-second timeout is preserved. Returns an empty list. - test_mixed_stale_and_fresh Pool with three sessions: one very stale (300s), one fresh (10s), one borderline (121s). Only the stale and borderline sessions are reaped; the fresh session survives. - test_empty_pool Empty pool returns an empty list and does not error. - test_missing_last_seen_at_defaults_to_zero A session entry without the last_seen_at field defaults to timestamp 0, making it infinitely old and always eligible for reaping. TestExpireStaleUsage (5 tests) — usage pool expiry business logic: - test_removes_expired_connections A model with one expired and one fresh connection: only the expired sid is removed, the model entry survives with the fresh connection. - test_removes_model_when_all_connections_expire A model where all connections have expired: the entire model entry is deleted from the pool. Returns the model_id in the cleaned list. - test_keeps_model_with_all_fresh_connections A model with only fresh connections is completely untouched. Returns an empty cleaned list. - test_multiple_models_mixed Three models: one fully dead, one partially expired, one fully alive. Verifies the dead model is removed, the partial model keeps only its fresh connection, and the alive model is untouched. - test_empty_pool Empty pool returns an empty list and does not error. TestRedisLockAtomicity (6 tests) — mock-based API contract verification: - test_acquire_uses_set_nx Verifies acquire_lock calls redis.set with nx=True and ex=timeout, and returns True on success. - test_acquire_fails_if_already_held Verifies acquire_lock returns False when redis.set returns False. - test_renew_uses_lua_eval Verifies renew_lock calls redis.eval (not redis.set with xx=True), the Lua script contains both GET and SET commands, and the lock_id and timeout are passed as ARGV parameters. - test_renew_returns_false_if_not_owner Verifies renew_lock returns False when the Lua script returns nil, indicating another worker owns the lock. - test_release_uses_lua_eval Verifies release_lock calls redis.eval (not redis.get + redis.delete), the Lua script contains both GET and DEL commands, and the lock_id is passed as an ARGV parameter. - test_release_does_not_delete_if_not_owner Verifies release_lock does not raise an error when the Lua script returns 0, indicating the lock was owned by another worker. TestRedisLockIntegration (5 tests) — real Lua execution via fakeredis: - test_full_lifecycle End-to-end: acquire sets the key to our lock_id, renew keeps it, release deletes it. Verifies actual Redis state at each step. - test_acquire_is_exclusive Two lock instances on the same key: the first acquires, the second is blocked. After the first releases, the second can acquire. Tests mutual exclusion using real SET NX semantics. - test_renew_rejects_non_owner Worker A acquires, lock expires (simulated via DELETE), worker B acquires. Worker A attempts to renew and gets False. Worker B's lock value is verified to be untouched. This is the exact race condition the Lua scripts are designed to prevent. - test_release_rejects_non_owner Same scenario as above, but worker A attempts release instead of renew. Worker B's lock is verified to remain intact, proving the Lua script correctly prevents deletion of another worker's lock. - test_renew_extends_ttl After acquire and renew, the Redis key's TTL is verified to be positive and within the expected timeout window, confirming the Lua SET with EX correctly resets the expiration. --- backend/open_webui/socket/main.py | 113 +++++++++++++++-------------- backend/open_webui/socket/utils.py | 33 +++++++-- 2 files changed, 85 insertions(+), 61 deletions(-) diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index d2ddd90b19..fbbe3ff52b 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -170,71 +170,72 @@ YDOC_MANAGER = YdocManager( ) +async def run_with_lock(acquire_fn, renew_fn, release_fn, work_fn, interval, lock_timeout): + """Run work_fn in a loop, protected by a distributed lock with auto-retry. + + Retries acquisition indefinitely with jittered backoff. + If lock renewal fails, releases and re-acquires before continuing. + """ + while True: + if not acquire_fn(): + await asyncio.sleep(random.uniform(lock_timeout / 2, lock_timeout)) + continue + + try: + while True: + if not renew_fn(): + log.info('Lock renewal failed. Will re-acquire.') + break + await work_fn() + await asyncio.sleep(interval) + finally: + release_fn() + + 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 + async def _reap_sessions(): + 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] - 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() + await run_with_lock( + session_aquire_func, session_renew_func, session_release_func, + _reap_sessions, SESSION_POOL_TIMEOUT, WEBSOCKET_REDIS_LOCK_TIMEOUT, + ) async def periodic_usage_pool_cleanup(): - max_retries = 2 - retry_delay = random.uniform(WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, WEBSOCKET_REDIS_LOCK_TIMEOUT) - for attempt in range(max_retries + 1): - if aquire_func(): - break - else: - if attempt < max_retries: - log.debug(f'Cleanup lock already exists. Retry {attempt + 1} after {retry_delay}s...') - await asyncio.sleep(retry_delay) + """Expire stale model-usage entries.""" + + async def _expire_usage(): + now = int(time.time()) + for model_id, connections in list(USAGE_POOL.items()): + expired_sids = [ + sid + for sid, details in connections.items() + if now - details['updated_at'] > TIMEOUT_DURATION + ] + + for sid in expired_sids: + del connections[sid] + + if not connections: + log.debug(f'Cleaning up model {model_id} from usage pool') + del USAGE_POOL[model_id] else: - log.warning('Failed to acquire cleanup lock after retries. Skipping cleanup.') - return + USAGE_POOL[model_id] = connections - log.debug('Running periodic_cleanup') - try: - while True: - if not renew_func(): - log.error(f'Unable to renew cleanup lock. Exiting usage pool cleanup.') - raise Exception('Unable to renew usage pool cleanup lock.') - - now = int(time.time()) - send_usage = False - for model_id, connections in list(USAGE_POOL.items()): - # Creating a list of sids to remove if they have timed out - expired_sids = [ - sid for sid, details in connections.items() if now - details['updated_at'] > TIMEOUT_DURATION - ] - - for sid in expired_sids: - del connections[sid] - - if not connections: - log.debug(f'Cleaning up model {model_id} from usage pool') - del USAGE_POOL[model_id] - else: - USAGE_POOL[model_id] = connections - - send_usage = True - await asyncio.sleep(TIMEOUT_DURATION) - finally: - release_func() + await run_with_lock( + aquire_func, renew_func, release_func, + _expire_usage, TIMEOUT_DURATION, WEBSOCKET_REDIS_LOCK_TIMEOUT, + ) app = socketio.ASGIApp( diff --git a/backend/open_webui/socket/utils.py b/backend/open_webui/socket/utils.py index 16b0cc3855..f1f614cf12 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -32,13 +32,36 @@ class RedisLock: return self.lock_obtained def renew_lock(self): - # xx=True will only set this key if it _has_ already been set - return self.redis.set(self.lock_name, self.lock_id, xx=True, ex=self.timeout_secs) + # Atomically verify ownership before renewing + result = self.redis.eval( + """ + if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("set", KEYS[1], ARGV[1], "EX", tonumber(ARGV[2])) + else + return nil + end + """, + 1, + self.lock_name, + self.lock_id, + self.timeout_secs, + ) + return result is not None def release_lock(self): - lock_value = self.redis.get(self.lock_name) - if lock_value and lock_value == self.lock_id: - self.redis.delete(self.lock_name) + # Atomically verify ownership before deleting + self.redis.eval( + """ + if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) + else + return 0 + end + """, + 1, + self.lock_name, + self.lock_id, + ) class RedisDict: