modify: 排他锁

2026-08-27 09:42:21 +08:00
parent 9f7f71be7c
commit 0bdee195c1
7 changed files with 159 additions and 2 deletions
@@ -52,8 +52,41 @@ cache.close()                              # 释放连接
 with Cache.from_config(cfg) as cache: ...  # 上下文管理器，退出自动 close

 cache.client  # 底层客户端逃生口（diskcache.Cache / redis.Redis）
+
+# 原子写（分布式锁原语）：仅当 key 不存在时写入，成功返回 True，已存在返回 False，不覆盖
+ok = cache.backend.set_nx("leader", {"node": "gw-1"}, expire=30)
 ```

+### 分布式排他锁 `DistributedLock`
+
+基于 `backend.set_nx`（SET NX + TTL 租约）实现的**跨进程/跨主机排他锁**，用于集群协调（如保证全局只有一个 MQTT Bridge 消费）：
+
+```python
+import asyncio
+from chongming_cache import Cache, DistributedLock
+
+async def main():
+    cache = Cache.from_config(config["cache"])
+    lock = DistributedLock(cache, "mqtt_bridge", ttl=30, renew_interval=10)
+    if not lock.acquire():        # 全局唯一持锁方返回 True，其余返回 False
+        print("其他实例已持有，跳过")
+        return
+    lock.start()                  # 后台自动续约（防宕机/长暂停导致锁过期被抢）
+    try:
+        ...                       # 执行受保护的单例任务
+    finally:
+        await lock.stop()         # 停止续约
+        lock.release()            # 受保护释放
+    cache.close()
+
+asyncio.run(main())
+```
+
+- `acquire()`：SET NX，原子地仅在不存在时写入；key 带 TTL 租约，持有方宕机后到期自动让出，不会死锁。
+- `start()` / `stop()`：后台续约任务；续约前校验 token，仅在自己仍持有时 `touch`，锁丢失即停止续约。
+- `release()`：受保护释放——仅当 token 仍匹配（仍是自己持有）才删除，绝不误删其他进程/主机抢到的锁。
+- **后端适用性**：`redis` 真正跨主机共享；`diskcache` 进程安全（SQLite），可用于单机多进程或共享卷场景，不可跨无共享文件系统的主机。
+
 ---

 ## 后端配置（config.toml `[cache]` 段）
@@ -4,7 +4,7 @@ requires = ["hatchling"]

 [project]
 name = "chongming-cache"
-version = "0.1.0"
+version = "0.2.0"
 description = "统一缓存封装：兼容 diskcache / redis 后端"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -1,6 +1,14 @@
 from chongming_cache.backends import CacheBackend, DiskBackend, RedisBackend
 from chongming_cache.cache import Cache
+from chongming_cache.lock import DistributedLock

 __version__ = "0.1.0"

-__all__ = ["Cache", "CacheBackend", "DiskBackend", "RedisBackend", "__version__"]
+__all__ = [
+    "Cache",
+    "CacheBackend",
+    "DiskBackend",
+    "RedisBackend",
+    "DistributedLock",
+    "__version__",
+]
@@ -49,6 +49,14 @@ class CacheBackend(ABC):
    def set(self, key, value, expire: float | None = None) -> None:
        """写入，可带 TTL（秒）。"""

+    @abstractmethod
+    def set_nx(self, key, value, expire: float | None = None) -> bool:
+        """仅当 key 不存在时写入（SET NX 语义），可带 TTL（秒）。
+
+        返回 True 表示本次写入成功（此前不存在）；返回 False 表示 key 已存在，未覆盖。
+        实现必须保持原子性，用于构建分布式排他锁。
+        """
+
    @abstractmethod
    def touch(self, key, expire: float | None = None) -> bool:
        """刷新 key 的 TTL，返回是否命中（未命中返回 False）。"""
@@ -47,6 +47,10 @@ class DiskBackend(CacheBackend):
    def set(self, key, value, expire: float | None = None) -> None:
        self._cache.set(key, value, expire=expire)

+    def set_nx(self, key, value, expire: float | None = None) -> bool:
+        # diskcache.add 仅在 key 不存在时写入，返回是否成功（SET NX 语义）且进程安全
+        return bool(self._cache.add(key, value, expire=expire))
+
    def touch(self, key, expire: float | None = None) -> bool:
        return self._cache.touch(key, expire=expire)

@@ -100,6 +100,17 @@ class RedisBackend(CacheBackend):
    def set(self, key, value, expire: float | None = None) -> None:
        self._redis.set(self._key(key), pickle.dumps(value), ex=expire)

+    def set_nx(self, key, value, expire: float | None = None) -> bool:
+        # redis SET key val NX EX ttl：原子地仅在不存在时写入，用于分布式锁
+        return bool(
+            self._redis.set(
+                self._key(key),
+                pickle.dumps(value),
+                ex=expire,
+                nx=True,
+            )
+        )
+
    def touch(self, key, expire: float | None = None) -> bool:
        # expire=None 时 redis-py 发 PERSIST（永久），与 diskcache.touch 语义一致
        return bool(self._redis.expire(self._key(key), expire))
@@ -0,0 +1,93 @@
+"""基于缓存的分布式排他锁（租约 + 续约）。
+
+利用后端提供的 ``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)