Skip to content

Instantly share code, notes, and snippets.

@codeachange

codeachange/RateLimiter.java

Created December 29, 2017 06:58
Download ZIP

Revisions

  1. codeachange created this gist Dec 29, 2017.
    115 changes: 115 additions & 0 deletions RateLimiter.java
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,115 @@
    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);
    }
    }