如何使用SpringBoot + Redis實現介面限流-Redis-PHP中文網

配置

首先我們創建一個Spring Boot 工程，引入Web 和Redis 依賴，同時考慮到接口限流一般是透過註解來標記，而註解是透過AOP 來解析的，所以我們還需要加上AOP 的依賴，最終的依賴如下：

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

登入後複製

然後提前準備好一個Redis 實例，這裡我們專案配置好之後，直接配置一下Redis 的基本資訊即可，如下：

spring.redis.host=localhost
spring.redis.port=6379
spring.redis.password=123

登入後複製

限流註解

接下來我們建立一個限流註解，我們將限流分為兩種情況：

針對目前介面的全域性極限，例如該介面可以在1 分鐘內存取100 次。
針對某一個 IP 位址的限流，例如某個 IP 位址可以在 1 分鐘內存取 100 次。

針對這兩種情況，我們建立一個列舉類別：

public enum LimitType {
    /**
     * 默认策略全局限流
     */
    DEFAULT,
    /**
     * 根据请求者IP进行限流
     */
    IP
}

登入後複製

接下來我們來建立限流註解：

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter {
    /**
     * 限流key
     */
    String key() default "rate_limit:";

    /**
     * 限流时间,单位秒
     */
    int time() default 60;

    /**
     * 限流次数
     */
    int count() default 100;

    /**
     * 限流类型
     */
    LimitType limitType() default LimitType.DEFAULT;
}

登入後複製

第一個參數限流的key，這個只是一個前綴，將來完整的key 是這個前綴再加上介面方法的完整路徑，共同組成限流key，這個key 將被存入到Redis 中。

另外三個參數好理解，我就不多說了。

好了，將來哪個介面需要限流，就在哪個介面上加入 @RateLimiter 註解，然後設定相關參數即可。

客製化RedisTemplate

在Spring Boot 中，我們其實更習慣使用Spring Data Redis 來操作Redis，不過預設的RedisTemplate 有一個小坑，就是序列化用的是JdkSerializationRedisSerializer，不知道小夥伴們有沒有註意過，直接用這個序列化工具將來存到Redis 上的key 和value 都會莫名其妙多一些前綴，這就導致你用命令讀取的時候可能會出錯。

例如儲存的時候，key 是name，value 是test，但是當你在命令列操作的時候，get name 卻取得不到你想要的數據，原因就是存到redis 之後name 前面多了一些字符，此時只能繼續使用RedisTemplate 將之讀取出來。

我們用 Redis 做限流會用到 Lua 腳本，使用 Lua 腳本的時候，就會出現上面說的這種情況，所以我們需要修改 RedisTemplate 的序列化方案。

可能有小夥伴會說為什麼不用 StringRedisTemplate 呢？ StringRedisTemplate 確實不存在上面所說的問題，但是它能夠儲存的資料類型不夠豐富，所以這裡不考慮。

修改RedisTemplate 序列化方案，程式碼如下：

@Configuration
public class RedisConfig {

    @Bean
    public RedisTemplate<Object, Object> redisTemplate(RedisConnectionFactory connectionFactory) {
        RedisTemplate<Object, Object> redisTemplate = new RedisTemplate<>();
        redisTemplate.setConnectionFactory(connectionFactory);
        // 使用Jackson2JsonRedisSerialize 替换默认序列化(默认采用的是JDK序列化)
        Jackson2JsonRedisSerializer<Object> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer<>(Object.class);
        ObjectMapper om = new ObjectMapper();
        om.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
        om.enableDefaultTyping(ObjectMapper.DefaultTyping.NON_FINAL);
        jackson2JsonRedisSerializer.setObjectMapper(om);
        redisTemplate.setKeySerializer(jackson2JsonRedisSerializer);
        redisTemplate.setValueSerializer(jackson2JsonRedisSerializer);
        redisTemplate.setHashKeySerializer(jackson2JsonRedisSerializer);
        redisTemplate.setHashValueSerializer(jackson2JsonRedisSerializer);
        return redisTemplate;
    }
}

登入後複製

這個其實也沒啥好說的，key 和value 我們都使用Spring Boot 中預設的jackson 序列化方式來解決。

Lua 腳本

這個其實我在之前vhr 那一套影片中講過，Redis 中的一些原子操作我們可以藉助Lua 腳本來實現，想要呼叫Lua 腳本，我們有兩種不同的思路：

在Redis 服務端定義好Lua 腳本，然後計算出來一個雜湊值，在Java 程式碼中，透過這個雜湊值鎖定要執行哪個Lua腳本。
直接在 Java 程式碼中將 Lua 腳本定義好，然後傳送到 Redis 服務端去執行。

Spring Data Redis 中也提供了操作 Lua 腳本的接口，還是比較方便的，所以我們在這裡就採用第二種方案。

我們在 resources 目錄下新建 lua 資料夾專門用來存放 lua 腳本，腳本內容如下：

