(SpringBoot)服务端限流功能

首先解释下服务端限流,并不是不接受客户端的消息,只是不处理或者少处理客户端的请求.

本文涉及以下几个知识点,请同学们复习复习哦!本文也不会过多详细介绍

1. aop
2. 自定义注解
3. redis
4. 自定义异常以及全局异常捕获

正式进入主题,首先定义一个枚举变量作为限流类型使用

package com.zhcj.xzjh.config.data;

/**
 * @author Dan
 * @version 1.0
 * @date 2022/6/6 15:22
 * @info
 */
public enum LimitType {
    /**
     * 默认策略全局限流
     */
    DEFAULT,
    /**
     * 根据请求者IP进行限流
     */
    IP
}

接着定义一个注解类

package com.zhcj.xzjh.config.annotation;

import com.zhcj.xzjh.config.data.LimitType;

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;

/**
 * @author Dan
 * @version 1.0
 * @date 2022/6/6 15:23
 * @info
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter {
    /**
     * 限流key
     */
    String key() default "rate_limit:";

    /**
     * 限流时间,单位秒
     */
    int time() default 60;

    /**
     * 限流次数
     */
    int count() default 100;

    /**
     * 限流类型
     */
    LimitType limitType() default LimitType.DEFAULT;
}

然后就再定义Aspect类,限流逻辑处理

package com.zhcj.xzjh.config;

import com.zhcj.xzjh.config.annotation.RateLimiter;
import com.zhcj.xzjh.config.data.LimitType;
import com.zhcj.xzjh.exception.RateLimitException;
import com.zhcj.xzjh.util.IpUtil;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * @author Dan
 * @version 1.0
 * @date 2022/6/6 15:30
 * @info 限流拦截器
 */
@Slf4j
@Aspect
@Component
public class RateLimiterAspect {


    @Autowired
    private RedisTemplate redisTemplate;

    @Autowired
    private RedisScript<Long> limitScript;

    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable {
        String key = rateLimiter.key();
        int time = rateLimiter.time();
        int count = rateLimiter.count();

        String combineKey = getCombineKey(rateLimiter, point);
        List<Object> keys = Collections.singletonList(combineKey);
        Long number = (Long) redisTemplate.execute(limitScript, keys, count, time);
        if (number == null || number.intValue() > count) {
            throw new RateLimitException("访问过于频繁，请稍候再试");
        }
        log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), key);
    }

    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
        StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
        if (rateLimiter.limitType() == LimitType.IP) {
            stringBuffer.append(IpUtil.getIpAddr(((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest())).append("-");
        }
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
        return stringBuffer.toString();
    }
}

在定义一个自定义异常,方便异常捕获处理

package com.zhcj.xzjh.exception;

public class RateLimitException extends RuntimeException {
    public RateLimitException(String message) {
        super(message);
    }
}

这里使用到redis时,还会用到lua脚本,在项目的resource里创建lua保存以下代码

local key = KEYS[1]
local count = tonumber(ARGV[1])
local time = tonumber(ARGV[2])
local current = redis.call('get', key)
if current and tonumber(current) > count then
    return tonumber(current)
end
current = redis.call('incr', key)
if tonumber(current) == 1 then
    redis.call('expire', key, time)
end
return tonumber(current)

代码意思也不复杂,做计数计算.

再写个配置类,注入bean

package com.zhcj.xzjh.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;

/**
 * @author Dan
 * @version 1.0
 * @date 2022/6/6 15:27
 * @info 限流配置
 */
@Configuration
public class MyLimitConfig {
    @Bean
    public DefaultRedisScript<Long> limitScript() {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
        redisScript.setResultType(Long.class);
        return redisScript;
    }
}

全局异常处理

package com.zhcj.xzjh.config;

import com.zhcj.xzjh.common.Constants;
import com.zhcj.xzjh.exception.RateLimitException;
import com.zhcj.xzjh.vo.ResponseBean;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.ResponseStatus;
import org.springframework.web.bind.annotation.RestControllerAdvice;

import javax.servlet.http.HttpServletRequest;

@RestControllerAdvice
@Slf4j
public class GlobalExceptionHandler<T> {

    /**
     * 通用异常处理
     */
    @ExceptionHandler(Exception.class)
    public ResponseBean handleServiceMessageException(HttpServletRequest request, Exception ex) {
        //这里可以通用日志记录下错误,我就不多做处理了
        log.error(ex.toString(), ex);
        return ResponseBean.error(Constants.ERROE);
    }

    /**
     * 用于处理限流
     */
    @ExceptionHandler(RateLimitException.class)
    @ResponseStatus(HttpStatus.BAD_REQUEST)
    public ResponseBean bindException(RateLimitException e) {
        return new ResponseBean(503, "data caps", null);
    }
}

这样就完成所需的业务拉,代码比较多,但都很好理解~

接下来就是实战拉!

随便写一个controller测测,(5秒内超过3次请求触发限流机制)

package com.zhcj.xzjh.controller;

import com.zhcj.xzjh.config.annotation.RateLimiter;
import com.zhcj.xzjh.config.data.LimitType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.Date;

@RestController
public class HelloController {
    @GetMapping("/hello")
    @RateLimiter(time = 5, count = 3, limitType = LimitType.IP)
    public String hello() {
        return "hello>>>" + new Date();
    }
}

测试结果:

 非常好用~~~