SpringBoot如何使用自定义注解实现接口限流

目录

• 使用自定义注解实现接口限流
• 1、自定义限流注解
• 2、限流类型枚举类
• 3、限流 Lua 脚本
• 4、限流切面处理类
• 5、使用与测试
• SpringBoot工程中限流方式
• 1、google的guava,令牌桶算法实现限流
• 2、interceptor+redis根据注解限流

使用自定义注解实现接口限流

在高并发系统中，保护系统的三种方式分别为：缓存，降级和限流。

限流的目的是通过对并发访问请求进行限速或者一个时间窗口内的的请求数量进行限速来保护系统，一旦达到限制速率则可以拒绝服务、排队或等待。

1、自定义限流注解

import com.asurplus.common.enums.LimitType;
import java.lang.annotation.*;
/**
 * 限流注解
 *
 * @author Lizhou
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface Limit {
    /**
     * 限流key前缀
     */
    String prefix() default "limit:";
    /**
     * 限流时间,单位秒
     */
    int time() default 60;
    /**
     * 限流次数
     */
    int count() default 10;
    /**
     * 限流类型
     */
    LimitType type() default LimitType.DEFAULT;
}

2、限流类型枚举类

/**
 * 限流类型
 *
 * @author Lizhou
 */
public enum LimitType {
    /**
     * 默认策略全局限流
     */
    DEFAULT,
    /**
     * 根据请求者IP进行限流
     */
    IP
}

我们定义了两种限流类型，分别为全局限流和 IP 限流，全局限流对访问接口的所有用户进行限流保护，IP 限流对每个 IP 请求用户进行单独限流保护。

3、限流 Lua 脚本

1、由于我们使用 Redis 进行限流，我们需要引入 Redis 的 maven 依赖，同时需要引入 aop 的依赖

<!-- aop依赖 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
</dependency>
<!-- redis依赖 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

在配置文件中配置 Redis 的连接信息，具体参考：SpringBoot中整合Redis

2、限流脚本

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import java.util.Arrays;
import java.util.List;
/**
 * 接口限流
 */
@Slf4j
@Component
public class RedisLimitUtil {
    @Autowired
    private StringRedisTemplate redisTemplate;
    /**
     * 限流
     *
     * @param key   键
     * @param count 限流次数
     * @param times 限流时间
     * @return
     */
    public boolean limit(String key, int count, int times) {
        try {
            String script = "local lockKey = KEYS[1]\n" +
                    "local lockCount = KEYS[2]\n" +
                    "local lockExpire = KEYS[3]\n" +
                    "local currentCount = tonumber(redis.call('get', lockKey) or \"0\")\n" +
                    "if currentCount < tonumber(lockCount)\n" +
                    "then\n" +
                    "    redis.call(\"INCRBY\", lockKey, \"1\")\n" +
                    "    redis.call(\"expire\", lockKey, lockExpire)\n" +
                    "    return true\n" +
                    "else\n" +
                    "    return false\n" +
                    "end";
            RedisScript<Boolean> redisScript = new DefaultRedisScript<>(script, Boolean.class);
            List<String> keys = Arrays.asList(key, String.valueOf(count), String.valueOf(times));
            return redisTemplate.execute(redisScript, keys);
        } catch (Exception e) {
            log.error("限流脚本执行失败：{}", e.getMessage());
        }
        return false;
    }
}

通过 Lua 脚本，根据 Redis 中缓存的键值判断限流时间（也是 key 的过期时间）内，访问次数是否超出了限流次数，没超出则访问次数 +1，返回 true，超出了则返回 false。

4、限流切面处理类

import com.asurplus.common.annotation.Limit;
import com.asurplus.common.enums.LimitType;
import com.asurplus.common.exception.CustomException;
import com.asurplus.common.ip.IpUtil;
import com.asurplus.common.redis.RedisLimitUtil;
import com.asurplus.common.utils.HttpRequestUtil;
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.stereotype.Component;
import java.lang.reflect.Method;
/**
 * 限流处理
 *
 * @author Lizhou
 */
@Slf4j
@Aspect
@Component
public class LimitAspect {
    @Autowired
    private RedisLimitUtil redisLimitUtil;
    /**
     * 前置通知，判断是否超出限流次数
     *
     * @param point
     */
    @Before("@annotation(limit)")
    public void doBefore(JoinPoint point, Limit limit) {
        try {
            // 拼接key
            String key = getCombineKey(limit, point);
            // 判断是否超出限流次数
            if (!redisLimitUtil.limit(key, limit.count(), limit.time())) {
                throw new CustomException("访问过于频繁，请稍候再试");
            }
        } catch (CustomException e) {
            throw e;
        } catch (Exception e) {
            throw new RuntimeException("接口限流异常，请稍候再试");
        }
    }
    /**
     * 根据限流类型拼接key
     */
    public String getCombineKey(Limit limit, JoinPoint point) {
        StringBuilder sb = new StringBuilder(limit.prefix());
        // 按照IP限流
        if (limit.type() == LimitType.IP) {
            sb.append(IpUtil.getIpAddr(HttpRequestUtil.getRequest())).append("-");
        }
        // 拼接类名和方法名
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        sb.append(targetClass.getName()).append("-").append(method.getName());
        return sb.toString();
    }
}

1、使用我们刚刚的 Lua 脚本判断是否超出了限流次数，超出了限流次数后返回一个自定义异常，然后在全局异常中去捕捉异常，返回 JSON 数据。

2、根据注解参数，判断限流类型，拼接缓存 key 值

5、使用与测试

1、测试方法