local key = KEYS[1]
local count = tonumber(ARGV[1])
local time = tonumber(ARGV[2])
local current = redis.call(&#39;get&#39;, key)
if current and tonumber(current) > count then
    return tonumber(current)
end
current = redis.call(&#39;incr&#39;, key)
if tonumber(current) == 1 then
    redis.call(&#39;expire&#39;, key, time)
end
return tonumber(current)

登入後複製

這個腳本其實不難，大概瞅一眼就知道乾啥用的。 KEYS 和ARGV 都是一會呼叫時候傳進來的參數，tonumber 就是把字串轉為數字，redis.call 就是執行具體的redis 指令，具體流程是這樣：

首先取得到傳進來的key 以及限流的count 和時間time。
透過 get 取得到這個 key 對應的值，這個值就是目前時間窗內這個介面可以存取多少次。
如果是第一次訪問，此時拿到的結果為nil，否則拿到的結果應該是數字，所以接下來就判斷，如果拿到的結果是一個數字，而這個數字還大於count，那就表示已經超過流量限制了，那麼直接回傳查詢的結果即可。
如果拿到的結果為 nil，表示是第一次訪問，此時就給目前 key 自增 1，然後設定一個過期時間。
最後把自增 1 後的值回傳就好了。

其實這段 Lua 腳本很好理解。

接下來我們在一個 Bean 中來載入這段 Lua 腳本，如下：

@Bean
public DefaultRedisScript<Long> limitScript() {
    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
    redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
    redisScript.setResultType(Long.class);
    return redisScript;
}

登入後複製

可以啦，我們的 Lua 腳本現在就準備好了。

註解解析

接下來我們就需要自訂切面，來解析這個註解了，我們來看看切面的定義：

@Aspect
@Component
public class RateLimiterAspect {
    private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);
    @Autowired
    private RedisTemplate<Object, Object> redisTemplate;
    @Autowired
    private RedisScript<Long> limitScript;
    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable {
        String key = rateLimiter.key();
        int time = rateLimiter.time();
        int count = rateLimiter.count();

        String combineKey = getCombineKey(rateLimiter, point);
        List<Object> keys = Collections.singletonList(combineKey);
        try {
            Long number = redisTemplate.execute(limitScript, keys, count, time);
            if (number==null || number.intValue() > count) {
                throw new ServiceException("访问过于频繁，请稍候再试");
            }
            log.info("限制请求&#39;{}&#39;,当前请求&#39;{}&#39;,缓存key&#39;{}&#39;", count, number.intValue(), key);
        } catch (ServiceException e) {
            throw e;
        } catch (Exception e) {
            throw new RuntimeException("服务器限流异常，请稍候再试");
        }
    }

    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
        StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
        if (rateLimiter.limitType() == LimitType.IP) {
            stringBuffer.append(IpUtils.getIpAddr(((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest())).append("-");
        }
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
        return stringBuffer.toString();
    }
}

登入後複製

這個切面就是攔截所有加了 @RateLimiter 註解的方法，在前置通知中處理註解。

首先获取到注解中的 key、time 以及 count 三个参数。
获取一个组合的 key，所谓的组合的 key，就是在注解的 key 属性基础上，再加上方法的完整路径，如果是 IP 模式的话，就再加上 IP 地址。以 IP 模式为例，最终生成的 key 类似这样：rate_limit:127.0.0.1-org.javaboy.ratelimiter.controller.HelloController-hello（如果不是 IP 模式，那么生成的 key 中就不包含 IP 地址）。
将生成的 key 放到集合中。
通过 redisTemplate.execute 方法取执行一个 Lua 脚本，第一个参数是脚本所封装的对象，第二个参数是 key，对应了脚本中的 KEYS，后面是可变长度的参数，对应了脚本中的 ARGV。
判断 Lua 脚本执行后的结果是否超过 count，若超过则视为过载，抛出异常处理即可。

接口测试

接下来我们就进行接口的一个简单测试，如下：

@RestController
public class HelloController {
    @GetMapping("/hello")
    @RateLimiter(time = 5,count = 3,limitType = LimitType.IP)
    public String hello() {
        return "hello>>>"+new Date();
    }
}

登入後複製

每一个 IP 地址，在 5 秒内只能访问 3 次。

这个自己手动刷新浏览器都能测试出来。

全局异常处理

由于过载的时候是抛异常出来，所以我们还需要一个全局异常处理器，如下：

@RestControllerAdvice
public class GlobalException {
    @ExceptionHandler(ServiceException.class)
    public Map<String,Object> serviceException(ServiceException e) {
        HashMap<String, Object> map = new HashMap<>();
        map.put("status", 500);
        map.put("message", e.getMessage());
        return map;
    }
}

登入後複製

我将这句话重写成如下：这个 demo 很小，所以我没有定义实体类，而是直接使用 Map 来返回 JSON。最后我们看看过载时的测试效果：

如何使用SpringBoot + Redis實現介面限流