令牌桶限流算法

令牌桶算法是一個桶，勻速向桶里放令牌，控制桶最大容量（令牌最大數）和放入令牌速率（生成令牌/秒）。所有的請求在處理之前都需要拿到一個可用的令牌才會被處理，如果桶里面沒有令牌的話，則拒絕服務；

 

• 接口限制 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可用