Leaky bucket algorithm and token bucket algorithm

catalogue

1, Introduction to leaky bucket and token bucket

2, Redis + lua script + token bucket algorithm to realize current limiting control

1. Customize an annotation to label the current limiting method

2. Write lua script

3. Read lua script

4. Create an interceptor to intercept the method with this annotation

5. Register this interceptor in WebConfig

6. Annotation use

1, Introduction to leaky bucket and token bucket

Bucket leakage algorithm And Token Bucket It looks similar on the surface, so it's easy to confuse the two. But in fact, they have very different characteristics and are used for different purposes. Leaky bucket algorithm And Token Bucket The difference is: l Leaky bucket algorithm Ability to impose restrictions on data transmission speed . l Token Bucket Can limit the average of data transmission speed Same as

Some degree of Burst transmission . It should be noted that in some cases, leaky bucket algorithm can not effectively use network resources. Because the leakage rate of leaky bucket is fixed, even if there is no congestion in the network, leaky bucket algorithm can not make a single data flow reach the port rate. Therefore, the leaky bucket algorithm is suitable for the problems with burst characteristics

Inefficient in terms of traffic. The token bucket algorithm can meet these traffic with burst characteristics. Generally, the leaky bucket algorithm is combined with the token bucket algorithm network flow Provide more efficient control.

There are two common current limiting algorithms: leaky bucket algorithm and token bucket algorithm.

The idea of leaky bucket algorithm is very simple. Water (request) enters the leaky bucket first, and the leaky bucket leaves the water at a certain speed. When the water inflow speed is too high, it will overflow directly. It can be seen that the leaky bucket algorithm can forcibly limit the data transmission rate.

Schematic diagram of leaky bucket algorithm

For many application scenarios, in addition to limiting the average transmission rate of data, it is also required to allow some degree of burst transmission. At this time, the leaky bucket algorithm may not be suitable, and the token bucket algorithm is more suitable. As shown in Figure 2, the principle of token bucket algorithm is that the system will put tokens into the bucket at a constant speed, and if the request needs to be processed

If there is no token in the bucket, the service will be denied.

Schematic diagram of token bucket algorithm

It does not mean that the token bucket must be better than the vulnerability. Their use scenarios are different. The token bucket can be used to protect yourself. It is mainly used to limit the current of the caller's frequency in order not to be defeated. Therefore, if you have processing capacity, if the traffic bursts (the actual consumption capacity is stronger than the configured traffic limit), the actual processing rate can be increased

To exceed the configured limit. The leaky bucket algorithm is used to protect others, that is, to protect the system he calls. The main scenario is that when the calling third-party system itself has no protection mechanism or has traffic restrictions, our call speed cannot exceed its restrictions. Because we cannot change the third-party system, we can only control it at the main caller.

At this time, even if the traffic is sudden, it must be abandoned. Because the consumption power is determined by a third party.

To sum up: if you want to keep your system from breaking down, use token bucket. If you ensure that other people's system will not be broken, use the leaky bucket algorithm.

2, Redis + lua script + token bucket algorithm to realize current limiting control

1. Customize an annotation to label the current limiting method

@Target({ElementType.TYPE, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    //Unique indication of current limit
    String key() default "";
 
    //Current limiting unit time (unit: s)
    int time() default 1;
 
    //Limited number of visits per unit time
    int count();
 
    //Restrict ip
    boolean ipLimit() default false;
}

2. Write lua script

according to key(Parameter) query the corresponding value(Number of tokens)
	If yes null Explain the key It's the first time to enter 
	{
		Number of initialization token buckets (parameters);Record initialization time ->Returns the number of remaining tokens
	} 
	
	If not null
	{
		judge value Is it greater than 1 
		{
			Greater than 1  ->value - 1  -> Returns the number of remaining tokens
			Less than 1  -> Determine whether the time interval for replenishing tokens is sufficient
			{
				enough -> Supplementary token; Update replenishment token time-> Returns the number of remaining tokens
				Not enough	-> return -1 (Description: the current limit access times have been exceeded)
			}
		}
	}
	
redis.replicate_commands();
-- Passed in parameter key
local key = KEYS[1]
-- Minimum time interval for token bucket filling
local update_len = tonumber(ARGV[1])
-- Record current key The time when the token bucket was last updated key
local key_time = 'ratetokenprefix'..key
-- Get current time(there curr_time_arr The first is the number of seconds, and the second is the number of milliseconds after the number of seconds),Since I calculate by seconds, here is only curr_time_arr[1](be careful: redis Array subscripts start with 1)
--If the number of milliseconds is required; otherwise tonumber(arr[1]*1000 + arr[2])
local curr_time_arr = redis.call('TIME')
-- Current time seconds
local nowTime = tonumber(curr_time_arr[1])
-- from redis Get current key Corresponding to the last updated token bucket key Corresponding value
local curr_key_time = tonumber(redis.call('get',KEYS[1]) or 0)
-- Get current key Number of tokens in the corresponding token bucket
local token_count = tonumber(redis.call('get',KEYS[1]) or -1)
-- Current token bucket capacity
local token_size = tonumber(ARGV[2])
-- The number of token buckets is less than 0, indicating that the token bucket is not initialized
if token_count < 0 then
	redis.call('set',key_time,nowTime)
	redis.call('set',key,token_size -1)
	return token_size -1
else
	if token_count > 0 then --There are enough tokens in the current token bucket
		redis.call('set',key,token_count - 1)
		return token_count -1   --Returns the number of remaining tokens
	else    --The number of tokens in the current token bucket has been emptied
		if curr_key_time + update_len < nowTime then    --Judge whether the interval between the current time seconds and the last update time seconds is greater than the specified time interval( update_len)
			redis.call('set',key,token_size -1)
			return token_size - 1
		else
			return -1
		end
	end
end

3. Read lua script

@Component
public class CommonConfig {
    /**
     * Read current limiting script
     */
    @Bean
    public DefaultRedisScript<Number> redisluaScript() {
        DefaultRedisScript<Number> redisScript = new DefaultRedisScript<>();
        //The path of the script here is path for source root
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("myLua.lua"))); 
        redisScript.setResultType(Number.class);
        return redisScript;
    }
    /**
     * RedisTemplate
     */
    @Bean
    public RedisTemplate<String, Serializable> limitRedisTemplate(LettuceConnectionFactory redisConnectionFactory) {
        RedisTemplate<String, Serializable> template = new RedisTemplate<String, Serializable>();
        template.setKeySerializer(new StringRedisSerializer());
        template.setValueSerializer(new GenericJackson2JsonRedisSerializer());
        template.setConnectionFactory(redisConnectionFactory);
        return template;
    }
}

