一個輕量級的基於RateLimiter的分散式限流實現

【雨歌】發表於2020-07-31

上篇文章（限流演算法與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 = "@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 = "@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

【分散式架構】(10)---基於Redis元件的特性，實現一個分散式限流
2019-06-05
分散式架構Redis元件
基於 Redis 實現分散式應用限流
2017-09-07
Redis分散式
一個輕量級RPC的實現
2018-01-26
RPC
輕量級分散式鎖的設計原理分析與實現
2020-04-05
分散式
基於kubernetes的分散式限流
2022-04-09
分散式
基於Redis實現一個分散式鎖
2021-11-14
Redis分散式
PHP實現一個輕量級容器
2019-05-12
PHP
手把手教你實現一個基於Redis的分散式鎖
2018-09-08
Redis分散式
一個輕量級的分散式定時任務排程平臺-Cloudtask
2018-03-23
分散式Cloud
詳解Redisson分散式限流的實現原理
2023-02-12
Redis分散式
基於 Zookeeper 的分散式鎖實現
2022-12-05
分散式
基於redis的分散式鎖實現
2018-01-06
Redis分散式
ActorLite：一個輕量級Actor模型實現（中）
2009-05-14
模型
使用 Redis 實現分散式系統輕量級協調技術
2015-08-11
Redis分散式
教你一招：基於Redis實現一個分散式鎖
2022-09-11
Redis分散式
Guava RateLimiter限流
2019-02-26
GuavaMIT
分散式鎖與實現(一)基於Redis實現！
2019-03-22
分散式Redis
Spring的輕量級實現
2021-11-13
Spring
基於ObjC的輕量級Promise庫——RWPromiseKit
2016-03-21
OBJPromise
tiny4j:一個輕量級的類似Spring的實現
2016-12-04
Spring
基於gogs和drone的一個輕量級的開發部署自動化流程
2022-08-14
Go
基於Redis的分散式鎖的簡單實現
2018-01-27
Redis分散式
輕量級分散式檔案系統fastDFS
2016-12-11
分散式AST
基於 Redis 實現簡單的分散式鎖
2016-05-21
Redis分散式
實現一個redis的分散式鎖
2021-12-22
Redis分散式
基於redis實現分散式鎖
2018-11-19
Redis分散式
基於ZK實現分散式鎖
2018-05-26
分散式
理論+演算法+實戰，教你如何實現億級流量下的分散式限流
2022-02-17
演算法分散式
開源一個基於dotnet standard的輕量級的ORM框架-Light.Data
2019-05-22
ORM框架
Dapr實現一個簡單的基於.net分散式事務之Saga模式
2022-02-23
分散式模式
Uvicorn 初體驗-一個基於 asyncio 開發的一個輕量級高效的 Web 伺服器框架
2018-02-26
Web伺服器框架
[分散式限流] 滑動視窗演算法的 Golang 實現
2019-11-27
分散式演算法Golang
MySQL 中基於 XA 實現的分散式事務
2020-04-05
MySql分散式
基於快取或zookeeper的分散式鎖實現
2019-07-14
快取分散式
java 實現開箱即用基於 redis 的分散式鎖
2022-12-14
JavaRedis分散式
基於redis和zookeeper的分散式鎖實現方式
2018-06-15
Redis分散式
Fastflow——基於golang的輕量級工作流框架
2022-05-23
ASTGolang框架
基於RocketMQ實現分散式事務
2024-03-12
MQ分散式