接口限流——令牌桶算法
简介:
在网络中传输数据时,为了防止网络拥塞,需限制流出网络的流量,使流量以比较均匀的速度向外发送。令牌桶算法就实现了这个功能,可控制发送到网络上数据的数目,并允许突发数据的发送。
令牌桶算法是网络流量整形(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);
}
}
结语
令牌桶算法相关代码如上