引入依赖包
关于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();
}
}