接口限流——令牌桶算法

接口限流——令牌桶算法

簡介：

在網絡中傳輸數據時，爲了防止網絡擁塞，需限制流出網絡的流量，使流量以比較均勻的速度向外發送。令牌桶算法就實現了這個功能，可控制發送到網絡上數據的數目，並允許突發數據的發送。

令牌桶算法是網絡流量整形（Traffic Shaping）和速率限制（Rate Limiting）中最常使用的一種算法。典型情況下，令牌桶算法用來控制發送到網絡上的數據的數目，並允許突發數據的發送。

大小固定的令牌桶可自行以恆定的速率源源不斷地產生令牌。如果令牌不被消耗，或者被消耗的速度小於產生的速度，令牌就會不斷地增多，直到把桶填滿。後面再產生的令牌就會從桶中溢出。最後桶中可以保存的最大令牌數永遠不會超過桶的大小。

傳送到令牌桶的數據包需要消耗令牌。不同大小的數據包，消耗的令牌數量不一樣。

令牌桶這種控制機制基於令牌桶中是否存在令牌來指示什麼時候可以發送流量。令牌桶中的每一個令牌都代表一個字節。如果令牌桶中存在令牌，則允許發送流量；而如果令牌桶中不存在令牌，則不允許發送流量。因此，如果突發門限被合理地配置並且令牌桶中有足夠的令牌，那麼流量就可以以峯值速率發送。

令牌桶算法的基本過程如下：

假如用戶配置的平均發送速率爲r，則每隔1/r秒一個令牌被加入到桶中；

假設桶最多可以存發b個令牌。如果令牌到達時令牌桶已經滿了，那麼這個令牌會被丟棄；

當一個n個字節的數據包到達時，就從令牌桶中刪除n個令牌，並且數據包被髮送到網絡；

如果令牌桶中少於n個令牌，那麼不會刪除令牌，並且認爲這個數據包在流量限制之外；

算法允許最長b個字節的突發，但從長期運行結果看，數據包的速率被限制成常量r。對於在流量限制外的數據包可以以不同的方式處理：

它們可以被丟棄；

它們可以排放在隊列中以便當令牌桶中累積了足夠多的令牌時再傳輸；

它們可以繼續發送，但需要做特殊標記，網絡過載的時候將這些特殊標記的包丟棄。

注意：令牌桶算法不能與另外一種常見算法“漏桶算法（Leaky Bucket）”相混淆。這兩種算法的主要區別在於“漏桶算法”能夠強行限制數據的傳輸速率，而“令牌桶算法”在能夠限制數據的平均傳輸速率外，還允許某種程度的突發傳輸。在“令牌桶算法”中，只要令牌桶中存在令牌，那麼就允許突發地傳輸數據直到達到用戶配置的門限，因此它適合於具有突發特性的流量。

示意圖

spring boot 結合 redis 結合 lua 實現

pom 引入相關依賴

<properties>
    <java.version>1.8</java.version>
    <spring-cloud.version>Greenwich.SR2</spring-cloud.version>
  </properties>

  <dependencies>
    <dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-web</artifactId>
    </dependency>
    <dependency>
      <groupId>org.springframework.cloud</groupId>
      <artifactId>spring-cloud-starter-config</artifactId>
    </dependency>
    <dependency>
      <groupId>org.springframework.cloud</groupId>
      <artifactId>spring-cloud-starter-netflix-eureka-client</artifactId>
    </dependency>

    <dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-actuator</artifactId>
    </dependency>

    <dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-data-redis</artifactId>
    </dependency>

    <dependency>
      <groupId>cn.texous.demo</groupId>
      <artifactId>demo-common-util</artifactId>
      <version>0.0.1-SNAPSHOT</version>
    </dependency>

    <dependency>
      <groupId>org.projectlombok</groupId>
      <artifactId>lombok</artifactId>
      <optional>true</optional>
    </dependency>
    <dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-test</artifactId>
      <scope>test</scope>
    </dependency>
  </dependencies>

