package com.yxcheng.token.common.limit;



import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/**
 * redis 分布式集群限流工具类,采用滑动时间窗口的算法来进行计数统计和限流判断
 * 
 * @author wuyuegang
 * 
 */
@Component
public class RedisLimiter
{
    private final Logger logger = LoggerFactory.getLogger(this.getClass());
    
    private static final String LIMIT_TIME_KEY_SUFFIX = ".time_key";
    
    private static final String LIMIT_COUNTER_KEY_SUFFIX = ".counter_key";
    
    /**
     * 类锁，用于并发控制
     */
    public static Object lock = new Object();
    
    /**
     * 注入redis操作工具类
     */
    @Autowired
    private RedisUtil redisUtil;
    
    /**
     * 根据key值向redis申请资源，如果超过限流值，直接返回false
     * 
     * @param key 限流关键字，通常为 className+"."+methodName
     * @param count 限流数量
     * @param unitTime 限流单位时间
     * @return
     */
    public boolean tryAcquire(String key, int count, long unitTime)
    {
        // 定义时间窗口key,规则：key+常量后缀
        String timeKey = key + LIMIT_TIME_KEY_SUFFIX;
        // 定义计数器key，规则：key+常量后缀
        String counterKey = key + LIMIT_COUNTER_KEY_SUFFIX;
        long start = System.currentTimeMillis();
        try
        {
            if (!redisUtil.exists(timeKey))
            { // 如果不存在时间窗口key值，初始化一个时间窗口，放到redis缓存
                synchronized (lock) // 为了线程安全，加类锁判断，确保多线程环境下，只初始化一次
                {
                    if (!redisUtil.exists(timeKey))
                    {
                        redisUtil.set(timeKey, "0", unitTime);
                        // 同时置空计数器
                        redisUtil.set(counterKey, "0");// count reset to 0
                        if (logger.isDebugEnabled())
                        {
                            logger.debug("RedisLimiter init timeKey[{}] counterKey[{}]", timeKey, counterKey);
                        }
                    }
                }
                
            }
            if (redisUtil.exists(timeKey) && redisUtil.incrby(counterKey, (long)1) > count)
            {
                if (logger.isInfoEnabled())
                {
                    logger.info("RedisLimiter tryAcquire key[{}] exceed count [{}], return false, cost[{}]ms",
                        timeKey,
                        count,
                        (System.currentTimeMillis() - start));
                }
                // 如果在当前时间窗口下，计数器超过允许的限流count，直接返回false
                return false;
            }
        }
        catch (Exception e)
        {// 遇到任何异常，都应放过限流，不影响业务，并记录错误日志
            logger.error("RedisLimiter error when tryAcquire,key[{}] count [{}] cost[{}]ms errormsg [{}]",
                timeKey,
                count,
                (System.currentTimeMillis() - start),
                e.getMessage());
        }
        // 如果没有超过限流值，返回true
        return true;
    }
}
