Redis实现分布式限流

我们做一个简单的封装，把限流器定义成一个注解，然后定义2个属性，时间和次数，这也是计数器的2个核心属性

import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;

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

/**
 * 通用的方法限流
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Documented
@Order(Ordered.HIGHEST_PRECEDENCE)
public @interface SmartRateLimit {

    /**
     * 允许访问的次数，默认值MAX_VALUE
     */
    long limitCount() default Long.MAX_VALUE;

    /**
     * 时间段，单位秒,默认每秒限流大小
     */
    long timeRange() default 1;
}

注解的实现

import com.smart.server.ratelimit.RequestLimitException;
import com.smart.server.ratelimit.annotation.SmartRateLimit;

import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.aop.support.AopUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.EnableAspectJAutoProxy;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

import lombok.extern.slf4j.Slf4j;

/**
 * 通用方法限流实现
 **/
@Aspect
@Component
@EnableAspectJAutoProxy(proxyTargetClass = true)
@Slf4j
public class SmartRateLimitInterceptor {

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    @Pointcut("@annotation(com.smart.server.ratelimit.annotation.SmartRateLimit)")
    private void limit() {
    }

    @Before("limit()")
    public void before(JoinPoint joinPoint) throws RequestLimitException {

        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        Method targetMethod = AopUtils.getMostSpecificMethod(methodSignature.getMethod(), joinPoint.getTarget().getClass());
        String targetName = joinPoint.getTarget().getClass().getName();
        String methodName = joinPoint.getSignature().getName();

        SmartRateLimit smartRateLimit = targetMethod.getAnnotation(SmartRateLimit.class);
        long limitCount = smartRateLimit.limitCount();
        long timeRange = smartRateLimit.timeRange();
        String redisKey = getRedisKey(targetName, methodName);

        long count = redisTemplate.opsForValue().increment(redisKey, 1);
        log.info("count:" + count);
        if (count == 1) {
            redisTemplate.expire(redisKey, timeRange, TimeUnit.SECONDS);
        }
        if (count > limitCount) {
            log.error("访问方法" + redisKey + "]超过了限定的次数[" + limitCount + "]");
            throw new RequestLimitException("请求超出限制：" + redisKey + ", 当前次数：" + count);
        }

    }

    /**
     * 获取缓存的key值
     */
    private String getRedisKey(String targetName, String methodName) {
        StringBuilder sb = new StringBuilder("");
        sb.append("limitrate.").append(targetName).append(".").append(methodName);
        return sb.toString();
    }
}

利用Redis的increment自增操作记录单位时间内的请求数，Redis的单线程和increment的原子操作能够做到多线程下计数的准确性

当超过我们设计的阈值时，我们抛出一个超限异常，然后使用全局异常进行拦截处理

public class RequestLimitException extends Exception {
    private static final long serialVersionUID = 1364225358754654702L;

    public RequestLimitException() {
        super("请求超出设定的限制");
    }

    public RequestLimitException(String message) {
        super(message);
    }

}

import com.smart.server.base.BaseJsonResult;
import com.smart.server.ratelimit.RequestLimitException;
import com.smart.service.base.BusinessErrorMsg;
import com.smart.service.base.BusinessException;

import org.springframework.http.HttpStatus;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.bind.annotation.RestControllerAdvice;

import java.util.Arrays;
import java.util.Map;

import javax.servlet.http.HttpServletRequest;

import lombok.extern.slf4j.Slf4j;

/**
 * 全局异常处理
 */
@RestControllerAdvice
@Slf4j
public class GlobalExceptionHandler {
    /**
     * 限流异常
     */
    @ExceptionHandler(value = RequestLimitException.class)
    @ResponseBody
    public BaseJsonResult<Object> limitExceptionHandler(HttpServletRequest req, Exception e) throws Exception {
        printMethodParameters(req);
        BaseJsonResult<Object> baseJsonResult = new BaseJsonResult<>();
        baseJsonResult.setStatus(HttpStatus.INTERNAL_SERVER_ERROR.value());
        baseJsonResult.setMsg("系统繁忙,请稍后重试！");
        return baseJsonResult;
    }

}

至此，限流的处理已经完成，我们只需要在方法上加上注解即可

@RequestMapping(value = "ratelimit", method = RequestMethod.GET)
@SmartRateLimit(limitCount = 5)
public BaseJsonResult rateLimit() throws Exception {
    log.info("*********************************");
    return successNullDataResult();
}