redis计数器,用于控制流量

import time
from datetime import datetime, timedelta
import aioredis
from typing import List, Optional, Dict, Tuple

class UserAccessRecord:
    def __init__(self):
        self.redis = aioredis.Redis(
            host='localhost',
            port=6379,
            db=0,
            decode_responses=True
        )
        self.KEY_PREFIX = "user_access:"
        self.COUNTER_PREFIX = "user_access_counter:"  # 存储访问次数的ZSET
        self.FIVE_MINUTES = 5 * 60  # 5分钟的秒数

    async def record_access(self, key: str):
        """记录key的访问,包括访问时间戳和更新计数"""
        try:
            current_time = time.time()
            # 记录访问时间戳
            time_key = f"{self.KEY_PREFIX}{key}"
            await self.redis.zadd(time_key, {str(current_time): current_time})
            
            # 清理5分钟前的数据
            five_mins_ago = current_time - self.FIVE_MINUTES
            await self.redis.zremrangebyscore(time_key, 0, five_mins_ago)
            
            # 获取5分钟内的访问次数
            count = await self.redis.zcard(time_key)
            
            # 更新计数器ZSET
            counter_key = f"{self.COUNTER_PREFIX}all"
            await self.redis.zadd(counter_key, {key: count}, nx=False, xx=True, ch=True)
            
            # 清理计数为0的记录
            if count == 0:
                await self.redis.zrem(counter_key, key)
                
            # 设置过期时间
            await self.redis.expire(time_key, self.FIVE_MINUTES)
            await self.redis.expire(counter_key, self.FIVE_MINUTES)
            
        except Exception as e:
            print(f"Record access error: {e}")

    async def get_most_accessed_keys(self, limit: int = 10) -> List[Tuple[str, int]]:
        """获取访问次数最多的key"""
        try:
            counter_key = f"{self.COUNTER_PREFIX}all"
            # 获取访问次数最多的key
            result = await self.redis.zrevrange(
                counter_key,
                0,
                limit - 1,
                withscores=True
            )
            return [(key, int(score)) for key, score in result]
            
        except Exception as e:
            print(f"Get most accessed keys error: {e}")
            return []

    async def refresh_access_counts(self):
        """刷新所有key的访问次数"""
        try:
            current_time = time.time()
            five_mins_ago = current_time - self.FIVE_MINUTES
            
            # 获取所有key
            pattern = f"{self.KEY_PREFIX}*"
            all_keys = await self.redis.keys(pattern)
            
            counter_key = f"{self.COUNTER_PREFIX}all"
            pipeline = self.redis.pipeline()
            
            for time_key in all_keys:
                # 清理旧数据
                pipeline.zremrangebyscore(time_key, 0, five_mins_ago)
                # 获取当前计数
                pipeline.zcard(time_key)
                
            results = await pipeline.execute()
            
            # 更新计数器ZSET
            counts = {}
            for i in range(0, len(results), 2):
                key = all_keys[i//2].replace(self.KEY_PREFIX, "")
                count = results[i+1]
                if count > 0:
                    counts[key] = count
            
            if counts:
                await self.redis.zadd(counter_key, counts)
            
            # 清理计数为0的记录
            keys_to_remove = [key for key, count in counts.items() if count == 0]
            if keys_to_remove:
                await self.redis.zrem(counter_key, *keys_to_remove)
                
        except Exception as e:
            print(f"Refresh access counts error: {e}")
posted @ 2025-06-23 10:54  wsl-hitsz  阅读(22)  评论(0)    收藏  举报