配置 redis

package cn.texous.demo.dnw.config;

import org.springframework.cache.CacheManager;
import org.springframework.cache.annotation.CachingConfigurerSupport;
import org.springframework.cache.annotation.EnableCaching;
import org.springframework.cache.interceptor.KeyGenerator;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.cache.RedisCacheConfiguration;
import org.springframework.data.redis.cache.RedisCacheManager;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.serializer.GenericJackson2JsonRedisSerializer;
import org.springframework.data.redis.serializer.RedisSerializationContext;
import org.springframework.data.redis.serializer.StringRedisSerializer;

import java.lang.reflect.Method;

/**
 * Created by admin on 2017-11-21.
 */
@Configuration
@EnableCaching
public class RedisConfig extends CachingConfigurerSupport {

    /**
     * redis模板操作類,類似於jdbcTemplate的一個類;
     * 雖然CacheManager也能獲取到Cache對象，但是操作起來沒有那麼靈活；*
     * 這裏在擴展下：RedisTemplate這個類不見得很好操作，我們可以在進行擴展一個我們
     * 自己的緩存類，比如：RedisStorage類;
     *
     * @param factory : 通過Spring進行注入，參數在application.properties進行配置；
     * @return
     */
    @Bean
    public RedisTemplate<String, String> redisTemplate(RedisConnectionFactory factory) {
        StringRedisTemplate template = new StringRedisTemplate();
        template.setConnectionFactory(factory);
        StringRedisSerializer stringRedisSerializer = new StringRedisSerializer();
        template.setDefaultSerializer(new GenericJackson2JsonRedisSerializer());
        template.setKeySerializer(stringRedisSerializer);
        template.setValueSerializer(stringRedisSerializer);
        template.setHashKeySerializer(stringRedisSerializer);
        template.setHashValueSerializer(stringRedisSerializer);
        template.afterPropertiesSet();
        return template;
    }

    @Override
    @Bean
    public KeyGenerator keyGenerator() {
        return new KeyGenerator() {
            @Override
            public Object generate(Object target, Method method, Object... params) {
                StringBuilder sb = new StringBuilder();
                sb.append(target.getClass().getName());
                sb.append(method.getName());
                for (Object obj : params) {
                    sb.append(obj.toString());
                }
                return sb.toString();
            }
        };
    }

    @Bean
    public CacheManager cacheManager(RedisConnectionFactory factory) {
        RedisCacheConfiguration redisCacheConfiguration = RedisCacheConfiguration.defaultCacheConfig();
        StringRedisSerializer stringRedisSerializer = new StringRedisSerializer();
        redisCacheConfiguration.disableCachingNullValues();
        RedisSerializationContext.SerializationPair pair = RedisSerializationContext.SerializationPair.fromSerializer(stringRedisSerializer);
        redisCacheConfiguration.serializeKeysWith(pair);
        redisCacheConfiguration.serializeValuesWith(pair);
        return RedisCacheManager.builder(factory).cacheDefaults(redisCacheConfiguration).build();
    }

    @Bean("ratelimitLua")
    public DefaultRedisScript getRedisScript() {
        DefaultRedisScript redisScript = new DefaultRedisScript();
        redisScript.setLocation(new ClassPathResource("lua/tokenBucketRatelimit.lua"));
        redisScript.setResultType(java.lang.Long.class);
        return redisScript;
    }
    @Bean("ratelimitInitLua")
    public DefaultRedisScript getInitRedisScript() {
        DefaultRedisScript redisScript = new DefaultRedisScript();
        redisScript.setLocation(new ClassPathResource("lua/tokenBucketRatelimitInit.lua"));
        redisScript.setResultType(java.lang.Long.class);
        return redisScript;
    }

