Recently, you need to do a current limiting function by yourself , Other business codes are easy to say . The only end point is the realization of current limiting , The thought of redis Token bucket can be realized . A slap on the forehead , Just use it ！

Scene description

In real development, I found that what I thought was too simple , If it is based on redis If the provided command is called in the code , Efficiency is a small thing . Atomicity cannot be guaranteed ！ as follows ：

1. Threads 1 The total number of token buckets obtained is 10
2. Threads 1 Consume 1 A token , The remaining 9 A token
3. Threads 2 The total number of token buckets obtained is 10
4. Threads 1 The refresh token is 9 individual
5. Threads 2 Consume 1 A token , The remaining 9 A token
6. Threads 2 The refresh token is 9 individual

Atomicity cannot be guaranteed at all . Lock ？ Don't multiple nodes need to introduce distributed locks ？ It seems that the implementation in the server is not desirable , If atomic , Is bound to write LUA The code executed . Read some online demo, I didn't find what I wanted . Think of what you built spring cloud gateway It is implemented with a built-in token bucket . Start browsing the source code , Finally find the best solution .

gateway in redis The token bucket implementation class is ：org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter

public Mono<Response> isAllowed(String routeId, String id) {
    
			...
			//  This line is through lua Judge whether the current is limited 
			Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
			...
	}

Follow this class to find lua The code is from org.springframework.cloud.gateway.config.GatewayRedisAutoConfiguration Infused with

public RedisScript redisRequestRateLimiterScript() {
    
		DefaultRedisScript redisScript = new DefaultRedisScript<>();
		redisScript.setScriptSource(
				new ResourceScriptSource(new ClassPathResource("META-INF/scripts/request_rate_limiter.lua")));
		redisScript.setResultType(List.class);
		return redisScript;
	}

The core in the request_rate_limiter.lua In this file

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

local fill_time = capacity/rate
local ttl = math.floor(fill_time*2)

--redis.log(redis.LOG_WARNING, "rate " .. ARGV[1])
--redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2])
--redis.log(redis.LOG_WARNING, "now " .. ARGV[3])
--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
--redis.log(redis.LOG_WARNING, "ttl " .. ttl)

local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end
--redis.log(redis.LOG_WARNING, "last_tokens " .. last_tokens)

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end
--redis.log(redis.LOG_WARNING, "last_refreshed " .. last_refreshed)

local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end

--redis.log(redis.LOG_WARNING, "delta " .. delta)
--redis.log(redis.LOG_WARNING, "filled_tokens " .. filled_tokens)
--redis.log(redis.LOG_WARNING, "allowed_num " .. allowed_num)
--redis.log(redis.LOG_WARNING, "new_tokens " .. new_tokens)

if ttl > 0 then
  redis.call("setex", tokens_key, ttl, new_tokens)
  redis.call("setex", timestamp_key, ttl, now)
end

-- return { allowed_num, new_tokens, capacity, filled_tokens, requested, new_tokens }
return {
     allowed_num, new_tokens }

Solution

Okay , You can start writing your own current limiting tool class .

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.time.Instant;
import java.util.Arrays;
import java.util.List;

/** * redis Current limiter  * * @Author: dong * @Date: 2021/11/26 14:08 */
public class RedisLimiter{
    

    private RedisTemplate redisTemplate;

    private static final Long SUCCESS_FLAG = 1L;


    /** *  Determine if access is allowed  *@Author dong *@Date 2022/4/26 15:30 *@param id  Get the token bucket this time id *@param rate  Fill rate per second  *@param capacity  Maximum capacity of token bucket  *@param tokens  Each access consumes several tokens  *@return true  allow access to  false  Access not allowed  */
    public boolean isAllowed(String id,int rate,int capacity,int tokens){
    
        RedisScript<Long> redisScript = new DefaultRedisScript<>(SCRIPT,Long.class);
        Object result = redisTemplate.execute(redisScript,
                getKey(id),rate, capacity,
                Instant.now().getEpochSecond(), tokens);
        return SUCCESS_FLAG.equals(result);
    }

    private List<String> getKey(String id){
    
        String prefix = "limiter:"+id;
        String tokenKey = prefix + ":tokens";
        String timestampKey = prefix + ":timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    private static final String SCRIPT = "local tokens_key = KEYS[1]\n" +
            "local timestamp_key = KEYS[2]\n" +
            "local rate = tonumber(ARGV[1])\n" +
            "local capacity = tonumber(ARGV[2])\n" +
            "local now = tonumber(ARGV[3])\n" +
            "local requested = tonumber(ARGV[4])\n" +
            "local fill_time = capacity/rate\n" +
            "local ttl = math.floor(fill_time*2)\n" +
            "local last_tokens = tonumber(redis.call('get', tokens_key))\n" +
            "if last_tokens == nil then\n" +
            " last_tokens = capacity\n" +
            "end\n" +
            "local last_refreshed = tonumber(redis.call('get', timestamp_key))\n" +
            "if last_refreshed == nil then\n" +
            " last_refreshed = 0\n" +
            "end\n" +
            "local diff_time = math.max(0, now-last_refreshed)\n" +
            "local filled_tokens = math.min(capacity, last_tokens+(diff_time*rate))\n" +
            "local allowed = filled_tokens >= requested\n" +
            "local new_tokens = filled_tokens\n" +
            "local allowed_num = 0\n" +
            "if allowed then\n" +
            " new_tokens = filled_tokens - requested\n" +
            " allowed_num = 1\n" +
            "end\n" +
            "if ttl > 0 then\n" +
            " redis.call('setex', tokens_key, ttl, new_tokens)\n" +
            " redis.call('setex', timestamp_key, ttl, now)\n" +
            "end\n" +
            "return allowed_num\n";
}

When you , This token bucket implementation is completely reference spring cloud gateway, It's safe to eat .
Briefly describe the effect of this tool class .

