令牌桶限流算法

令牌桶限流算法

令牌桶算法是一个桶，匀速向桶里放令牌，控制桶最大容量（令牌最大数）和放入令牌速率（生成令牌/秒）。所有的请求在处理之前都需要拿到一个可用的令牌才会被处理，如果桶里面没有令牌的话，则拒绝服务；

 

• 接口限制 t 秒内最大访问次数为 n，则每隔 t/n 秒会放一个 token 到桶中；
• 桶中最多可以存放 b 个 token，如果 token 到达时令牌桶已经满了，那么这个 token 会被丢弃；
• 接口请求会先从令牌桶中取 token，拿到 token 则处理接口请求，拿不到 token 则执行限流；
• 当一个n个字节的数据包到达时，就从令牌桶中删除n个令牌（不同大小的数据包，消耗的令牌数量不一样），并且数据包被发送到网络；

Rate limiting Spring Boot with bucket4j

 每分钟10个限流，每分钟10个速度放入令牌token

Refill refill = Refill.intervally(10, Duration.ofMinutes(1));
Bandwidth limit = Bandwidth.classic(10, refill);
Bucket bucket = Bucket4j.builder()
    .addLimit(limit)
    .build();
 
for (int i = 1; i <= 10; i++) {
    assertTrue(bucket.tryConsume(1));
}
assertFalse(bucket.tryConsume(1));

 

Use Spring MVC Interceptor

public class RateLimitInterceptor implements HandlerInterceptor {

    @Autowired
    private PricingPlanService pricingPlanService;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
            throws Exception {
        String apiKey = request.getHeader("X-api-key");
        if (apiKey == null || apiKey.isEmpty()) {
            response.sendError(HttpStatus.BAD_REQUEST.value(), "Missing Header: X-api-key");
            return false;
        }

        String url = request.getRequestURI();

        Bucket tokenBucket = pricingPlanService.resolveBucket(apiKey+"-"+url);
        ConsumptionProbe probe = tokenBucket.tryConsumeAndReturnRemaining(1);
        if (probe.isConsumed()) {
            response.addHeader("X-Rate-Limit-Remaining", String.valueOf(probe.getRemainingTokens()));
            return true;
        } else {
            long waitForRefill = probe.getNanosToWaitForRefill() / 1_000_000_000;
            response.addHeader("X-Rate-Limit-Retry-After-Seconds", String.valueOf(waitForRefill));
            response.sendError(HttpStatus.TOO_MANY_REQUESTS.value(),
                    "You have exhausted your API Request Quota");
            return false;
        }
    }
}

　　

public class PricingPlanService {

    private final Map<String, Bucket> cache = new ConcurrentHashMap<>();

    public Bucket resolveBucket(String apiKey) {
        return cache.computeIfAbsent(apiKey, this::newBucket);
    }

    private Bucket newBucket(String apiKey) {
        PricingPlan pricingPlan = PricingPlan.resolvePlanFromApiKey(apiKey);
        return Bucket4j.builder()
                .addLimit(pricingPlan.getLimit())
                .build();
    }
}

 

public enum PricingPlan implements BandwidthFactory {

    FREE(){
        @Override
        public Bandwidth getLimit() {
            return Bandwidth.classic(20, Refill.intervally(20, Duration.ofHours(1)));
        }
    },
    BASIC(){
        @Override
        public Bandwidth getLimit() {
            return Bandwidth.classic(40, Refill.intervally(40, Duration.ofHours(1)));
        }
    },
    PROFESSIONAL(){
        @Override
        public Bandwidth getLimit() {
            return Bandwidth.classic(100, Refill.intervally(100, Duration.ofHours(1)));
        }
    };
}

　　

@SpringBootConfiguration
public class AppConfig implements WebMvcConfigurer {

    @Autowired
    private RateLimitInterceptor interceptor;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(interceptor)
                .addPathPatterns("/api/v1/area/**");
    }
}

 

    @PostMapping("/api/v1/area/rectangle3")
    public ResponseEntity<AreaV1> rectangle3(@RequestHeader(value = "X-api-key") String apiKey,
                                            @RequestBody RectangleDimensionsV1 dimensions) {
            return ResponseEntity.ok()
                    .body(new AreaV1("rectangle", dimensions.getLength() * dimensions.getWidth()));
    }

　　

发起测试请求

curl -v -X POST http://localhost:8071/api/v1/area/rectangle3 \
    -H "Content-Type: application/json" -H "X-api-key:FX001-99999" \
    -d '{ "length": 10, "width": 12 }'

　　

每小时20个token，剩余19个token可用