SpringBoot使用自定义注解＋拦截器＋Redis实现高并发接口限流

需求：使用自定义注解＋拦截器和Redis实现高并发接口限流

整体代码结构：就是一个简单的单体项目

依赖

         <!-- Spring Boot Starter Web -->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
        <!-- Spring Boot Starter Data Redis -->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-data-redis</artifactId>
        </dependency>
        <!-- Jedis -->
        <dependency>
            <groupId>redis.clients</groupId>
            <artifactId>jedis</artifactId>
        </dependency>

下面是每个包下的具体代码：

RateLimit：

package com.xxxx.annotation;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    int value(); // 每秒允许的请求数
}

WebConfig：

package com.xxxx.config;

import com.xxxx.interceptor.RateLimitInterceptor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

/**
 * 配置类
 */
@Configuration
public class WebConfig implements WebMvcConfigurer {

    private final RateLimitInterceptor rateLimitInterceptor;

    @Autowired
    public WebConfig(RateLimitInterceptor rateLimitInterceptor) {
        this.rateLimitInterceptor = rateLimitInterceptor;
    }

    /**
     * 注册自定义拦截器
     * @param registry
     */
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(rateLimitInterceptor);
    }
}

TestController：

package com.xxxx.controller;

import com.xxxx.annotation.RateLimit;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
public class TestController {

    @RateLimit(5) // 每秒允许5次请求
    @GetMapping("/testLimit")
    public String test() {
        return "request success";
    }
}

RateLimitInterceptor：

package com.xxxx.interceptor;

import com.xxxx.annotation.RateLimit;
import com.xxxx.service.RedisService;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

/**
 * 自定义拦截器：限流拦截器
 */
@Component
public class RateLimitInterceptor implements HandlerInterceptor {

    private final  RedisService redisService;

    public RateLimitInterceptor(RedisService redisService) {
        this.redisService = redisService;
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws IOException {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            RateLimit rateLimit = handlerMethod.getMethodAnnotation(RateLimit.class);

            if (rateLimit != null) {
                String key = handlerMethod.getBeanType().getName() + ":" + handlerMethod.getMethod().getName();
                int maxRequests = rateLimit.value();

                if (!redisService.isAllowed(key, maxRequests)) {
                    response.setStatus(429); // Too Many Requests
                    response.getWriter().write("Too many requests");
                    return false;
                }
            }
        }
        return true;
    }
}

RedisService：

package com.xxxx.service;

import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

import java.time.Duration;

@Service
public class RedisService {

    private final StringRedisTemplate redisTemplate;

    public RedisService(StringRedisTemplate redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 判断是否超过单秒最大请求数
     * @param key          接口的名称
     * @param maxRequests  单秒最大请求数
     * @return
     */
    public boolean isAllowed(String key, int maxRequests) {

        //获取目前接口的请求次数
        String count = redisTemplate.opsForValue().get(key);

        if (count == null) {
            redisTemplate.opsForValue().set(key, "1", Duration.ofSeconds(1));
            return true;
        } else if (Integer.parseInt(count) < maxRequests) {
            redisTemplate.opsForValue().increment(key);
            return true;
        } else {
            return false;
        }
    }
}

启动类：RateLimitDemoApplication

package com.xxxx;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

@SpringBootApplication
public class RateLimitDemoApplication {

    public static void main(String[] args) {
        SpringApplication.run(RateLimitDemoApplication.class, args);
        System.out.println("============= RateLimitDemoApplication Start ============");
    }
}

配置文件application.yml

spring:
  redis:
    port: 6379
    host: 127.0.0.1
    database: 1
    timeout: 10000

测试：