Jedis uses lua script to complete token bucket throttling

Jedis uses lua script to complete token bucket throttling

1, Simple syntax for lua scripts

KEYS[1]
ARGV[1]
These two parameters represent the first element of the key array and the first element of the arg array
Let's look at a simple use
eval "return redis.call('get',KEYS[1])" 1 one

First, eval is the keyword for parsing Lua script. The written Lua script is in the string, where
redis.call() means to call redis. 1 indicates the first bit of the array, and one indicates that KEY[1] is passed into one.

Here is an example of jedis calling lua script

    public static void main(String[] args) throws IOException {
        Jedis jedis = new Jedis("127.0.0.1", 6379);
        ClassPathResource classPathResource = new ClassPathResource("/META-INF/script/redis_limiter.lua");
        byte[] buffer = new byte[(int)classPathResource.getFile().length()];
        classPathResource.getInputStream().read(buffer);
        String script = new String(buffer);
        System.out.println(script);
        String lua = jedis.scriptLoad(script);
        Object evalsha = jedis.evalsha(lua, getKeys(), getArgs());
        System.out.println(evalsha);
    }

    public static List<String> getKeys() {
        return Arrays.asList("one");
    }

    public static List<String> getArgs() {
        return Arrays.asList("50");
    }
local key=KEYS[1]
local arg=ARGV[1]

local value=tonumber(arg)
local currentValue=tonumber(redis.call('get', key))
if currentValue==nil then
    currentValue=0
end
local new_value=(currentValue+value)

redis.call('set',key,new_value)

In the Lue script above, the key value and parameter value are defined, and tonumber() means to convert to a number.
The general meaning of this paragraph is to get the value of key and determine whether it is empty or empty. It means that the value starts from 0 and then calls the redis function to assign a new value to key.

2, Token bucket current limiting

1. Conception

First of all, it is necessary to clarify what is token bucket current limit. In fact, it is to issue tokens to the bucket per unit time. If there is a request, request to apply for a token. If there is a token, you can get the token and carry out normal operation. If you fail to obtain the token, you will refuse access. (the capacity in the bucket should be limited, otherwise no one will access it for a period of time, and the token will be used all the time. When the traffic peak reaches, there are many tokens in the bucket, which does not limit the flow.)

For this purpose, we can make two points clear:

  • Issuing tokens at a constant rate
  • Token bucket has maximum capacity

2. Realization

With Lua script, we can achieve atomicity, because redis processes business in a single thread mode (another thread will be created during master-slave replication)
Then we can design it like this
First, a key value is stored in redis to represent the token bucket and the number of tokens in the token bucket
Then store a timestamp of the last access
In the lua script, we pass in four parameters

  • Rate of token buckets
  • Capacity of token bucket
  • Request time stamp
  • Number of tokens requested
    Use the redis.call() function call to get the current number of token buckets and the timestamp of the last refresh
    We can calculate the difference between this time and the last time, calculate how many tokens are generated between the two accesses, and then use the number of token buckets generated by the last access + the capacity generated by the intermediate access = the number in the current token bucket.
    Then we can judge whether the number of tokens currently requested is greater than the number in the current bucket. If it is greater than return false and less than return true, we can refresh the number of token buckets and timestamp stored in redis before returning
    Now let's fight
--Incoming token bucket key And timestamp key
local token_key=KEYS[1]
local time_key=KEYS[2]

-- Incoming rate, capacity, current time and number of requested tokens respectively
local rate=tonumber(ARGV[1])
local capacity=tonumber(ARGV[2])
local now_time=tonumber(ARGV[3])
local requestNum=tonumber(ARGV[4])
--Gets the token tree and time of the last visit
local last_tokens=tonumber(redis.call('get',token_key))
if last_tokens==nil then
    last_tokens=capacity
end
local last_time=tonumber(redis.call('get',time_key))
if last_time==nil then
    last_time=0
end
--1´╝îCalculate time difference
local time_del=math.max(0,(now_time-last_time))
--2,Calculate how many tokens should be in the current token bucket
local current_tokens=math.min(capacity,(last_tokens+time_del*rate))
--3,Determine whether access is allowed
local acuire=0
if (requestNum <= current_tokens) then
    acuire=1
    current_tokens=(current_tokens-requestNum)
end

-- Calculate the expiration time (the default is the time to fill the token)*2)
local ttl = 60
--4´╝îimprove a record
    redis.call('setex',token_key,ttl,current_tokens)
    redis.call('setex',time_key,ttl,now_time)

return { acuire }

java code
Current limiting interceptor

package com.xzq.config;

@Component
public class LimiterIntercepter implements HandlerInterceptor, InitializingBean {
    private Logger logger = LoggerFactory.getLogger(LimiterIntercepter.class);
    private static String TOKEN_BUCKET = "TOKEN_BUCKET";
    private static String REFRESH_TIME = "REFRESH_TIME";
    private String scriptLua;
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (isallow()) {
            logger.info("Allow entry");
            return true;
        }else{
            logger.info("Restricted access");
            response.setStatus(500);
            return false;
        }
    }

    public boolean isallow() {
        Jedis jedis = new Jedis("127.0.0.1", 6379);
        long result = (long)((List) jedis.evalsha(scriptLua, getKeys(), getArgs())).get(0);
        return result == 1L;
    }

    public static List<String> getKeys() {
        return Arrays.asList(TOKEN_BUCKET, REFRESH_TIME);
    }

    public static List<String> getArgs() {
        String now = String.valueOf(System.currentTimeMillis() / 1000);
        // Rate: 1, capacity: 5, current time second value: now, number of request Tokens: 1
        return Arrays.asList("1", "5", now, "1");
    }

    @Override
    public void afterPropertiesSet() throws Exception {
	   Jedis jedis = new Jedis("127.0.0.1", 6379);

        ClassPathResource classPathResource = new ClassPathResource("/META-INF/script/redis_limiter.lua");
        byte[] buffer = new byte[(int)classPathResource.getFile().length()];
        classPathResource.getInputStream().read(buffer);
        scriptLua = jedis.scriptLoad(new String(buffer));
    }
}

mvc configuration

package com.xzq.config;
@Component
public class MvcConfig implements WebMvcConfigurer {

    @Autowired
    private LimiterIntercepter limiterIntercepter;
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(limiterIntercepter);
    }
}

3, Jemeter pressure measuring tool test

Test with Jemeter pressure measuring tool

You can see that only five requests are allowed to enter in one second, because the token bucket capacity is 5, the initial value is 5, and then one request is entered in one second, because we set the rate to 1.

Tags: Redis

Posted on Fri, 26 Nov 2021 17:45:06 -0500 by ghe