AiOpsWechatiPad/common/redis_helper.py

import redis
import os
from config import conf
import uuid
import time
import threading
# 定义全局 redis_helper
redis_helper = None

class RedisHelper:
    def __init__(self, host='localhost', port=6379, password=None ,db=0):
        # 初始化 Redis 连接
        self.client = redis.Redis(host=host, port=port,db=db,password=password)
        self.lock_renewal_thread = None

    def set_hash(self, hash_key, data, timeout=None):
        """添加或更新哈希，并设置有效期"""
        self.client.hset(hash_key, mapping=data)
        if timeout:
            # 设置有效期（单位：秒）
            self.client.expire(hash_key, timeout)

    def get_hash(self, hash_key):
        """获取整个哈希表数据"""
        result = self.client.hgetall(hash_key)
        # 将字节数据解码成字符串格式返回
        return {k.decode('utf-8'): v.decode('utf-8') for k, v in result.items()}

    def get_hash_field(self, hash_key, field):
        """获取哈希表中的单个字段值"""
        result = self.client.hget(hash_key, field)
        return result.decode('utf-8') if result else None

    def delete_hash(self, hash_key):
        """删除整个哈希表"""
        self.client.delete(hash_key)
    
    def delete_hash_field(self, hash_key, field):
        """删除哈希表中的某个字段"""
        self.client.hdel(hash_key, field)

    def update_hash_field(self, hash_key, field, value):
        """更新哈希表中的某个字段"""
        self.client.hset(hash_key, field, value)

    def acquire_lock(self, lock_name, timeout=60):
        """
        尝试获取分布式锁，成功返回 True，失败返回 False
        :param lock_name: 锁的名称
        :param timeout: 锁的超时时间（秒）
        :return: bool
        """
        #print('获取锁')
        identifier = str(time.time())  # 使用时间戳作为唯一标识
        #if self.client.set(lock_name, identifier, nx=True, ex=timeout):
        if self.client.setnx(lock_name, identifier):
            self.client.expire(lock_name, timeout)
            self.lock_renewal_thread = threading.Thread(target=self.renew_lock, args=(lock_name, identifier, timeout))
            self.lock_renewal_thread.start()
            return True
        return False

    def renew_lock(self, lock_name, identifier, timeout):
        """
        锁的自动续期
        :param lock_name: 锁的名称
        :param identifier: 锁的唯一标识
        :param timeout: 锁的超时时间（秒）
        """
        while True:
            time.sleep(timeout / 2)
            if self.client.get(lock_name) == identifier.encode():
                self.client.expire(lock_name, timeout)
            else:
                break

    def release_lock(self, lock_name, identifier):
        """
        释放分布式锁
        :param lock_name: 锁的名称
        :param identifier: 锁的唯一标识
        """
        if self.client.get(lock_name) == identifier.encode():
            self.client.delete(lock_name)
            if self.lock_renewal_thread:
                self.lock_renewal_thread.join()

def start():
    global redis_helper
    
    host=conf().get("redis_host")
    port=conf().get("redis_port")
    password=conf().get("redis_password")
    db=conf().get("redis_db")
    redis_helper = RedisHelper(host=host,port=port,password=password,db=db)