基于Spring AOP 和redis的zset开发自定义接口限流注解

引入依赖包

关于Spring AOP的依赖（starter）,提供@Aspect支持。引入SpringBoo的data的starter提供数据库操作的支持。

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

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

注解定义及参数设置

import java.lang.annotation.*;

@Retention(RetentionPolicy.RUNTIME)//
@Target(ElementType.METHOD)//指定注解可以使用在方法上
public @interface RateLimiter {
    long period() default 30; // 默认时间段，秒
    long count() default 3000; // 默认最大访问次数
}

注解实现

按照ip来对访问的用户进行拦截，zSetOperations 是按照当前时间 System.currentTimeMillis() 来作为score，向redis中存入数据，每次拦截都会删除period时间段之前的记录：

 zSetOperations.removeRangeByScore(key, 0, currentMs - period * 1000);

被拦截则不会插入数据，没有被拦截则会插入数据，所以拦截的代码放在插入数据之前。

import com.des.data.annotation.RateLimiter;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import java.util.concurrent.TimeUnit;

@Aspect
@Component
@Slf4j
public class RateLimiterAspect {

    @Resource
    private RedisTemplate redisTemplate;

    @Pointcut("@annotation(rateLimiter)")
    public void controllerAspect(RateLimiter rateLimiter) {}
    
    @Around("controllerAspect(rateLimiter)")
    public Object doArround(ProceedingJoinPoint joinPoint, RateLimiter rateLimiter) throws Throwable {
        //获取注解参数
        long period = rateLimiter.period();
        long count = rateLimiter.count();
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = attributes.getRequest();
        String ip = request.getRemoteAddr();
        String uri = request.getRequestURI();
        String key = "req_rateLimiter_".concat(uri).concat(ip);
        ZSetOperations<String, Long> zSetOperations = redisTemplate.opsForZSet();
        if (zSetOperations == null) {
            return joinPoint.proceed();
        }
        long currentMs = System.currentTimeMillis();
        long countReq = zSetOperations.zCard(key);
        zSetOperations.removeRangeByScore(key, 0, currentMs - period * 1000);
        System.out.println("请求数：" + countReq);
        if (countReq > count) {
            throw new RuntimeException("请求超限");
        }
        zSetOperations.add(key, currentMs, currentMs);
        redisTemplate.expire(key, period, TimeUnit.SECONDS);
        return joinPoint.proceed();
    }

}

使用示例

import java.util.Date;

@RestController
@RequestMapping("data/")
public class RateLimteController {

    @RateLimiter(count = 300,period = 30)
    @PostMapping("invoke")
    public Result invoke(){
        System.out.println(new Date());
        return Result.ok();
    }
}