4. Create an interceptor to intercept the method with this annotation

@Component
public class RateLimitInterceptor implements HandlerInterceptor {
    private final Logger LOG = LoggerFactory.getLogger(this.getClass());
    
    @Autowired
    private RedisTemplate<String, Serializable> limitRedisTemplate;
 
    @Autowired
    private DefaultRedisScript<Number> redisLuaScript;
    
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        assert handler instanceof HandlerMethod;
        HandlerMethod method = (HandlerMethod) handler;
        RateLimit rateLimit = method.getMethodAnnotation(RateLimit.class);
        //There are our custom annotations on the current method
        if (rateLimit != null) {
            //Get the limited number of visits per unit time
            int count = rateLimit.count();
            String key = rateLimit.key();
            //Obtain current limiting unit time (unit: s)
            int time = rateLimit.time();
            boolean ipLimit = rateLimit.ipLimit();
            //Splicing key s in redis
            StringBuilder sb = new StringBuilder();
            sb.append(Constants.RATE_LIMIT_KEY).append(key).append(":");
            //If you need to limit ip
            if(ipLimit){
                sb.append(getIpAddress(request)).append(":");
            }
            List<String> keys = Collections.singletonList(sb.toString());
           //Execute lua script
            Number execute = limitRedisTemplate.execute(redisLuaScript, keys, time, count);
            assert execute != null;
            if (-1 == execute.intValue()) {
                ResultModel resultModel = ResultModel.error_900("Interface call exceeds current limit");
                response.setStatus(901);
                response.setCharacterEncoding("utf-8");
                response.setContentType("application/json");
                response.getWriter().write(JSONObject.toJSONString(resultModel));
                response.getWriter().flush();
                response.getWriter().close();
                LOG.info("The current interface call exceeds the current limit within the time period,key:{}", sb.toString());
                return false;
            } else {
                LOG.info("Remaining in the current access period{}Number of visits", execute.toString());
            }
        }
        return true;
    }
 
    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
 
    }
 
    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
 
    }
    
    public static String getIpAddr(HttpServletRequest request) {
        String ipAddress = null;
        try {
            ipAddress = request.getHeader("x-forwarded-for");
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getRemoteAddr();
            }
            // For the case of passing through multiple agents, the first IP is the real IP of the client, and multiple IPS are divided according to ','
            // "***.***.***.***".length()
            if (ipAddress != null && ipAddress.length() > 15) { 
                // = 15
                if (ipAddress.indexOf(",") > 0) {
                    ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
                }
            }
        } catch (Exception e) {
            ipAddress = "";
        }
        return ipAddress;
    }
 
}

A custom constant is used as the redis prefix

public class Constants {
    public static final String RATE_LIMIT_KEY = "rateLimit:";
}

5. Register this interceptor in WebConfig

@Configuration
@EnableWebMvc
public class WebConfig extends WebMvcConfigurerAdapter {
 
    @Autowired
    private RateLimitInterceptor rateLimitInterceptor;
 
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(rateLimitInterceptor);
        super.addInterceptors(registry);
    }
}

6. Annotation use

@RestController
@RequestMapping(value = "/test")
public class TestController {
 
    //The flow restriction rule is that only five requests can be sent by the same ip within one second
    @RateLimit(key = "testGet",time = 1,count = 5,ipLimit = true)
    @RequestMapping(value = "/get")
    public ResultModel testGet(){
        return ResultModel.ok_200();
    }
 
}

 

Keywords: Java

Added by SuisydeKing on Wed, 09 Feb 2022 10:41:42 +0200