#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简易文件锁(多进程通用)
- 基于 fcntl.flock(POSIX 系统:Linux / macOS)
- 支持超时等待、轮询间隔
- 支持 with 上下文管理
"""
import os
import time
import errno
import fcntl
from typing import Optional
class FileLockTimeoutError(TimeoutError):
"""获取文件锁超时异常"""
pass
class FileLock:
"""
文件锁(进程间互斥)
使用方式:
from file_lock import FileLock
lock = FileLock("/tmp/order_commander.lock", timeout=10)
with lock:
# 在此代码块中,多进程互斥
do_something()
锁语义:
- 同一时刻只有一个进程能持有锁
- 进程退出或文件描述符关闭时,锁自动释放
- 其他进程会阻塞等待(或在 timeout 后抛异常)
"""
def __init__(
self,
lock_file: str = "/tmp/schwab_order_commander.lock",
timeout: Optional[float] = None,
poll_interval: float = 0.1,
):
"""
:param lock_file: 锁文件路径(建议放 /tmp 或你的项目 runtime 目录)
:param timeout: 获取锁的最大等待时间(秒)。
- None 表示一直等
- 0 表示立即返回,获取不到则抛 FileLockTimeoutError
:param poll_interval: 轮询间隔(秒),仅当 timeout 非 None 且 > 0 时生效
"""
self.lock_file = os.path.abspath(lock_file)
self.timeout = timeout
self.poll_interval = poll_interval
self._fd: Optional[int] = None
self._is_locked: bool = False
def acquire(self, blocking: bool = True) -> None:
"""
获取锁。
:param blocking: 是否阻塞等待(一般保持 True 即可)
:raises FileLockTimeoutError: 在 timeout 内仍未获取到锁
:raises OSError: 其他底层系统错误
"""
if self._is_locked:
# 已经持有锁,直接返回
return
# 确保目录存在
lock_dir = os.path.dirname(self.lock_file)
if lock_dir and not os.path.isdir(lock_dir):
os.makedirs(lock_dir, exist_ok=True)
# 以 append 模式打开,避免截断文件
fd = os.open(self.lock_file, os.O_RDWR | os.O_CREAT, 0o644)
self._fd = fd
start_time = time.time()
if not blocking and self.timeout is None:
# 非阻塞且没有 timeout 逻辑,直接尝试一次
self._try_lock_once(non_blocking=True)
self._is_locked = True
self._write_owner_info()
return
# 带 timeout 的阻塞模式
while True:
try:
# 非阻塞方式申请锁;失败会抛 BlockingIOError
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
self._is_locked = True
self._write_owner_info()
return
except OSError as e:
if e.errno not in (errno.EAGAIN, errno.EACCES):
# 其他错误直接抛出
raise
# 锁被其他进程占用
if self.timeout is not None:
elapsed = time.time() - start_time
if elapsed >= self.timeout:
# 超时,释放 fd
os.close(fd)
self._fd = None
raise FileLockTimeoutError(
f"获取文件锁超时: {self.lock_file}, timeout={self.timeout}s"
)
# 等待后重试
time.sleep(self.poll_interval)
def _try_lock_once(self, non_blocking: bool) -> None:
"""内部单次加锁(用于非阻塞模式)"""
if self._fd is None:
raise RuntimeError("文件描述符未打开")
flags = fcntl.LOCK_EX
if non_blocking:
flags |= fcntl.LOCK_NB
try:
fcntl.flock(self._fd, flags)
except OSError as e:
if e.errno in (errno.EAGAIN, errno.EACCES):
raise FileLockTimeoutError(f"锁已被占用: {self.lock_file}") from e
raise
def _write_owner_info(self) -> None:
"""
写入简单的 owner 信息(可选功能):
- PID
- 获取时间
方便调试和排查问题。
"""
if self._fd is None:
return
try:
# 清空原内容再写
os.lseek(self._fd, 0, os.SEEK_SET)
os.ftruncate(self._fd, 0)
info = f"pid={os.getpid()} time={time.strftime('%Y-%m-%d %H:%M:%S')}\n"
os.write(self._fd, info.encode("utf-8"))
os.fsync(self._fd)
except OSError:
# 写失败无所谓,锁本身仍然有效
pass
def release(self) -> None:
"""
释放锁。
- 正常情况下会显式调用;
- 进程退出时,操作系统也会自动释放锁(fd 关闭)。
"""
if not self._is_locked:
return
if self._fd is not None:
try:
fcntl.flock(self._fd, fcntl.LOCK_UN)
finally:
os.close(self._fd)
self._fd = None
self._is_locked = False
def __enter__(self):
self.acquire()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.release()
@property
def locked(self) -> bool:
"""当前进程是否持有锁"""
return self._is_locked
class MultiFileLock:
def __init__(self, lock_files: list[str], timeout=10):
self.locks = [FileLock(path, timeout=timeout) for path in sorted(lock_files)]
self.acquired = False
def acquire(self):
try:
for lock in self.locks:
lock.acquire()
self.acquired = True
except Exception as e:
# 失败时释放已经获取的锁,避免半锁状态
self.release()
raise e
def release(self):
for lock in reversed(self.locks):
try:
lock.release()
except Exception:
pass
self.acquired = False
def __enter__(self):
self.acquire()
return self
def __exit__(self, exc_type, exc, tb):
self.release()
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import time
import uuid
import threading
from typing import Optional, List
from redis import Redis
class RedisLockTimeoutError(TimeoutError):
pass
class RedisDistributedLock:
"""
Redis 分布式锁(可跨机器 / 多进程 / 多服务)
- 安全释放:仅持锁者才能释放(通过随机 token 检查)
- 支持超时等待
- 支持自动续期 watchdog
"""
LUA_RELEASE_SCRIPT = """
if redis.call("GET", KEYS[1]) == ARGV[1] then
return redis.call("DEL", KEYS[1])
else
return 0
end
"""
def __init__(
self,
redis_client: Redis,
lock_key: str,
ttl: int = 1000 * 10, # 10 秒(毫秒)
timeout: Optional[float] = None, # 等锁最大等待时间(秒)
renew_interval: float = 2.0, # 每 2 秒续期一次
poll_interval: float = 0.1, # 抢锁失败时每 0.1 秒重试
):
self.redis = redis_client
self.lock_key = lock_key
self.ttl = ttl
self.timeout = timeout
self.poll_interval = poll_interval
self.renew_interval = renew_interval
self.token = str(uuid.uuid4())
self._is_locked = False
self._stop_event = threading.Event()
self._renew_thread: Optional[threading.Thread] = None
# ---- 获取锁 ----
def acquire(self):
# 每次重新开始加锁前重置事件与线程句柄,保证 watchdog 正常工作
self._stop_event = threading.Event()
self._renew_thread = None
start_time = time.time()
while True:
res = self.redis.set(
self.lock_key, self.token, nx=True, px=self.ttl
)
if res:
self._is_locked = True
self._start_watchdog()
return
if self.timeout is not None:
elapsed = time.time() - start_time
if elapsed >= self.timeout:
raise RedisLockTimeoutError(
f"获取分布式锁超时: {self.lock_key}"
)
time.sleep(self.poll_interval)
# ---- 释放锁 ----
def release(self):
if not self._is_locked:
return
self._stop_event.set()
if self._renew_thread:
self._renew_thread.join(timeout=1)
script = self.redis.register_script(self.LUA_RELEASE_SCRIPT)
script(keys=[self.lock_key], args=[self.token])
self._is_locked = False
# ---- 续期 watchdog ----
def _watchdog(self):
while not self._stop_event.wait(self.renew_interval):
if self._is_locked:
# 只有持锁者才能续期
current = self.redis.get(self.lock_key)
if current and current.decode("utf-8") == self.token:
self.redis.pexpire(self.lock_key, self.ttl)
def _start_watchdog(self):
self._renew_thread = threading.Thread(
target=self._watchdog, daemon=True
)
self._renew_thread.start()
# ---- 上下文管理器支持 ----
def __enter__(self):
self.acquire()
return self
def __exit__(self, exc_type, exc, tb):
self.release()
@property
def locked(self) -> bool:
return self._is_locked
class MultiRedisDistributedLock:
"""
多锁组合(分布式锁版)
- 支持跨机器 / 多进程 / 多服务一致互斥
- 必须全部加锁成功才允许执行
- 防止死锁(锁按 key 排序)
"""
def __init__(
self,
redis_client: Redis,
lock_keys: List[str],
ttl: int = 1000 * 10, # 毫秒
timeout: float = None, # 等所有锁的最大时间
renew_interval: float = 2.0, # 续期周期
poll_interval: float = 0.1, # 重试间隔
):
if not lock_keys:
raise ValueError("必须至少传入一个 lock_key")
# 按 key 字典序排序,避免死锁
lock_keys_sorted = sorted(lock_keys)
self.locks = [
RedisDistributedLock(
redis_client=redis_client,
lock_key=key,
ttl=ttl,
timeout=timeout,
renew_interval=renew_interval,
poll_interval=poll_interval,
)
for key in lock_keys_sorted
]
self.acquired = False
# ---- 获取所有锁 ----
def acquire(self):
try:
for lock in self.locks:
lock.acquire()
self.acquired = True
except Exception as e:
# 回滚:释放已获取的锁
self.release()
if isinstance(e, RedisLockTimeoutError):
raise
raise RedisLockTimeoutError(f"Multi-lock 获取失败: {e}")
# ---- 释放所有锁(逆序,最安全)----
def release(self):
for lock in reversed(self.locks):
try:
lock.release()
except Exception:
pass
self.acquired = False
# ---- 支持 with ----
def __enter__(self):
self.acquire()
return self
def __exit__(self, exc_type, exc, tb):
self.release()
@property
def locked(self):
return self.acquired