	isAllowed("testId1",1,60,1);

The above description represents ： Token bucket testId1. Every minute you can visit 60 Time .

Of course, this is the ideal situation . In extreme cases , It should be accessible 120 Time of . The extreme scenario is as follows

• testId1 The time when the token bucket is not used >=60 second
• testId1 The token bucket starts to be used at a certain moment
• The token is consumed and filled at the same time
• Originally used 60 Seconds , The token bucket initially has 60 A token , The service life has been filled 60 individual

I personally understand , Token bucket is mainly used to ensure the usage rate . For the above scenario , Whether or not bug. Look at everyone's use . However, I have solved the above situation , Only the following minor changes are needed .

1. The maximum number of traffic in the incoming cycle time （ Cycle time ： Barrel capacity 60, Filling rate 1/s, So cycle time =60s）
2. obtain key Method to increase the number of times in a recording cycle key
3. modify lua Script

The changes are as follows ：

/** *  Determine if access is allowed  *@Author dong *@Date 2022/4/26 15:30 *@param id  Get the token bucket this time id *@param rate  Fill rate per second  *@param capacity  Maximum capacity of token bucket  *@param tokens  Each access consumes several tokens  *@param maxCount  Maximum visits in cycle time  *@return true  allow access to  false  Access not allowed  */
    public boolean isAllowed(String id,int rate,int capacity,int tokens,int maxCount){
    
        RedisScript<Long> redisScript = new DefaultRedisScript<>(SCRIPT,Long.class);
        Object result = redisTemplate.execute(redisScript,
                getKey(id),rate, capacity,
                Instant.now().getEpochSecond(), tokens,maxCount);
        return SUCCESS_FLAG.equals(result);
    }

    private List<String> getKey(String id){
    
        String prefix = "limiter:"+id;
        String tokenKey = prefix + ":tokens";
        String timestampKey = prefix + ":timestamp";
        String countKey = prefix + ":count";
        return Arrays.asList(tokenKey, timestampKey,countKey);
    }

    private static final String SCRIPT = "local tokens_key = KEYS[1]\n" +
            "local timestamp_key = KEYS[2]\n" +
            "local count_key = KEYS[3]\n" +
            "local rate = tonumber(ARGV[1])\n" +
            "local capacity = tonumber(ARGV[2])\n" +
            "local now = tonumber(ARGV[3])\n" +
            "local requested = tonumber(ARGV[4])\n" +
            "local min_max = tonumber(ARGV[5])\n" +
            "local fill_time = capacity/rate\n" +
            "local ttl = math.floor(fill_time*2)\n" +
            "local has_count = tonumber(redis.call('get', count_key))\n" +
            "if has_count == nil then\n" +
            " has_count = 0\n" +
            "end\n" +
            "if has_count >= min_max then\n" +
            "return 0\n" +
            "end\n" +
            "local last_tokens = tonumber(redis.call('get', tokens_key))\n" +
            "if last_tokens == nil then\n" +
            " last_tokens = capacity\n" +
            "end\n" +
            "local last_refreshed = tonumber(redis.call('get', timestamp_key))\n" +
            "if last_refreshed == nil then\n" +
            " last_refreshed = 0\n" +
            "end\n" +
            "local diff_time = math.max(0, now-last_refreshed)\n" +
            "local filled_tokens = math.min(capacity, last_tokens+(diff_time*rate))\n" +
            "local allowed = filled_tokens >= requested\n" +
            "local new_tokens = filled_tokens\n" +
            "local allowed_num = 0\n" +
            "if allowed then\n" +
            " new_tokens = filled_tokens - requested\n" +
            " allowed_num = 1\n" +
            "end\n" +
            "if ttl > 0 then\n" +
            " redis.call('setex', tokens_key, ttl, new_tokens)\n" +
            " redis.call('setex', timestamp_key, ttl, now)\n" +
            "end\n" +
            "local count_ttl = tonumber(redis.call('ttl',count_key))\n" +
            "if count_ttl < 0 then\n" +
            " count_ttl = fill_time\n" +
            "end\n" +
            "redis.call('setex', count_key,count_ttl , has_count+1)\n" +
            "return allowed_num\n";
}

Such changes can keep the access rate and throughput controllable , But whether it is necessary depends on your own needs .

Finally, I want to say , Excellent open source projects are really treasures . Learn to use it , Then stand on the shoulders of giants and move forward