coding++：高并发解决方案限流技术-使用RateLimiter实现令牌桶限流-Demo

Posted 2020-11-23 codingmode

RateLimiter是guava提供的基于令牌桶算法的实现类，可以非常简单的完成限流特技，并且根据系统的实际情况来调整生成token的速率。

通常可应用于抢购限流防止冲垮系统；限制某接口、服务单位时间内的访问量，譬如一些第三方服务会对用户访问量进行限制；限制网速，单位时间内只允许上传下载多少字节等。

guava的maven依赖

<dependency>
     <groupId>com.google.guava</groupId>
     <artifactId>guava</artifactId>
     <version>25.1-jre</version>
 </dependency>

令牌桶的原理，有一个独立线程一直以一个固定的速率往桶中存放令牌  客户端去桶中获取令牌，获取到令牌，就可以访问，获取不到，说明请求过多，需要服务降级。

package com.aiyuesheng.controller;

import java.util.concurrent.TimeUnit;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.aiyuesheng.hystrix.OrderHystrixCommand;
import com.aiyuesheng.service.OrderService;
import com.aiyuesheng.utils.LimitService;
import com.alibaba.fastjson.JSONObject;
import com.google.common.util.concurrent.RateLimiter;

@RestController
public class Index {// 令牌桶：1.0 表示 每秒中生成1个令牌存放在桶中
    RateLimiter rateLimiter = RateLimiter.create(1.0);

    @Autowired
    private OrderService orderService;
//令牌桶限流
    @RequestMapping("/searchCustomerInfoByRateLimiter")
    public Object searchCustomerInfoByRateLimiter() {
        // 1.限流判断
        // 如果在0.5秒内 没有获取不到令牌的话，则会一直等待
        System.out.println("生成令牌等待时间:" + rateLimiter.acquire());
        boolean acquire = rateLimiter.tryAcquire(500, TimeUnit.MILLISECONDS); // 每次发送请求，愿意等待0.5秒，如果设为1秒，每次都能查询成功，因为没秒中都会放入一个令牌到桶中
        if (!acquire) {
            System.out.println("稍后再试！");
            return "稍后再试！";
        }
        // 2.如果没有达到限流的要求,直接调用接口查询
        System.out.println(orderService.searchCustomerInfo());
        return orderService.searchCustomerInfo();
    }

}

基于 AOP 实现：

package com.tree.ztree_demo.currentlimiting;

import java.lang.annotation.*;

/**
 * @version V1.0
 * @Title: RateLimit.java
 * @Package mlq.pic.picback.currentlimiting
 * @Description: 限流注解
 * @author: MLQ
 * @date: 2019/11/14 14:52
 */
@Inherited
@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {

    /**
     * 默认每秒支持2个
     *
     * @return
     */
    int limintNum() default 2;


}

-----

package com.tree.ztree_demo.currentlimiting;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.util.concurrent.RateLimiter;
import groovy.util.logging.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @version V1.0
 * @Title: RateLimitAspect.java
 * @Package mlq.pic.picback.currentlimiting
 * @Description: 增强限流拦截
 * @author: MLQ
 * @date: 2019/11/14 14:54
 */
@Component
@Scope
@Aspect
@Slf4j
public class RateLimitAspect {

    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimitAspect.class);

    /**
     * 用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
     */
    private ConcurrentHashMap<String, RateLimiter> map = new ConcurrentHashMap<>();

    private static final String POINT = "execution (* com.tree.ztree_demo..*.abc*(..))";

    private static ObjectMapper objectMapper = new ObjectMapper();

    private RateLimiter rateLimiter;

    @Autowired
    private HttpServletResponse response;

    @Pointcut(POINT)
    public void serviceLimit() {
    }

    @Around("serviceLimit()")
    public Object around(ProceedingJoinPoint joinPoint) throws NoSuchMethodException {
        Object obj = null;
        //获取拦截的方法名
        Signature sig = joinPoint.getSignature();
        //获取拦截的方法名
        MethodSignature msig = (MethodSignature) sig;
        //返回被织入增加处理目标对象
        Object target = joinPoint.getTarget();
        //为了获取注解信息
        Method currentMethod = target.getClass().getMethod(msig.getName(), msig.getParameterTypes());
        //获取注解信息
        RateLimit annotation = currentMethod.getAnnotation(RateLimit.class);
        if (!ObjectUtils.isEmpty(annotation)) {
            //获取注解每秒加入桶中的token
            int limitNum = annotation.limintNum();

            // 注解所在方法名区分不同的限流策略
            String functionName = msig.getName();

            //获取rateLimiter
            if (map.containsKey(functionName)) {
                rateLimiter = map.get(functionName);
            } else {
                map.put(functionName, RateLimiter.create(limitNum));
                rateLimiter = map.get(functionName);
            }
            // 如果在0.5秒内 没有获取不到令牌的话，则会一直等待
            System.out.println("生成令牌等待时间:" + rateLimiter.acquire());
            try {
                // 每次发送请求，愿意等待0.5秒，如果设为1秒，每次都能查询成功，因为没秒中都会放入一个令牌到桶中
                if (rateLimiter.tryAcquire(500, TimeUnit.MILLISECONDS)) {
                    //执行方法
                    obj = joinPoint.proceed();
                } else {
                    Map<String, Object> map = new HashMap<>();
                    map.put("code", 100001);
                    map.put("message", "系统繁忙，请稍后再试！");
                    //拒绝了请求（服务降级）
                    String result = objectMapper.writeValueAsString(map);
                    LOGGER.info("拒绝了请求：" + result);
                    outErrorResult(result);
                }
            } catch (Throwable throwable) {
                throwable.printStackTrace();
            }
        } else {
            Map<String, Object> map = new HashMap<>();
            map.put("code", 403);
            map.put("message", "签名问题");
            //拒绝了请求（服务降级）
            String result = null;
            try {
                result = objectMapper.writeValueAsString(map);
            } catch (JsonProcessingException e) {
                e.printStackTrace();
            }
            LOGGER.info("拒绝了请求：" + result);
            outErrorResult(result);
        }
        return obj;
    }

    //将结果返回
    public void outErrorResult(String result) {
        response.setContentType("application/json;charset=UTF-8");
        try (ServletOutputStream outputStream = response.getOutputStream()) {
            outputStream.write(result.getBytes("utf-8"));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    static {
        objectMapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
    }

}

使用方式：在controller 访问地址加上注解即可

@RateLimit(limintNum = 100)
@RequestMapping("abc")
public Object abc() {
    Map<String, Object> map = new HashMap<>();
    map.put("code", "200");
    map.put("message", "OK");
    return JSON.toJSONString(map);
}