Merge pull request #758 from SuanmoSuanyangTechnology/refactor/redis-lock

fix(redis_lock): refactor RedisFairLock to use ZSET for queue management and fix loop shutdown

api/app/tasks.py

@@ -1176,6 +1176,7 @@ def write_message_task(
 
     redis_client = get_sync_redis_client()
     lock = None
+    loop = None
     if redis_client is not None:
         lock = RedisFairLock(
             key=f"memory_write:{end_user_id}",
@@ -1196,6 +1197,7 @@ def write_message_task(
             }
 
     try:
+        task_start_time = int(time.time())
         loop = set_asyncio_event_loop()
 
         result = loop.run_until_complete(_run())
@@ -1219,6 +1221,7 @@ def write_message_task(
         return {
             "status": "SUCCESS",
             "result": result,
+            "start_at": task_start_time,
             "end_user_id": end_user_id,
             "config_id": config_id,
             "elapsed_time": elapsed_time,
@@ -1252,7 +1255,8 @@ def write_message_task(
                 logger.warning(f"[CELERY WRITE] 释放锁失败: {e}")
         # Gracefully shutdown the event loop to prevent
         # 'RuntimeError: Event loop is closed' from httpx.AsyncClient.__del__
-        _shutdown_loop_gracefully(loop)
+        if loop:
+            _shutdown_loop_gracefully(loop)
 
 
 # unused task

api/app/utils/redis_lock.py

@@ -1,14 +1,15 @@
-import redis
-import uuid
-import time
 import threading
+import time
+import uuid
+
+import redis
 
 UNLOCK_SCRIPT = """
 if redis.call("get", KEYS[1]) == ARGV[1] then
     return redis.call("del", KEYS[1])
-else
-    return 0
 end
+
+return 0
 """
 
 RENEW_SCRIPT = """
@@ -19,38 +20,44 @@ else
 end
 """
 
-CLEANUP_DEAD_HEAD_SCRIPT = """
+ACQUIRE_SCRIPT = """
 local queue_key = KEYS[1]
 local lock_key = KEYS[2]
 
-local first = redis.call("lindex", queue_key, 0)
-if not first then
-    return 0
+local client_id = ARGV[1]
+local expire = tonumber(ARGV[2])
+local time_out = tonumber(ARGV[3])
+
+local now = tonumber(redis.call("time")[1])
+
+if redis.call("zscore", queue_key, client_id) == false then
+    redis.call("zadd", queue_key, now, client_id)
 end
 
-if redis.call("exists", lock_key) == 1 then
-    return 0
+local expired = redis.call("zrangebyscore", queue_key, 0, now - time_out)
+
+for _, v in ipairs(expired) do
+    redis.call("zrem", queue_key, v)
 end
 
-redis.call("lpop", queue_key)
-return 1
-"""
+local first = redis.call("zrange", queue_key, 0, 0)[1]
+if first == client_id then
 
-SAFE_RELEASE_QUEUE_SCRIPT = """
-local queue_key = KEYS[1]
-local value = ARGV[1]
+    if redis.call("set", lock_key, client_id, "NX", "EX", expire) then
+        redis.call("zrem", queue_key, client_id)
+        return 1
+    end
 
-local first = redis.call("lindex", queue_key, 0)
-if first == value then
-    redis.call("lpop", queue_key)
-    return 1
+    if redis.call("get", lock_key) == client_id then
+        redis.call("expire", lock_key, expire)
+        return 1
+    end
 end
 return 0
 """
 
 
 def _ensure_str(val):
-    """统一将 Redis 返回值转为 str，兼容 decode_responses=True/False"""
     if val is None:
         return None
     if isinstance(val, bytes):
@@ -59,18 +66,21 @@ def _ensure_str(val):
 
 
 class RedisFairLock:
+    # ZOMBIE CLEAN BUFFER
+    CLEANUP_BUFFER = 30
+
     def __init__(
             self,
             key: str,
             redis_client: redis.StrictRedis,
             expire: int = 30,
-            retry_interval: float = 0.05,
+            retry_interval: float = 1,
             timeout: float = 600,
             auto_renewal: bool = True
     ):
         self.key = key
-        self.queue_key = f"{key}:queue"
-        self.value = str(uuid.uuid4())
+        self.queue_key = f"{key}:zset"
+        self.value = f"{uuid.uuid4().hex}:{int(time.time())}"
         self.expire = expire
         self.retry_interval = retry_interval
         self.timeout = timeout
@@ -83,25 +93,25 @@ class RedisFairLock:
     def acquire(self):
         start = time.time()
 
-        self.redis.rpush(self.queue_key, self.value)
-
         while True:
-            first = _ensure_str(self.redis.lindex(self.queue_key, 0))
+            ok = self.redis.eval(
+                ACQUIRE_SCRIPT,
+                2,
+                self.queue_key,
+                self.key,
+                self.value,
+                str(self.expire),
+                str(self.timeout + self.CLEANUP_BUFFER)
+            )
 
-            if first == self.value:
-                ok = self.redis.set(self.key, self.value, nx=True, ex=self.expire)
-                if ok:
-                    self._locked = True
-
-                    if self.auto_renewal:
-                        self._start_renewal()
-                    return True
-
-            if first:
-                self.redis.eval(CLEANUP_DEAD_HEAD_SCRIPT, 2, self.queue_key, self.key)
+            if ok == 1:
+                self._locked = True
+                if self.auto_renewal:
+                    self._start_renewal()
+                return True
 
             if time.time() - start > self.timeout:
-                self.redis.lrem(self.queue_key, 0, self.value)
+                self.redis.zrem(self.queue_key,  self.value)
                 return False
 
             time.sleep(self.retry_interval)
@@ -112,13 +122,15 @@ class RedisFairLock:
             if self._stop_renew.is_set():
                 break
 
-            self.redis.eval(
+            success = self.redis.eval(
                 RENEW_SCRIPT,
                 1,
                 self.key,
                 self.value,
                 str(self.expire)
             )
+            if not success:
+                break
 
     def _start_renewal(self):
         self._stop_renew = threading.Event()
@@ -139,8 +151,6 @@ class RedisFairLock:
 
         self.redis.eval(UNLOCK_SCRIPT, 1, self.key, self.value)
 
-        self.redis.eval(SAFE_RELEASE_QUEUE_SCRIPT, 1, self.queue_key, self.value)
-
         self._locked = False
 
     def __enter__(self):