49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
"""Redis 分布式锁自动续期工具。"""
|
|
|
|
import asyncio
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
|
|
class RedisLockRenewer:
|
|
"""在后台周期性续期 Redis 分布式锁，防止长处理超时。
|
|
|
|
用法::
|
|
|
|
renewer = RedisLockRenewer(redis, lock_key, ttl=30, interval=10)
|
|
async with renewer:
|
|
... # 长处理
|
|
"""
|
|
|
|
def __init__(self, redis, lock_key: str, *, ttl: int, interval: int) -> None:
|
|
self._redis = redis
|
|
self._lock_key = lock_key
|
|
self._ttl = ttl
|
|
self._interval = interval
|
|
self._stop_event = asyncio.Event()
|
|
self._task: asyncio.Task | None = None
|
|
|
|
async def __aenter__(self) -> RedisLockRenewer:
|
|
self._task = asyncio.create_task(self._renew_loop())
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
self._stop_event.set()
|
|
if self._task is not None:
|
|
self._task.cancel()
|
|
try:
|
|
await self._task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._task = None
|
|
|
|
async def _renew_loop(self) -> None:
|
|
while not self._stop_event.is_set():
|
|
try:
|
|
await asyncio.wait_for(self._stop_event.wait(), timeout=self._interval)
|
|
except TimeoutError:
|
|
try:
|
|
await self._redis.expire(self._lock_key, self._ttl)
|
|
except Exception:
|
|
logger.warning("Failed to renew lock %s", self._lock_key)
|