@Limit(type = LimitType.DEFAULT, time = 10, count = 2)
@GetMapping("test")
public String test() {
    return "请求成功：" + System.currentTimeMillis();
}

使用自定义注解 @Limit，限制为 10 秒内，允许访问 2 次

2、测试结果

第一次

第二次

第三次

可以看出，前面两次都成功返回了请求结果，第三次超出了接口限流次数，返回了自定义异常信息。

SpringBoot工程中限流方式

限流，是防止用户恶意刷新接口。常见的限流方式有阿里开源的sentinel、redis等。

1、google的guava,令牌桶算法实现限流

Guava的 RateLimiter提供了令牌桶算法实现：平滑突发限流(SmoothBursty)和平滑预热限流(SmoothWarmingUp)实现。

// RateLimiter提供了两个工厂方法，最终会调用下面两个函数，生成RateLimiter的两个子类。
static RateLimiter create(SleepingStopwatch stopwatch, double permitsPerSecond) {
     RateLimiter rateLimiter = new SmoothBursty(stopwatch, 1.0 /* maxBurstSeconds */);
    rateLimiter.setRate(permitsPerSecond);
     return rateLimiter;
}
static RateLimiter create(
     SleepingStopwatch stopwatch, double permitsPerSecond, long warmupPeriod, TimeUnit unit,
 double coldFactor) {
     RateLimiter rateLimiter = new SmoothWarmingUp(stopwatch, warmupPeriod, unit, coldFactor);
    rateLimiter.setRate(permitsPerSecond);
     return rateLimiter;
}

• 平滑突发限流：使用 RateLimiter的静态方法创建一个限流器，设置每秒放置的令牌数为10个。返回的RateLimiter对象可以保证1秒内不会给超过10个令牌，并且以固定速率进行放置，达到平滑输出的效果。
• 平滑预热限流：RateLimiter的 SmoothWarmingUp是带有预热期的平滑限流，它启动后会有一段预热期，逐步将分发频率提升到配置的速率。

@RestController
public class HomeController {
    // 这里的10表示每秒允许处理的量为10个
    private RateLimiter limiter = RateLimiter.create(10);
    private RateLimiter limiter2 = RateLimiter.create(2, 1000, TimeUnit.SECONDS);
    //permitsPerSecond: 表示 每秒新增 的令牌数;warmupPeriod: 表示在从 冷启动速率 过渡到 平均速率 的时间间隔
    @GetMapping("/test/{name}")
    public String test(@PathVariable("name") String name) {
        // 请求RateLimiter, 超过permits会被阻塞
        final double acquire = limiter.acquire();
        System.out.println("acquire=" + acquire);
        //判断double是否为空或者为0
        if (acquire == 0) {
            return name;
        } else {
            return "操作太频繁";
        }
    }
    @AccessLimit(limit = 2, sec = 10)
    @GetMapping("/test2/{name}")
    public String test2(@PathVariable("name") String name) {
        return name;
    }
}

2、interceptor+redis根据注解限流

public class AccessLimitInterceptor implements HandlerInterceptor {
    @Resource
    private RedisTemplate<String, Integer> redisTemplate;
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = handlerMethod.getMethod();
            if (!method.isAnnotationPresent(AccessLimit.class)) {
                return true;
            }
            AccessLimit accessLimit = method.getAnnotation(AccessLimit.class);
            if (accessLimit == null) {
                return true;
            }
            int limit = accessLimit.limit();
            int sec = accessLimit.sec();
            String key = IPUtil.getIpAddress(request) + request.getRequestURI();
            //资源唯一标识
            Integer maxLimit = redisTemplate.opsForValue().get(key);
            if (maxLimit == null) {
                //set时一定要加过期时间
                redisTemplate.opsForValue().set(key, 1, sec, TimeUnit.SECONDS);
            } else if (maxLimit < limit) {
                redisTemplate.opsForValue().set(key, maxLimit + 1, sec, TimeUnit.SECONDS);
            } else {
                output(response, "请求太频繁!");
                return false;
            }
        }
        return true;
    }
    public void output(HttpServletResponse response, String msg) throws IOException {
        response.setContentType("application/json;charset=UTF-8");
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(msg.getBytes("UTF-8"));
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            outputStream.flush();
            outputStream.close();
        }
    }
    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
    }
    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
    }
}

@Configuration
public class InterceptorConfig extends WebMvcConfigurationSupport {
    @Bean
    public AccessLimitInterceptor accessLimitInterceptor() {
        return new AccessLimitInterceptor();
    }
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        //addPathPatterns 添加拦截规则
        registry.addInterceptor(accessLimitInterceptor()).addPathPatterns("/**");
    }
    @Override
    public void addViewControllers(ViewControllerRegistry registry) {
        registry.addViewController("/");
    }
}

@Configuration
public class RedisConfig {
    @Bean
    public RedisTemplate<Object, Object> redisTemplate(RedisConnectionFactory redisConnectionFactory) {
        RedisTemplate<Object, Object> template = new RedisTemplate<Object, Object>();
        template.setConnectionFactory(redisConnectionFactory);
        template.setKeySerializer(new StringRedisSerializer());
        template.setValueSerializer(new GenericJackson2JsonRedisSerializer());
        template.setHashKeySerializer(new GenericJackson2JsonRedisSerializer());
        template.setHashValueSerializer(new GenericJackson2JsonRedisSerializer());
        template.afterPropertiesSet();
        return template;
    }
}

限流方式还有很多，后续继续尝试。

以上为个人经验，希望能给大家一个参考，也希望大家多多支持我们。