上两个GPT写的锁,一个是文件锁,一个是Redis锁,写的那是相当的完美

#!/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

 

posted @ 2025-11-28 19:48  就是想学习  阅读(0)  评论(0)    收藏  举报