LQAgentPlatform/foundation/infrastructure/cache/redis_lock.py

# !/usr/bin/python
# -*- encoding: utf-8 -*-
"""
@Time    :   2025/07/30 14:40
@Author  :    
@File    :   RedisLock.py
@Software:   VScode
@Desc    :   None
"""


import time
import uuid

class RedisLock:
    """
    Redis 锁类
    """
    
    def __init__(self, redis_client, lock_name, expire_time=30):
        """
        :param redis_client: Redis 客户端连接
        :param lock_name: 锁的名称
        :param expire_time: 锁的过期时间(秒)
        """
        self.redis = redis_client
        self.lock_name = lock_name
        self.expire_time = expire_time
        self.identifier = str(uuid.uuid4())  # 唯一标识，用于安全释放锁

    def acquire(self, timeout=10):
        """
        获取锁
        :param timeout: 获取锁的超时时间(秒)
        :return: 是否成功获取锁
        """
        end = time.time() + timeout
        while time.time() < end:
            
            # 尝试获取锁
            if self.redis.set(self.lock_name, self.identifier, nx=True, ex=self.expire_time):
                return True
            time.sleep(0.001)  # 短暂等待后重试
        return False

    def release(self):
        """
        释放锁
        """
        # 使用 Lua 脚本保证原子性
        unlock_script = """
        if redis.call("get", KEYS[1]) == ARGV[1] then
            return redis.call("del", KEYS[1])
        else
            return 0
        end
        """
        self.redis.eval(unlock_script, 1, self.lock_name, self.identifier)


    def __enter__(self):
        if not self.acquire():
            raise Exception("Could not acquire lock")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.release()