diff --git a/backend/open_webui/socket/utils.py b/backend/open_webui/socket/utils.py index 010d1f34f6..5e72753f42 100644 --- a/backend/open_webui/socket/utils.py +++ b/backend/open_webui/socket/utils.py @@ -7,8 +7,8 @@ import json import uuid import pycrdt as Y -from open_webui.utils.redis import get_redis_connection from open_webui.env import REDIS_KEY_PREFIX +from open_webui.utils.redis import get_redis_connection YDOC_KEY_PREFIX = f'{REDIS_KEY_PREFIX}:ydoc:documents' @@ -16,6 +16,19 @@ YDOC_KEY_PREFIX = f'{REDIS_KEY_PREFIX}:ydoc:documents' class RedisLock: """Distributed lock backed by a Redis SET with NX/EX semantics.""" + _RENEW_SCRIPT = """ + if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('expire', KEYS[1], ARGV[2]) + end + return 0 + """ + _RELEASE_SCRIPT = """ + if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('del', KEYS[1]) + end + return 0 + """ + def __init__( self, redis_url, @@ -41,13 +54,10 @@ 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) + return bool(self.redis.eval(self._RENEW_SCRIPT, 1, self.lock_name, self.lock_id, self.timeout_secs)) 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) + self.redis.eval(self._RELEASE_SCRIPT, 1, self.lock_name, self.lock_id) class RedisDict: