From f5f0bc0d4d714f13dad19a61e3edd3510600028c Mon Sep 17 00:00:00 2001 From: Jacob Leksan Date: Tue, 31 Mar 2026 15:15:05 -0400 Subject: [PATCH] fix: atomic Redis lock release and renew Use Lua for compare-and-delete (release) and compare-and-expire (renew) so lock ownership is verified atomically and stale renews cannot steal another holder's token after TTL races. Made-with: Cursor --- backend/open_webui/socket/main.py | 14 +++++----- backend/open_webui/socket/utils.py | 41 +++++++++++++++++++++++++----- 2 files changed, 42 insertions(+), 13 deletions(-) diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 33c9ffea05..2633e05a5d 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -140,7 +140,7 @@ if WEBSOCKET_MANAGER == 'redis': redis_sentinels=redis_sentinels, redis_cluster=WEBSOCKET_REDIS_CLUSTER, ) - aquire_func = clean_up_lock.aquire_lock + acquire_func = clean_up_lock.acquire_lock renew_func = clean_up_lock.renew_lock release_func = clean_up_lock.release_lock @@ -151,7 +151,7 @@ if WEBSOCKET_MANAGER == 'redis': redis_sentinels=redis_sentinels, redis_cluster=WEBSOCKET_REDIS_CLUSTER, ) - session_aquire_func = session_cleanup_lock.aquire_lock + session_acquire_func = session_cleanup_lock.acquire_lock session_renew_func = session_cleanup_lock.renew_lock session_release_func = session_cleanup_lock.release_lock else: @@ -160,8 +160,8 @@ else: SESSION_POOL = {} USAGE_POOL = {} - aquire_func = release_func = renew_func = lambda: True - session_aquire_func = session_release_func = session_renew_func = lambda: True + acquire_func = release_func = renew_func = lambda: True + session_acquire_func = session_release_func = session_renew_func = lambda: True YDOC_MANAGER = YdocManager( @@ -172,8 +172,8 @@ 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.') + if not session_acquire_func(): + log.debug("Session cleanup lock held by another node. Skipping.") return try: @@ -197,7 +197,7 @@ 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(): + if acquire_func(): break else: if attempt < max_retries: diff --git a/backend/open_webui/socket/utils.py b/backend/open_webui/socket/utils.py index 16b0cc3855..7807f738b6 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -5,6 +5,25 @@ from open_webui.env import REDIS_KEY_PREFIX from typing import Optional, List, Tuple import pycrdt as Y +# Redis has no single command for compare-and-delete or compare-and-expire; Lua runs atomically. +_RELEASE_LOCK_IF_OWNER_SCRIPT = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +""" + +# SET key id XX EX would only verify the key exists, not that we still own it — another +# holder's token could be overwritten after TTL races. Only extend TTL when value matches. +_RENEW_LOCK_IF_OWNER_SCRIPT = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("expire", KEYS[1], tonumber(ARGV[2])) +else + return 0 +end +""" + class RedisLock: def __init__( @@ -25,20 +44,30 @@ class RedisLock: redis_cluster=redis_cluster, decode_responses=True, ) + self._release_if_owner = self.redis.register_script(_RELEASE_LOCK_IF_OWNER_SCRIPT) + self._renew_if_owner = self.redis.register_script(_RENEW_LOCK_IF_OWNER_SCRIPT) - def aquire_lock(self): + def acquire_lock(self): # nx=True will only set this key if it _hasn't_ already been set self.lock_obtained = self.redis.set(self.lock_name, self.lock_id, nx=True, ex=self.timeout_secs) 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) + # Must not use SET ... XX alone: that only checks key existence, so a stale renew + # could steal the lock from another holder after expiry/reacquisition. + result = self._renew_if_owner( + keys=[self.lock_name], + args=[self.lock_id, str(self.timeout_secs)], + ) + ok = result == 1 + self.lock_obtained = ok + return ok 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) + try: + self._release_if_owner(keys=[self.lock_name], args=[self.lock_id]) + except Exception: + pass # Best-effort; TTL will clear the key if we crash class RedisDict: