一個輕量級的基于RateLimiter的分布式限流實作-有解無憂

上篇文章（限流演算法與Guava RateLimiter決議）對常用的限流演算法及Google Guava基于令牌桶演算法的實作RateLimiter進行了介紹，RateLimiter通過執行緒鎖控制同步，只適用于單機應用，在分布式環境下，雖然有像阿里Sentinel的限流開源框架，但對于一些小型應用來說未免過重，但限流的需求在小型專案中也是存在的，比如獲取手機驗證碼的控制，對資源消耗較大操作的訪問頻率控制等，本文介紹最近寫的一個基于RateLimiter，適用于分布式環境下的限流實作，并使用spring-boot-starter的形式發布，比較輕量級且“開箱即用”，

本文限流實作包括兩種形式：

基于RateLimiter令牌桶演算法的限速控制（嚴格限制訪問速度）
基于Lua腳本的限量控制（限制一個時間視窗內的訪問量，對訪問速度沒有嚴格限制）

限速控制

1. 令牌桶模型

首先定義令牌桶模型，與RateLimiter中類似，包括幾個關鍵屬性與關鍵方法，其中關鍵屬性定義如下，

@Data
public class RedisPermits {

    /**
     * 最大存盤令牌數
     */
    private double maxPermits;
    /**
     * 當前存盤令牌數
     */
    private double storedPermits;
    /**
     * 添加令牌的時間間隔/毫秒
     */
    private double intervalMillis;
    /**
     * 下次請求可以獲取令牌的時間，可以是過去（令牌積累）也可以是將來的時間（令牌預消費）
     */
    private long nextFreeTicketMillis;

    //...

關鍵方法定義與RateLimiter也大同小異，方法注釋基本已描述各方法用途，不再贅述，

    /**
     * 構建Redis令牌資料模型
     *
     * @param permitsPerSecond     每秒放入的令牌數
     * @param maxBurstSeconds      maxPermits由此欄位計算，最大存盤maxBurstSeconds秒生成的令牌
     * @param nextFreeTicketMillis 下次請求可以獲取令牌的起始時間，默認當前系統時間
     */
    public RedisPermits(double permitsPerSecond, double maxBurstSeconds, Long nextFreeTicketMillis) {
        this.maxPermits = permitsPerSecond * maxBurstSeconds;
        this.storedPermits = maxPermits;
        this.intervalMillis = TimeUnit.SECONDS.toMillis(1) / permitsPerSecond;
        this.nextFreeTicketMillis = nextFreeTicketMillis;
    }

    /**
     * 基于當前時間，若當前時間晚于nextFreeTicketMicros，則計算該段時間內可以生成多少令牌，將生成的令牌加入令牌桶中并更新資料
     */
    public void resync(long nowMillis) {
        if (nowMillis > nextFreeTicketMillis) {
            double newPermits = (nowMillis - nextFreeTicketMillis) / intervalMillis;
            storedPermits = Math.min(maxPermits, storedPermits + newPermits);
            nextFreeTicketMillis = nowMillis;
        }
    }

    /**
    * 保留指定數量令牌，并回傳需要等待的時間
    */
    public long reserveAndGetWaitLength(long nowMillis, int permits) {
        resync(nowMillis);
        double storedPermitsToSpend = Math.min(permits, storedPermits); // 可以消耗的令牌數
        double freshPermits = permits - storedPermitsToSpend; // 需要等待的令牌數
        long waitMillis = (long) (freshPermits * intervalMillis); // 需要等待的時間

        nextFreeTicketMillis = LongMath.saturatedAdd(nextFreeTicketMillis, waitMillis);
        storedPermits -= storedPermitsToSpend;
        return waitMillis;
    }

    /**
    * 在超時時間內，是否有指定數量的令牌可用
    */
    public boolean canAcquire(long nowMillis, int permits, long timeoutMillis) {
        return queryEarliestAvailable(nowMillis, permits) <= timeoutMillis;
    }

    /**
     * 指定數量令牌數可用需等待的時間
     *
     * @param permits 需保留的令牌數
     * @return 指定數量令牌可用的等待時間，如果為0或負數，表示當前可用
     */
    private long queryEarliestAvailable(long nowMillis, int permits) {
        resync(nowMillis);
        double storedPermitsToSpend = Math.min(permits, storedPermits); // 可以消耗的令牌數
        double freshPermits = permits - storedPermitsToSpend; // 需要等待的令牌數
        long waitMillis = (long) (freshPermits * intervalMillis); // 需要等待的時間

        return LongMath.saturatedAdd(nextFreeTicketMillis - nowMillis, waitMillis);
    }

2. 令牌桶控制類

Guava RateLimiter中的控制都在RateLimiter及其子類中（如SmoothBursty），本處涉及到分布式環境下的同步，因此將其解耦，令牌桶模型存盤于Redis中，對其同步操作的控制放置在如下控制類，其中同步控制使用到了前面介紹的分布式鎖（參考基于Redis分布式鎖的正確打開方式）

@Slf4j
public class RedisRateLimiter {

    /**
     * 獲取一個令牌，阻塞一直到獲取令牌，回傳阻塞等待時間
     *
     * @return time 阻塞等待時間/毫秒
     */
    public long acquire(String key) throws IllegalArgumentException {
        return acquire(key, 1);
    }

    /**
     * 獲取指定數量的令牌，如果令牌數不夠，則一直阻塞，回傳阻塞等待的時間
     *
     * @param permits 需要獲取的令牌數
     * @return time 等待的時間/毫秒
     * @throws IllegalArgumentException tokens值不能為負數或零
     */
    public long acquire(String key, int permits) throws IllegalArgumentException {
        long millisToWait = reserve(key, permits);
        log.info("acquire {} permits for key[{}], waiting for {}ms", permits, key, millisToWait);
        try {
            Thread.sleep(millisToWait);
        } catch (InterruptedException e) {
            log.error("Interrupted when trying to acquire {} permits for key[{}]", permits, key, e);
        }
        return millisToWait;
    }

    /**
     * 在指定時間內獲取一個令牌，如果獲取不到則一直阻塞，直到超時
     *
     * @param timeout 最大等待時間（超時時間），為0則不等待立即回傳
     * @param unit    時間單元
     * @return 獲取到令牌則true，否則false
     * @throws IllegalArgumentException
     */
    public boolean tryAcquire(String key, long timeout, TimeUnit unit) throws IllegalArgumentException {
        return tryAcquire(key, 1, timeout, unit);
    }

    /**
     * 在指定時間內獲取指定數量的令牌，如果在指定時間內獲取不到指定數量的令牌，則直接回傳false，
     * 否則阻塞直到能獲取到指定數量的令牌
     *
     * @param permits 需要獲取的令牌數
     * @param timeout 最大等待時間（超時時間）
     * @param unit    時間單元
     * @return 如果在指定時間內能獲取到指定令牌數，則true,否則false
     * @throws IllegalArgumentException tokens為負數或零，拋出例外
     */
    public boolean tryAcquire(String key, int permits, long timeout, TimeUnit unit) throws IllegalArgumentException {
        long timeoutMillis = Math.max(unit.toMillis(timeout), 0);
        checkPermits(permits);

        long millisToWait;
        boolean locked = false;
        try {
            locked = lock.lock(key + LOCK_KEY_SUFFIX, WebUtil.getRequestId(), 60, 2, TimeUnit.SECONDS);
            if (locked) {
                long nowMillis = getNowMillis();
                RedisPermits permit = getPermits(key, nowMillis);
                if (!permit.canAcquire(nowMillis, permits, timeoutMillis)) {
                    return false;
                } else {
                    millisToWait = permit.reserveAndGetWaitLength(nowMillis, permits);
                    permitsRedisTemplate.opsForValue().set(key, permit, expire, TimeUnit.SECONDS);
                }
            } else {
                return false;  //超時獲取不到鎖，也回傳false
            }
        } finally {
            if (locked) {
                lock.unLock(key + LOCK_KEY_SUFFIX, WebUtil.getRequestId());
            }
        }
        if (millisToWait > 0) {
            try {
                Thread.sleep(millisToWait);
            } catch (InterruptedException e) {

            }
        }
        return true;
    }

    /**
     * 保留指定的令牌數待用
     *
     * @param permits 需保留的令牌數
     * @return time 令牌可用的等待時間
     * @throws IllegalArgumentException tokens不能為負數或零
     */
    private long reserve(String key, int permits) throws IllegalArgumentException {
        checkPermits(permits);
        try {
            lock.lock(key + LOCK_KEY_SUFFIX, WebUtil.getRequestId(), 60, 2, TimeUnit.SECONDS);
            long nowMillis = getNowMillis();
            RedisPermits permit = getPermits(key, nowMillis);
            long waitMillis = permit.reserveAndGetWaitLength(nowMillis, permits);
            permitsRedisTemplate.opsForValue().set(key, permit, expire, TimeUnit.SECONDS);
            return waitMillis;
        } finally {
            lock.unLock(key + LOCK_KEY_SUFFIX, WebUtil.getRequestId());
        }
    }

    /**
     * 獲取令牌桶
     *
     * @return
     */
    private RedisPermits getPermits(String key, long nowMillis) {
        RedisPermits permit = permitsRedisTemplate.opsForValue().get(key);
        if (permit == null) {
            permit = new RedisPermits(permitsPerSecond, maxBurstSeconds, nowMillis);
        }
        return permit;
    }

    /**
     * 獲取redis服務器時間
     */
    private long getNowMillis() {
        String luaScript = "return redis.call('time')";
        DefaultRedisScript<List> redisScript = new DefaultRedisScript<>(luaScript, List.class);
        List<String> now = (List<String>)stringRedisTemplate.execute(redisScript, null);
        return now == null ? System.currentTimeMillis() : Long.valueOf(now.get(0))*1000+Long.valueOf(now.get(1))/1000;
    }

    //...
}

其中：

acquire 是阻塞方法，如果沒有可用的令牌，則一直阻塞直到獲取到令牌，
tryAcquire 則是非阻塞方法，如果在指定超時時間內獲取不到指定數量的令牌，則直接回傳false，不阻塞等待，
getNowMillis 獲取Redis服務器時間，避免業務服務器時間不一致導致的問題，如果業務服務器能保障時間同步，則可從本地獲取提高效率，

3. 令牌桶控制工廠類

工廠類負責管理令牌桶控制類，將其快取在本地，這里使用了Guava中的Cache，一方面避免每次都新建控制類提高效率，另一方面通過控制快取的最大容量來避免像用戶粒度的限流占用過多的記憶體，

public class RedisRateLimiterFactory {

    private PermitsRedisTemplate permitsRedisTemplate;
    private StringRedisTemplate stringRedisTemplate;
    private DistributedLock distributedLock;

    private Cache<String, RedisRateLimiter> cache = CacheBuilder.newBuilder()
            .initialCapacity(100)  //初始大小
            .maximumSize(10000) // 快取的最大容量
            .expireAfterAccess(5, TimeUnit.MINUTES) // 快取在最后一次訪問多久之后失效
            .concurrencyLevel(Runtime.getRuntime().availableProcessors()) // 設定并發級別
            .build();

    public RedisRateLimiterFactory(PermitsRedisTemplate permitsRedisTemplate, StringRedisTemplate stringRedisTemplate, DistributedLock distributedLock) {
        this.permitsRedisTemplate = permitsRedisTemplate;
        this.stringRedisTemplate = stringRedisTemplate;
        this.distributedLock = distributedLock;
    }

    /**
     * 創建RateLimiter
     *
     * @param key              RedisRateLimiter本地快取key
     * @param permitsPerSecond 每秒放入的令牌數
     * @param maxBurstSeconds  最大存盤maxBurstSeconds秒生成的令牌
     * @param expire           該令牌桶的redis tty/秒
     * @return RateLimiter
     */
    public RedisRateLimiter build(String key, double permitsPerSecond, double maxBurstSeconds, int expire) {
        if (cache.getIfPresent(key) == null) {
            synchronized (this) {
                if (cache.getIfPresent(key) == null) {
                    cache.put(key, new RedisRateLimiter(permitsRedisTemplate, stringRedisTemplate, distributedLock, permitsPerSecond,
                            maxBurstSeconds, expire));
                }
            }
        }
        return cache.getIfPresent(key);
    }
}

4. 注解支持

定義注解 @RateLimit 如下，表示以每秒rate的速率放置令牌，最多保留burst秒的令牌，取令牌的超時時間為timeout，limitType用于控制key型別，目前支持：

IP, 根據客戶端IP限流
USER, 根據用戶限流，對于Spring Security可從SecurityContextHolder中獲取當前用戶資訊，如userId
METHOD, 根據方法名全域限流，className.methodName，注意避免同時對同一個類中的同名方法做限流控制，否則需要修改獲取key的邏輯
CUSTOM，自定義，支持運算式決議，如#{id}, #{user.id}

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimit {
    String key() default "";
    String prefix() default "rateLimit:"; //key前綴
    int expire() default 60; // 表示令牌桶模型RedisPermits redis key的過期時間/秒
    double rate() default 1.0; // permitsPerSecond值
    double burst() default 1.0; // maxBurstSeconds值
    int timeout() default 0; // 超時時間/秒
    LimitType limitType() default LimitType.METHOD;
}

通過切面的前置增強來為添加了 @RateLimit 注解的方法提供限流控制，如下

@Aspect
@Slf4j
public class RedisLimitAspect {
    //...

    @Before(value = "https://www.cnblogs.com/spec-dog/p/@annotation(rateLimit)")
    public void rateLimit(JoinPoint  point, RateLimit rateLimit) throws Throwable {
        String key = getKey(point, rateLimit.limitType(), rateLimit.key(), rateLimit.prefix());
        RedisRateLimiter redisRateLimiter = redisRateLimiterFactory.build(key, rateLimit.rate(), rateLimit.burst(), rateLimit.expire());
        if(!redisRateLimiter.tryAcquire(key, rateLimit.timeout(), TimeUnit.SECONDS)){
            ExceptionUtil.rethrowClientSideException(LIMIT_MESSAGE);
        }
    }

    //...

限量控制

1. 限量控制類

限制一個時間視窗內的訪問量，可使用計數器演算法，借助Lua腳本執行的原子性來實作，

Lua腳本邏輯：

以需要控制的物件為key（如方法，用戶ID，或IP等），當前訪問次數為Value，時間視窗值為快取的過期時間
如果key存在則將其增1，判斷當前值是否大于訪問量限制值，如果大于則回傳0，表示該時間視窗內已達訪問量上限，如果小于則回傳1表示允許訪問
如果key不存在，則將其初始化為1，并設定過期時間，回傳1表示允許訪問

public class RedisCountLimiter {

    private StringRedisTemplate stringRedisTemplate;

    private static final String LUA_SCRIPT = "local c \nc = redis.call('get',KEYS[1]) \nif c and redis.call('incr',KEYS[1]) > tonumber(ARGV[1]) then return 0 end"
            + " \nif c then return 1 else \nredis.call('set', KEYS[1], 1) \nredis.call('expire', KEYS[1], tonumber(ARGV[2])) \nreturn 1 end";

    private static final int SUCCESS_RESULT = 1;
    private static final int FAIL_RESULT = 0;

    public RedisCountLimiter(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    /**
     * 是否允許訪問
     *
     * @param key redis key
     * @param limit 限制次數
     * @param expire 時間段/秒
     * @return 獲取成功true，否則false
     * @throws IllegalArgumentException
     */
    public boolean tryAcquire(String key, int limit, int expire) throws IllegalArgumentException {
        RedisScript<Number> redisScript = new DefaultRedisScript<>(LUA_SCRIPT, Number.class);
        Number result = stringRedisTemplate.execute(redisScript, Collections.singletonList(key), String.valueOf(limit), String.valueOf(expire));
        if(result != null && result.intValue() == SUCCESS_RESULT) {
            return true;
        }
        return false;
    }

}

2. 注解支持

定義注解 @CountLimit 如下，表示在period時間視窗內，最多允許訪問limit次，limitType用于控制key型別，取值與 @RateLimit 同，

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface CountLimit {
    String key() default "";
    String prefix() default "countLimit:"; //key前綴
    int limit() default 1;  // expire時間段內限制訪問次數
    int period() default 1; // 表示時間段/秒
    LimitType limitType() default LimitType.METHOD;
}

同樣采用前值增強來為添加了 @CountLimit 注解的方法提供限流控制，如下

@Before(value = "https://www.cnblogs.com/spec-dog/p/@annotation(countLimit)")
public void countLimit(JoinPoint  point, CountLimit countLimit) throws Throwable {
    String key = getKey(point, countLimit.limitType(), countLimit.key(), countLimit.prefix());
    if (!redisCountLimiter.tryAcquire(key, countLimit.limit(), countLimit.period())) {
        ExceptionUtil.rethrowClientSideException(LIMIT_MESSAGE);
    }
}

使用示例

1.添加依賴

<dependencies>
    <dependency>
        <groupId>cn.jboost.springboot</groupId>
        <artifactId>limiter-spring-boot-starter</artifactId>
        <version>1.3-SNAPSHOT</version>
    </dependency>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
    </dependency>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-redis</artifactId>
    </dependency>
</dependencies>

2.配置redis相關引數

spring:
  application:
    name: limiter-demo
  redis:
    #資料庫索引
    database: 0
    host: 192.168.40.92
    port: 6379
    password: password
    #連接超時時間
    timeout: 2000

3.測驗類

@RestController
@RequestMapping("limiter")
public class LimiterController {

    /**
     * 注解形式
     * @param key
     * @return
     */
    @GetMapping("/count")
    @CountLimit(key = "#{key}", limit = 2, period = 10, limitType = LimitType.CUSTOM)
    public String testCountLimit(@RequestParam("key") String key){
        return "test count limiter...";
    }

    /**
     * 注解形式
     * @param key
     * @return
     */
    @GetMapping("/rate")
    @RateLimit(rate = 1.0/5, burst = 5.0, expire = 120, timeout = 0)
    public String testRateLimit(@RequestParam("key") String key){
        return "test rate limiter...";
    }

    @Autowired
    private RedisRateLimiterFactory redisRateLimiterFactory;
    /**
     * 代碼段形式
     * @param
     * @return
     */
    @GetMapping("/rate2")
    public String testRateLimit(){
        RedisRateLimiter limiter = redisRateLimiterFactory.build("LimiterController.testRateLimit", 1.0/30, 30, 120);
        if(!limiter.tryAcquire("app.limiter", 0, TimeUnit.SECONDS)) {
            System.out.println(LocalDateTime.now());
            ExceptionUtil.rethrowClientSideException("您的訪問過于頻繁，請稍后重試");
        }
        return "test rate limiter 2...";
    }
}

4.驗證

啟動測驗專案，瀏覽器中訪問 http://localhost:8080/limiter/rate?key=test ，第一次訪問成功，如圖

ratelimiter1

持續重繪，將回傳如下錯誤，直到5s之后再回傳成功，限制5秒1次的訪問速度

ratelimiter2

注解的使用

限流型別LimitType支持IP（客戶端IP）、用戶（userId）、方法（className.methodName）、自定義（CUSTOM）幾種形式，默認為METHOD
LimitType為CUSTOM時，需要手動指定key（其它key自動為ip，userid，或methodname），key支持運算式形式，如#{id}, #{user.id}
針對某個時間視窗內限制訪問一次的場景，既可以使用 @CountLimit，也可以使用 @RateLimit，比如驗證碼一分鐘內只允許獲取一次，以下兩種形式都能達到目的

//同一個手機號碼60s內最多訪問一次
@CountLimit(key = "#{params.phone}", limit = 1, period = 60, limitType = LimitType.CUSTOM)
//以1/60的速度放置令牌，最多保存60s的令牌（也就是最多保存一個），控制訪問速度為1/60個每秒（1個每分鐘）
@RateLimit(key = "#{params.phone}", rate = 1.0/60, burst = 60, expire = 120, limitType = LimitType.CUSTOM)

總結

本文介紹了適用于分布式環境的基于RateLimiter令牌桶演算法的限速控制與基于計數器演算法的限量控制，可應用于中小型專案中有相關需求的場景（注：本實作未做壓力測驗，如果用戶并發量較大需驗證效果），

本文完整代碼見：https://github.com/ronwxy/base-spring-boot ，目錄 spring-boot-autoconfigure/src/main/java/cn/jboost/springboot/autoconfig/limiter
示例專案代碼地址：https://github.com/ronwxy/springboot-demos/tree/master/springboot-limiter

如果覺得有幫助，別忘了給個star ^_，作者公眾號：半路雨歌，歡迎關注查看更多干貨文章，

[轉載請注明出處]
作者：雨歌
歡迎關注作者公眾號：半路雨歌，查看更多技術干貨文章
qrcode

轉載請註明出處，本文鏈接：https://www.uj5u.com/houduan/117751.html

標籤：Java

上一篇：【JAVA SE基礎篇】65.UDP TCP

下一篇：樂觀鎖的簡單使用通俗易懂超詳細