mirror of
https://github.com/fawney19/Aether.git
synced 2026-01-02 15:52:26 +08:00
refactor: 改进分布式任务锁的清理策略
实现两种锁清理模式: - 单实例模式(默认):启动时使用 Lua 脚本原子性清理旧锁,解决 worker 重启时���锁残留问题 - 多实例模式:使用 NX 选项竞争锁,依赖 TTL 处理异常退出 可通过 SINGLE_INSTANCE_MODE 环境变量控制模式选择。
This commit is contained in:
@@ -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,7 +50,46 @@ 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 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 锁独占执行")
|
||||
|
||||
Reference in New Issue
Block a user