93 lines
3.3 KiB
Python
93 lines
3.3 KiB
Python
|
"""基于缓存的分布式排他锁（租约 + 续约）。
|
|||
|
|
|
|||
|
|
利用后端提供的 ``set_nx`` 原子原语实现跨进程/跨主机的排他锁：
|
|||
|
|
|
|||
|
|
- ``acquire``（SET NX）保证全局只有一方成功，key 带 TTL 租约，防止持有方宕机导致死锁。
|
|||
|
|
- 持有方通过后台任务定期 ``touch`` 续约；续约前先校验 token，避免误续别人抢到的锁。
|
|||
|
|
- ``release`` 同样先校验 token，仅当仍是自己持有时才删除，杜绝误删他人锁。
|
|||
|
|
|
|||
|
|
适用后端：
|
|||
|
|
- ``redis``：真正的跨主机共享，多实例部署的推荐后端。
|
|||
|
|
- ``diskcache``：进程安全（SQLite），可在单机多进程/共享卷场景使用，不可跨无共享文件系统的主机。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import os
|
|||
|
|
import socket
|
|||
|
|
import uuid
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
from chongming_cache.cache import Cache
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DistributedLock:
|
|||
|
|
"""租约式分布式排他锁，基于集群共享缓存。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
cache: ``chongming_cache.Cache`` 实例（后端可为 redis 或 diskcache）。
|
|||
|
|
key: 锁标识（会被后端前缀命名空间隔离）。
|
|||
|
|
token: 持锁者标识；缺省生成本机唯一 token。
|
|||
|
|
ttl: 租约有效期（秒），期间未续约则锁自动失效。
|
|||
|
|
renew_interval: 续约周期（秒），须小于 ttl。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
cache: Cache,
|
|||
|
|
key: str,
|
|||
|
|
token: Any | None = None,
|
|||
|
|
ttl: float = 30.0,
|
|||
|
|
renew_interval: float = 10.0,
|
|||
|
|
):
|
|||
|
|
self._cache = cache
|
|||
|
|
self._key = key
|
|||
|
|
self._token = token if token is not None else self._new_token()
|
|||
|
|
self._ttl = ttl
|
|||
|
|
self._renew_interval = min(renew_interval, ttl / 2)
|
|||
|
|
self._renew_task: asyncio.Task | None = None
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _new_token() -> str:
|
|||
|
|
return f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4().hex}"
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def token(self) -> Any:
|
|||
|
|
return self._token
|
|||
|
|
|
|||
|
|
def acquire(self) -> bool:
|
|||
|
|
"""尝试获取锁。成功（全局唯一持有者）返回 True，否则返回 False。"""
|
|||
|
|
return bool(self._cache.backend.set_nx(self._key, self._token, expire=self._ttl))
|
|||
|
|
|
|||
|
|
def _holding(self) -> bool:
|
|||
|
|
return self._cache.get(self._key, None) == self._token
|
|||
|
|
|
|||
|
|
def start(self) -> None:
|
|||
|
|
"""获取成功后启动后台续约；已启动则忽略。"""
|
|||
|
|
if self._renew_task is None:
|
|||
|
|
self._renew_task = asyncio.create_task(self._renew_loop())
|
|||
|
|
|
|||
|
|
async def _renew_loop(self):
|
|||
|
|
while True:
|
|||
|
|
try:
|
|||
|
|
await asyncio.sleep(self._renew_interval)
|
|||
|
|
except asyncio.CancelledError:
|
|||
|
|
return
|
|||
|
|
# 仅在自己仍持有（token 匹配）时续约；锁已丢失则停止续约
|
|||
|
|
if not self._holding():
|
|||
|
|
return
|
|||
|
|
self._cache.touch(self._key, self._ttl)
|
|||
|
|
|
|||
|
|
async def stop(self) -> None:
|
|||
|
|
"""取消后台续约任务。"""
|
|||
|
|
if self._renew_task is not None:
|
|||
|
|
self._renew_task.cancel()
|
|||
|
|
try:
|
|||
|
|
await self._renew_task
|
|||
|
|
except (asyncio.CancelledError, Exception):
|
|||
|
|
pass
|
|||
|
|
self._renew_task = None
|
|||
|
|
|
|||
|
|
def release(self) -> None:
|
|||
|
|
"""受保护释放：仅当仍由本 token 持有时才删除锁。"""
|
|||
|
|
if self._holding():
|
|||
|
|
self._cache.delete(self._key)
|