注解使用redisson完成api限流

1. 注解类

package com.sanrex.annotation;

import org.redisson.api.RateIntervalUnit;

import java.lang.annotation.*;

/**
 * Redisson限流注解
 *
 * @author :chou
 * @date :2022-01-24 10:11
 **/
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface Limiter {

    /**
     * 限流器的 key
     *
     * @return key
     */
    //String key();

    /**
     * 限制数量
     *
     * @return 许可数量
     */
    long limitNum() default 20;

	/**
	 * 速率时间间隔
	 *
	 * @return 速率时间间隔
	 */
	long limitTimeInterval() default 1;

	/**
	 * 时间单位
	 *
	 * @return 时间
	 */
	RateIntervalUnit limitTime() default RateIntervalUnit.MINUTES;
}

3. aop切面类

package com.sanrex.aop;

import com.sanrex.annotation.Limiter;
import com.sanrex.quote.util.IPUtils;
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.aspectj.lang.reflect.MethodSignature;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;

/**
 * Redisson限流注解
 *
 * @author :chou
 * @date :2022-01-24 10:11
 **/
@Component
@Aspect
@Slf4j
public class LimiterAspect {

	private RedissonClient redisson;

	public LimiterAspect(RedissonClient redissonClient) {
		this.redisson = redissonClient;
	}

	@Pointcut("@annotation(com.sanrex.annotation.Limiter)")
	public void limitAspect() {}

	@Around("limitAspect()")
	public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
		Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
		Limiter limiterAnnotation = method.getAnnotation(Limiter.class);
		long limitNum = limiterAnnotation.limitNum();
		long limitTimeInterval = limiterAnnotation.limitTimeInterval();

		ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
		HttpServletRequest request = attributes.getRequest();
		String ip = IPUtils.getIpAddr(request);
		String key = "req:limit:" + ip + "-" + request.getRequestURI();

		RateIntervalUnit rateIntervalUnit = limiterAnnotation.limitTime();
		RRateLimiter limiter = redisson.getRateLimiter(key);
		// trySetRate方法只会执行一次，目前暂时没有提供方法重新设置速率，只能删除后重建 issue#2322
		limiter.trySetRate(RateType.OVERALL, limitNum, limitTimeInterval, rateIntervalUnit);
		boolean allow = limiter.tryAcquire();
		if (!allow) {
			String url = request.getRequestURL().toString();
			String unit = "";
			if (rateIntervalUnit == RateIntervalUnit.SECONDS) {
				unit = "秒";
			} else if (rateIntervalUnit == RateIntervalUnit.MINUTES) {
				unit = "分钟";
			} else if (rateIntervalUnit == RateIntervalUnit.HOURS) {
				unit = "小时";
			} else if (rateIntervalUnit == RateIntervalUnit.DAYS) {
				unit = "天";
			}
			String tooManyRequestMsg = String.format("用户IP[%s]访问地址[%s]时间间隔[%s %s]超过了限定的次数[%s]", ip, url, limitTimeInterval, unit, limitNum);
			log.info(tooManyRequestMsg);
			throw new RuntimeException(tooManyRequestMsg);
		}
		return joinPoint.proceed();
	}
}

3. 实际使用

	@Limiter(limitNum = 5, limitTimeInterval = 3)
	@PostMapping("/product_unit")
	public R apiMqPullData(@RequestBody Map queryMap) {
		Integer pullCount = provider.apiMqPullData(queryMap);
		return R.success(String.format("拉取到数据到MQ数量: %s", pullCount));
	}