RateLimiter频率限制器令牌桶算法

RateLimiter.java

package com.base.framework.utils;

import com.base.framework.exception.AppException;
import org.springframework.data.redis.core.RedisTemplate;

import java.util.concurrent.TimeUnit;

/**
 * 频率限制器
 * 令牌桶算法
 */
public class RateLimiterUtil {

    /**
     * 判断是否超出频率限制
     * @return true 表示超出限制了，false 表示没有超出限制
     */
    public static boolean isExceedRateLimit(RateLimit rateLimit, RedisTemplate<String, Object> redisTemplate, String redisKey) {
        RateLimiterUtil.Allowance allowance = (RateLimiterUtil.Allowance) redisTemplate.opsForValue().get(redisKey);
        if (null == allowance) {
            allowance = new RateLimiterUtil.Allowance(0, RateLimiterUtil.currentTimeSeconds());
        }

        RateLimiterUtil.Allowance newAllowance;
        boolean result;
        try {
            newAllowance = RateLimiterUtil.checkRateLimit(rateLimit, allowance);
            result = false;
        } catch (RateLimiterUtil.RateLimitExceedException e) {
            newAllowance = new RateLimiterUtil.Allowance(0, RateLimiterUtil.currentTimeSeconds());
            result = true;
        }
        redisTemplate.opsForValue().set(redisKey, newAllowance, 5, TimeUnit.HOURS);
        return result;
    }

    private static Allowance checkRateLimit(RateLimit rateLimit, Allowance allowance) throws RateLimitExceedException {
        Integer currentTimestamp = currentTimeSeconds();
        Integer newAllowance = allowance.getAllowance() + (currentTimestamp - allowance.getTimestamp()) * rateLimit.getLimit() / rateLimit.getSeconds();
        if (newAllowance > rateLimit.getLimit()) {
            newAllowance = rateLimit.getLimit();
        }
        if (newAllowance < 1) {
            throw new RateLimitExceedException();
        }
        return new Allowance(newAllowance - 1, currentTimestamp);
    }

    // 余量类，保存余量和检查时间
    public static class Allowance {
        private Integer allowance;
        private Integer timestamp;

        public Allowance() {
        }

        public Allowance(Integer allowance, Integer timestamp) {
            this.allowance = allowance;
            this.timestamp = timestamp;
        }

        public Integer getAllowance() {
            return allowance;
        }

        public void setAllowance(Integer allowance) {
            this.allowance = allowance;
        }

        public Integer getTimestamp() {
            return timestamp;
        }

        public void setTimestamp(Integer timestamp) {
            this.timestamp = timestamp;
        }
    }

    // 频率限制量， e.g. {limit=10, seconds=5} 表示5秒内限制10次
    public static class RateLimit {
        private Integer limit;
        private Integer seconds;

        public RateLimit(Integer limit, Integer seconds) {
            this.limit = limit;
            this.seconds = seconds;
        }

        public Integer getLimit() {
            return limit;
        }

        public void setLimit(Integer limit) {
            this.limit = limit;
        }

        public Integer getSeconds() {
            return seconds;
        }

        public void setSeconds(Integer seconds) {
            this.seconds = seconds;
        }
    }

    public static class RateLimitExceedException extends AppException {
        public RateLimitExceedException() {
            super("100107", "频率超出限制");
        }
    }

    public static Integer currentTimeSeconds() {
        return (int) (System.currentTimeMillis() / 1000);
    }
}