scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-09-10 22:43:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

fix: harden distributed lock lifecycle with atomic Lua scripts and resilient retry

Browse source 
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.
This commit is contained in:
DrMelone 2026-02-23 22:07:13 +01:00
parent f6b85700ea
commit 2ca7c3bfe6
2 changed files with 85 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

113
backend/open_webui/socket/main.py
View file

@ -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(

33
backend/open_webui/socket/utils.py
View file

@ -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: