1 Commits

Author SHA1 Message Date
fawney19
41719a00e7 refactor: 改进分布式任务锁的清理策略
实现两种锁清理模式:
- 单实例模式(默认):启动时使用 Lua 脚本原子性清理旧锁,解决 worker 重启时���锁残留问题
- 多实例模式:使用 NX 选项竞争锁,依赖 TTL 处理异常退出

可通过 SINGLE_INSTANCE_MODE 环境变量控制模式选择。
2025-12-28 21:34:43 +08:00

View File

@@ -1,8 +1,16 @@
"""分布式任务协调器,确保仅有一个 worker 执行特定任务"""
"""分布式任务协调器,确保仅有一个 worker 执行特定任务
锁清理策略:
- 单实例模式(默认):启动时使用原子操作清理旧锁并获取新锁
- 多实例模式:使用 NX 选项竞争锁,依赖 TTL 处理异常退出
使用方式:
- 默认行为:启动时清理旧锁(适用于单机部署)
- 多实例部署:设置 SINGLE_INSTANCE_MODE=false 禁用启动清理
"""
from __future__ import annotations
import asyncio
import os
import pathlib
import uuid
@@ -19,6 +27,10 @@ except ImportError: # pragma: no cover - Windows 环境
class StartupTaskCoordinator:
"""利用 Redis 或文件锁,保证任务只在单个进程/实例中运行"""
# 类级别标记:在当前进程中是否已尝试过启动清理
# 注意:这在 fork 模式下每个 worker 都是独立的
_startup_cleanup_attempted = False
def __init__(self, redis_client=None, lock_dir: Optional[str] = None):
self.redis = redis_client
self._tokens: Dict[str, str] = {}
@@ -26,6 +38,8 @@ class StartupTaskCoordinator:
self._lock_dir = pathlib.Path(lock_dir or os.getenv("TASK_LOCK_DIR", "./.locks"))
if not self._lock_dir.exists():
self._lock_dir.mkdir(parents=True, exist_ok=True)
# 单实例模式:启动时清理旧锁(适用于单机部署,避免残留锁问题)
self._single_instance_mode = os.getenv("SINGLE_INSTANCE_MODE", "true").lower() == "true"
def _redis_key(self, name: str) -> str:
return f"task_lock:{name}"
@@ -36,12 +50,51 @@ class StartupTaskCoordinator:
if self.redis:
token = str(uuid.uuid4())
try:
acquired = await self.redis.set(self._redis_key(name), token, nx=True, ex=ttl)
if acquired:
self._tokens[name] = token
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
return True
return False
if self._single_instance_mode:
# 单实例模式:使用 Lua 脚本原子性地"清理旧锁 + 竞争获取"
# 只有当锁不存在或成功获取时才返回 1
# 这样第一个执行的 worker 会清理旧锁并获取,后续 worker 会正常竞争
script = """
local key = KEYS[1]
local token = ARGV[1]
local ttl = tonumber(ARGV[2])
local startup_key = KEYS[1] .. ':startup'
-- 检查是否已有 worker 执行过启动清理
local cleaned = redis.call('GET', startup_key)
if not cleaned then
-- 第一个 worker删除旧锁标记已清理
redis.call('DEL', key)
redis.call('SET', startup_key, '1', 'EX', 60)
end
-- 尝试获取锁NX 模式)
local result = redis.call('SET', key, token, 'NX', 'EX', ttl)
if result then
return 1
end
return 0
"""
result = await self.redis.eval(
script, 2,
self._redis_key(name), self._redis_key(name),
token, ttl
)
if result == 1:
self._tokens[name] = token
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
return True
return False
else:
# 多实例模式:直接使用 NX 选项竞争锁
acquired = await self.redis.set(
self._redis_key(name), token, nx=True, ex=ttl
)
if acquired:
self._tokens[name] = token
logger.info(f"任务 {name} 通过 Redis 锁独占执行")
return True
return False
except Exception as exc: # pragma: no cover - Redis 异常回退
logger.warning(f"Redis 锁获取失败,回退到文件锁: {exc}")