高性能分布式限流：Redis+Lua真香！

-- 获取调用脚本时传入的第一个key值（用作限流的 key）
local key = KEYS[1]
-- 获取调用脚本时传入的第一个参数值（限流大小）
local limit = tonumber(ARGV[1])

-- 获取当前流量大小
local curentLimit = tonumber(redis.call('get', key) or "0")

-- 是否超出限流
if curentLimit + 1 > limit then
    -- 返回(拒绝)
    return 0
else
    -- 没有超出 value + 1
    redis.call("INCRBY", key, 1)
    -- 设置过期时间
    redis.call("EXPIRE", key, 2)
    -- 返回(放行)
    return 1
end

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>
<dependency>
    <groupId>redis.clients</groupId>
    <artifactId>jedis</artifactId>
</dependency>
<dependency>
    <groupId>com.google.guava</groupId>
    <artifactId>guava</artifactId>
    <version>21.0</version>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
</dependency>
<dependency>
    <groupId>org.apache.commons</groupId>
    <artifactId>commons-lang3</artifactId>
    <version>3.9</version>
</dependency>
<dependency>
    <groupId>org.projectlombok</groupId>
    <artifactId>lombok</artifactId>
    <optional>true</optional>
</dependency>

spring.redis.host=127.0.0.1
spring.redis.port=6379
spring.redis.database=0

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.connection.jedis.JedisConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.serializer.GenericJackson2JsonRedisSerializer;
import org.springframework.data.redis.serializer.StringRedisSerializer;

import java.io.Serializable;

/**
 * redis配置类.
 *
 * @author : 北哥 公众号：BiggerBoy
 * @version : 1.0 2022/09/09
 * @since : 1.0
 */
@Configuration
public class RedisConfig {
    @Bean
    JedisConnectionFactory jedisConnectionFactory() {
        return new JedisConnectionFactory();
    }

    @Bean
    public RedisTemplate<String, Serializable> limitRedisTemplate(RedisConnectionFactory redisConnectionFactory) {
        RedisTemplate<String, Serializable> template = new RedisTemplate<>();
        template.setKeySerializer(new StringRedisSerializer());
        template.setValueSerializer(new GenericJackson2JsonRedisSerializer());
        template.setConnectionFactory(redisConnectionFactory);
        return template;
    }
}

import com.biggerboy.redislimiter.enums.LimitType;
import java.lang.annotation.*;

/**
 * redis限流注解.
 *
 * @author : 北哥 公众号：BiggerBoy
 * @version : 1.0 2022/09/09
 * @since : 1.0
 */
@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface MyRedisLimiter {
    /**
     * 缓存到Redis的key
     */
    String key();
    /**
     * Key的前缀
     */
    String prefix() default "limiter:";
    /**
     * 给定的时间范围 单位(秒)
     * 默认1秒 即1秒内超过count次的请求将会被限流
     */
    int period() default 1;
    /**
     * 一定时间内最多访问的次数
     */
    int count();
    /**
     * 限流的维度(用户自定义key 或者 调用方ip)
     */
    LimitType limitType() default LimitType.CUSTOMER;
}

/**
 * @param pjp
 * @author 问北(北哥) 公众号：BiggerBoy
 * @description 切面
 * @date 2022-9-8 18:29:53
 */
@Around("execution(public * *(..)) && @annotation(com.wenbei.annotation.MyRedisLimiter)")
public Object limit(ProceedingJoinPoint pjp) {
    MethodSignature signature = (MethodSignature) pjp.getSignature();
    Method method = signature.getMethod();
    MyRedisLimiter limitAnnotation = method.getAnnotation(MyRedisLimiter.class);
    LimitType limitType = limitAnnotation.limitType();
    int limitPeriod = limitAnnotation.period();
    int limitCount = limitAnnotation.count();

    String key = getKey(limitAnnotation, limitType);
    ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix(), key));
    try {
        Number count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
        logger.info("try to access, this time count is {} for key: {}", count, key);
        if (count != null && count.intValue() <= limitCount) {
            return pjp.proceed();
        } else {
            demote();//降级
            return null;
        }
    } catch (Throwable e) {
        if (e instanceof RuntimeException) {
            throw new RuntimeException(e.getLocalizedMessage());
        }
        throw new RuntimeException("服务器出现异常，请稍后再试");
    }
}

private static final String LIMIT_LUA_PATH = "limit.lua";
private DefaultRedisScript<Number> redisScript;
@PostConstruct
public void init() {
    redisScript = new DefaultRedisScript<>();
    redisScript.setResultType(Number.class);
    ClassPathResource classPathResource = new ClassPathResource(LIMIT_LUA_PATH);
    try {
        classPathResource.getInputStream();//探测资源是否存在
        redisScript.setScriptSource(new ResourceScriptSource(classPathResource));
    } catch (IOException e) {
        logger.error("未找到文件：{}", LIMIT_LUA_PATH);
    }
}

local count
count = redis.call('get',KEYS[1])
--不超过最大值，则直接返回
if count and tonumber(count) > tonumber(ARGV[1]) then
    return count;
end
    --执行计算器自加
    count = redis.call('incr',KEYS[1])
if tonumber(count) == 1 then
    --从第一次调用开始限流，设置对应key的过期时间
    redis.call('expire',KEYS[1],ARGV[2])
end
return count;

/**
 * 降级策略
 * @author 北哥
 * @date 2020/4/8 13:24
 */
private void demote() {
    logger.info("try to access fail, this request will be demoted");
    //throw new RuntimeException("限流了");
    response.setHeader("Content-Type", "text/html;charset=UTF8");
    PrintWriter writer = null;
    try {
        writer = response.getWriter();
        writer.println("访问失败，请稍后再试...");
        writer.flush();
    } catch (Exception e) {
        e.printStackTrace();
    } finally {
        if (writer != null) {
            writer.close();
        }
    }
}

/** * 测试限流controller
 * @author 北哥 公众号：BiggerBoy
 * @date 2022年9月8日17:02:40
 */
@RestController
public class TestLimiterController {
    /**
     * @author 北哥
     * @description
     * @date 2020/4/8 13:42
     */
    @MyRedisLimiter(key = "limitTest", count = 2)
    @RequestMapping(value = "/limitTest")
    public Long limitTest() {
        System.out.println("limitTest");
        return 1L;
    }

    /**
     * @author 北哥
     * @description
     * @date 2020/4/8 13:42
     */
    @MyRedisLimiter(key = "customer_limit_test", period = 10, count = 3, limitType = LimitType.CUSTOMER)
    @GetMapping("/limitTest2")
    public Integer testLimiter2() {
        System.out.println("limitTest2");
        return 1;
    }

    /**
     * @author 北哥
     * @description
     * @date 2020/4/8 13:42
     */
    @MyRedisLimiter(key = "ip_limit_test", period = 10, count = 3, limitType = LimitType.IP)
    @GetMapping("/limitTest3")
    public Integer testLimiter3() {
        System.out.println("limitTest3");
        return 3;
    }
}