    @Bean("ratelimitByIdLua")
    public DefaultRedisScript getLimitByIdRedisScript() {
        DefaultRedisScript redisScript = new DefaultRedisScript();
        redisScript.setLocation(new ClassPathResource("lua/counterRatelimit.lua"));
        redisScript.setResultType(java.lang.Long.class);
        return redisScript;
    }

}

限流工具類編寫

package cn.texous.demo.dnw.utils;

import cn.texous.demo.dnw.common.constant.RedisKeyEnum;
import cn.texous.demo.dnw.common.constant.TokenEnum;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Collections;

@Component
public class RateLimitClient {

    @Autowired
    StringRedisTemplate stringRedisTemplate;
    @Qualifier("ratelimitLua")
    @Resource
    RedisScript<Long> ratelimitLua;

    @Qualifier("ratelimitInitLua")
    @Resource
    RedisScript<Long> ratelimitInitLua;

    @Qualifier("ratelimitByIdLua")
    @Resource
    RedisScript<Long> ratelimitByIdLua;

    public TokenEnum initToken(String key){
        TokenEnum token = TokenEnum.SUCCESS;
        Long currMillSecond = stringRedisTemplate.execute(
                (RedisCallback<Long>) redisConnection -> redisConnection.time()
        );
        /**
         * last_mill_second 最後時間毫秒
         * curr_permits 當前可用的令牌
         * max_burst 令牌桶最大值
         * rate 每秒生成幾個令牌
         * app 應用
         *
         * redis.pcall("HMSET",KEYS[1],
         "last_mill_second",ARGV[1],
         "curr_permits",ARGV[2],
         "max_burst",ARGV[3],
         "rate",ARGV[4],
         "app",ARGV[5])
         */
        Long accquire = stringRedisTemplate.execute(ratelimitInitLua,
                Collections.singletonList(getKey(key)), currMillSecond.toString(), "100", "100000", "10", "skynet");
        if (accquire == 1) {
            token = TokenEnum.SUCCESS;
        } else if (accquire == 0) {
            token = TokenEnum.SUCCESS;
        } else {
            token = TokenEnum.FAILED;
        }
        return token;
    }

    public Long accquireTokenCounter(String ip) {
        return stringRedisTemplate.execute(ratelimitByIdLua,
                Collections.singletonList(ip), "10", "1");
    }

    /**
     * 獲得key操作
     *
     * @param key
     * @return
     */
    public TokenEnum accquireToken(String key) {
        return accquireToken(key, 1);
    }

    public TokenEnum accquireToken(String key, Integer permits) {
        TokenEnum token = TokenEnum.SUCCESS;
        Long currMillSecond = stringRedisTemplate.execute(
                (RedisCallback<Long>) redisConnection -> redisConnection.time()
        );

        Long accquire = stringRedisTemplate.execute(ratelimitLua,
                Collections.singletonList(getKey(key)), permits.toString(), currMillSecond.toString());
        if (accquire == 1) {
            token = TokenEnum.SUCCESS;
        } else {
            token = TokenEnum.FAILED;
        }
        return token;
    }

    public boolean exists(String key) {
        return stringRedisTemplate.hasKey(key);
    }

    public boolean hexists(String key, String field) {
        return stringRedisTemplate.opsForHash().hasKey(key, field);
    }

    public String getKey(String key) {
        return RedisKeyEnum.RATELIMIT_KEY_PREFIX.getKey() + key;
    }

}

相關常量類

package cn.texous.demo.dnw.common.constant;

import lombok.AllArgsConstructor;
import lombok.Getter;

/**
 * insert description here
 *
 * @author Showa.L
 * @since 2019/8/21 16:41
 */
@Getter
@AllArgsConstructor
public enum RedisKeyEnum {

    RATELIMIT_KEY_PREFIX("ratelimit:", "限流key前綴");

    private String key;
    private String desc;

}

package cn.texous.demo.dnw.common.constant;

/**
 * insert description here
 *
 * @author Showa.L
 * @since 2019/8/21 16:43
 */
public enum TokenEnum {

    SUCCESS,
    FAILED;

    public boolean isSuccess() {
        return this.equals(SUCCESS);
    }

    public boolean isFailed() {
        return this.equals(FAILED);
    }

}

對應的 lua 腳本：

計數器算法

local key = "ratelimit:" .. KEYS[1]
local limit = tonumber(ARGV[1])
local expire_time = ARGV[2]

local is_exists = redis.call("EXISTS", key)
if is_exists == 1 then
    if redis.call("INCR", key) > limit then
        return 0
    else
        return 1
    end
else
    redis.call("SET", key, 1)
    redis.call("EXPIRE", key, expire_time)
    return 1
end

令牌桶算法

令牌桶初始化

local result=1
redis.pcall("HMSET",KEYS[1],
        "last_mill_second",ARGV[1],
        "curr_permits",ARGV[2],
        "max_burst",ARGV[3],
        "rate",ARGV[4],
        "app",ARGV[5])
return result

令牌桶算法

local ratelimit_info=redis.pcall("HMGET",KEYS[1],"last_mill_second","curr_permits","max_burst","rate","app")
local last_mill_second=ratelimit_info[1]
local curr_permits=tonumber(ratelimit_info[2])
local max_burst=tonumber(ratelimit_info[3])
local rate=tonumber(ratelimit_info[4])
local app=tostring(ratelimit_info[5])
if app == nil then
    return 0
end

local local_curr_permits=max_burst;

if(type(last_mill_second) ~='boolean' and last_mill_second ~=nil) then
    local reverse_permits=math.floor((ARGV[2]-last_mill_second)/1000)*rate
    if(reverse_permits>0) then
        redis.pcall("HMSET",KEYS[1],"last_mill_second",ARGV[2])
    end

    local expect_curr_permits=reverse_permits+curr_permits
    local_curr_permits=math.min(expect_curr_permits,max_burst);

else
    redis.pcall("HMSET",KEYS[1],"last_mill_second",ARGV[2])
end

local result=-1
if(local_curr_permits-ARGV[1]>=0) then
    result=1
    redis.pcall("HMSET",KEYS[1],"curr_permits",local_curr_permits-ARGV[1])
else
    redis.pcall("HMSET",KEYS[1],"curr_permits",local_curr_permits)
end

return result

測試類編寫

package cn.texous.demo.dnw.utils;

import cn.texous.demo.dnw.DemoNormalWebApplicationTests;
import cn.texous.demo.dnw.common.constant.TokenEnum;
import org.junit.Test;
import org.springframework.beans.factory.annotation.Autowired;

/**
 * insert description here
 *
 * @author Showa.L
 * @since 2019/8/21 17:56
 */
public class RateLimitClientTests extends DemoNormalWebApplicationTests {

    @Autowired
    private RateLimitClient rateLimitClient;

    @Test
    public void testGetToken() throws InterruptedException {
        int scount = 0;
        int fcount = 0;
        String base = "192.168.0.";
        String ip = base + "1";
        String key = rateLimitClient.getKey(ip);
        long start = System.currentTimeMillis();

        for (int i = 0; i < 10000; i++) {
//            Long success = rateLimitClient.acquireToken(ip);
//            key = rateLimitClient.getKey(ip);
            if (!rateLimitClient.exists(key))
                rateLimitClient.initToken(ip);

            TokenEnum success = rateLimitClient.accquireToken(ip);
            if (success.isSuccess())
                scount++;
            else {
//                ip = base + i;
                fcount++;
            }
//            if (i % 500 == 0)
//                Thread.sleep(1000);
        }
        System.out.println(System.currentTimeMillis() - start);
        System.out.println("success count: " + scount);
        System.out.println("failed count: " + fcount);
    }


}

結語

令牌桶算法相關代碼如上

參考文獻

《redis+lua 實現分佈式令牌桶，高併發限流》

《令牌桶算法》

《Redis + Lua Java實現限流》