限流算法之漏桶算法与令牌桶算法

背景：

      在是我们实际的生产场景中，我们需要对大到整个服务网关，细到某个接口方法做流量控制，限流，解决并发请求负载过高，导致大到可能压垮服务，导致某节点不可用，所以我们需要限流，通常做法是请求超出预算后排队，或者直接拒绝请求.

漏桶算法

      漏桶算法基本思路就是当请求到一个缓冲区，然后从缓冲区流出再去处理请求，缓冲区(桶)的大小，从缓冲区(请求响应速度)流出的速度是关键。其基本可以用一个缓冲队列来实现。

 

 

上面为简单的漏桶流程图

 

令牌桶算法

      令牌桶算法是是另外一种限流算法，顾名思义，当请求需要从桶中拿到令牌才能进行接口的访问，其中大体思路是，一个令牌队列(桶)，一个定时生成令牌放入桶中，请求过来需要经过这个桶，当从队列取得令牌后才能进行后续处理，否则排队，也可以丢弃请求，令牌的生成速度影响请求响应效率。

google的一个工具包中有一个RateLimiter类，它是令牌桶算法的一种实现，下面贴出它的依赖，以及使用场景

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

 令牌桶场景：假设我针对某个方法使用RateLimiter实现限流控制，该方法每秒只能一个一个请求进去

import java.util.UUID;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import com.space.rabbitmq.common.annotation.TokenBucketAnnotation;


@RestController
public class TokenBucketController {
    @TokenBucketAnnotation(description = "测试令牌桶")
    @GetMapping("/token")
    public String token(String message){
        String uuid = UUID.randomUUID().toString();
        return uuid;
    }
}

这个一个普通的controller， 我们针对真个controller限流，每秒只能一次访问，下面我门需要一个令牌桶注解，方便我们实现切面

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;

@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD, ElementType.TYPE})
public @interface TokenBucketAnnotation {
	String description() default "";
}

下面需要实现一个切面在调用之前判断是否有足够的令牌，如果令牌不足抛出拒绝请求

import com.google.common.util.concurrent.RateLimiter;
import com.space.rabbitmq.common.annotation.TokenBucketAnnotation;
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.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;
import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Aspect
public class TokenBucketAspect {
	
	private Map<Method, RateLimiter> tokenBucket = new ConcurrentHashMap<>();

    @Pointcut("@annotation(com.space.rabbitmq.common.annotation.TokenBucketAnnotation)")
    public void tokenBucketManager() {
    }
	
    @Around("tokenBucketManager()")
    public Object doAround(ProceedingJoinPoint joinPoint) throws Throwable {

        Class<?> aClass = joinPoint.getTarget().getClass();
        Signature signature = joinPoint.getSignature();
        String targetName = aClass.getName();
        String methodName = signature.getName();
        Class<?> targetClass = Class.forName(targetName);
        Method[] methods = targetClass.getMethods();
        String desc;
        Method method = Arrays.stream(methods).filter(m -> m.getName().equals(methodName)).findFirst().get();
        TokenBucketAnnotation annotation = method.getAnnotation(TokenBucketAnnotation.class);
        desc = annotation.description();
        System.out.println(desc);
        //log.....
        acquire(method);
        Object ret = null;
        try{
            ret = joinPoint.proceed();
        }catch (Throwable e){
        	//log......
        	//为了测试方便  可以自己自定义令牌不够丢弃请求的异常，这里为了测试方便就不写了
        	throw new Throwable(e.getMessage());
        }
        return ret;
    }
    
    private void acquire(Method method) throws Throwable{
    	RateLimiter limiter = tokenBucket.get(method);
    	if(null == limiter){
    		//假设这个方法最大并发数为1， 每秒生成一个令牌
    		limiter = RateLimiter.create(1);
    		tokenBucket.put(method, limiter);
    	}
    	//如果请求超过1，令牌不够 需要排队
    	//limiter.acquire();
    	//令牌不够的情况直接丢弃请求   tryAcquire方法会尝试获取令牌，如果令牌不足放回false
    	if(!limiter.tryAcquire()){
    		throw new Throwable("令牌不足");
    	}
    }

    @Before("tokenBucketManager()")
    public void before(){

    }
}

这样我们就可以针对某个方法实现限流，在浏览器中迅速刷新请求可以看出效果