限流算法（三）SpringBoot调用redis-lua脚本

引入redis依赖

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

资源目录新建scripts文件夹，将上一节讲解的lua脚本放进去

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 }

构造redis-script对象

springboot将每个lua脚本抽象为一个RedisScript对象，该类提供了两个方法，一个是设置lua脚本的io流，还有一个是直接将lua脚本以字符串的形式设置，这里用io流的形式。
该对象的泛型是lua脚本的返回值，我们的脚本返回的是两个long类型，所以使用List<Long>来接收。

    @Bean(name = "rateLimitRedisScript")
    public RedisScript<List<Long>> rateLimitRedisScript() {
        DefaultRedisScript redisScript = new DefaultRedisScript<>();
//        redisScript.setScriptText();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("scripts/redis_rate_limit.lua")));
        redisScript.setResultType(List.class);
        return redisScript;
    }

设置redis序列化规则

@Bean
    @ConditionalOnMissingBean(StringRedisTemplate.class)
    public StringRedisTemplate stringRedisTemplate(RedisConnectionFactory redisConnectionFactory) throws UnknownHostException {
        StringRedisTemplate template = new StringRedisTemplate();
        template.setConnectionFactory(redisConnectionFactory);
        return template;
    }

    @Bean
    @ConditionalOnMissingBean(RedisTemplate.class)
    public RedisTemplate<String, Object> redisTemplate(
            RedisConnectionFactory redisConnectionFactory)
            throws UnknownHostException {

        Jackson2JsonRedisSerializer<Object> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer<Object>(Object.class);
        ObjectMapper om = new ObjectMapper();
        om.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
        om.enableDefaultTyping(ObjectMapper.DefaultTyping.NON_FINAL);
        jackson2JsonRedisSerializer.setObjectMapper(om);

        RedisTemplate<String, Object> template = new RedisTemplate<String, Object>();
        template.setConnectionFactory(redisConnectionFactory);
        template.setKeySerializer(jackson2JsonRedisSerializer);
        template.setValueSerializer(jackson2JsonRedisSerializer);
        template.setHashKeySerializer(jackson2JsonRedisSerializer);
        template.setHashValueSerializer(jackson2JsonRedisSerializer);
        template.afterPropertiesSet();
        return template;
    }

调用lua脚本

    @Resource
    private RedisScript<List<Long>> rateLimitRedisScript;

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    @GetMapping
    public List<Long> userToken() {
        // 设置lua脚本的ARGV的值
        List<String> scriptArgs = Arrays.asList(
                1 + "",
                3 + "",
                (Instant.now().toEpochMilli()) + "",
                "1");
        // 设置lua脚本的KEYS值
        List<String> keys = getKeys("test");
        return stringRedisTemplate.execute(rateLimitRedisScript,keys, scriptArgs.toArray());
    }

    private List<String> getKeys(String id) {
        // use `{}` around keys to use Redis Key hash tags
        // this allows for using redis cluster

        // Make a unique key per user.
        String prefix = "request_rate_limiter.{" + id;

        // You need two Redis keys for Token Bucket.
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

调用成功后会返回两个数，一个是是否成功标志，0代表限流，1代表未限流，还有一个是令牌桶中剩余的令牌数

[
  1,
  2
]

注意

• 这里在拼接key的时候，对id使用了大括号{}进行了包裹，这是因为lua脚本执行成功的前提条件是所用使用到的redis健值必须在一个hash槽中，使用大括号对key进行包裹后，redis在对key进行hash时，指挥hash大括号内部的字符，这样就可以保证lua脚本中的使用的key-value在同一个槽内。这样就确保了cluster模式下正常执行redis-lua脚本，但是需要注意的是，这里大括号内包裹的内容不能是不变的，如果是不变的话，会有大量的key-value被分配到同一个槽里，导致hash倾斜，key-value分布不均匀。
• 这里使用的不是RedisTemplate,而是使用的StringRedisTemplate执行lua脚本的，使用RedisTemplate执行lua脚本的时候，会报错。

关于springboot执行lua脚本，就介绍到这，下一节会具体使用redis-lua结合aop做一个限流工具

限流算法（四）AOP+RedisLua对接口进行限流