io.github.bucket4j.Bucket Java Examples

@Test
public void givenBucketLimit_whenThrottleRequests_thenConsumeReturnsTrue() throws InterruptedException {
    Refill refill = Refill.intervally(1, Duration.ofSeconds(2));
    Bandwidth limit = Bandwidth.classic(1, refill);
    Bucket bucket = Bucket4j.builder()
        .addLimit(limit)
        .build();

    assertTrue(bucket.tryConsume(1));

    ScheduledExecutorService executor = Executors.newScheduledThreadPool(1);
    CountDownLatch latch = new CountDownLatch(1);

    executor.schedule(new AssertTryConsume(bucket, latch), 2, TimeUnit.SECONDS);

    latch.await();
}
 

@Test
public void givenMultipletLimits_whenExceedSmallerLimit_thenConsumeReturnsFalse() {
    Bucket bucket = Bucket4j.builder()
        .addLimit(Bandwidth.classic(10, Refill.intervally(10, Duration.ofMinutes(1))))
        .addLimit(Bandwidth.classic(5, Refill.intervally(5, Duration.ofSeconds(20))))
        .build();

    for (int i = 1; i <= 5; i++) {
        assertTrue(bucket.tryConsume(1));
    }
    assertFalse(bucket.tryConsume(1));
}
 

@Test
public void givenBucketLimit_whenExceedLimit_thenConsumeReturnsFalse() {
    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));
}
 

@Override
public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
    String ip = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
    Bucket bucket = CACHE.computeIfAbsent(ip, k -> createNewBucket());

    log.debug("IP: " + ip + "，TokenBucket Available Tokens: " + bucket.getAvailableTokens());
    if (bucket.tryConsume(1)) {
        return chain.filter(exchange);
    } else {
        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        return exchange.getResponse().setComplete();
    }
}
 

@Override
public Mono<Response> isAllowed(String routeId, String id) {
    Config routeConfig = getConfig().get(routeId);
    if (routeConfig == null) {
        throw new IllegalArgumentException("No Configuration found for route " + routeId);
    }

    // How many requests per second do you want a user to be allowed to do?
    int replenishRate = routeConfig.getReplenishRate();
    // How much bursting do you want to allow?
    int burstCapacity = routeConfig.getBurstCapacity();

    Bucket bucket = this.buckets.computeIfAbsent(id, k -> {
        Refill refill = Refill.greedy(replenishRate, Duration.ofSeconds(1));
        Bandwidth limit = Bandwidth.classic(burstCapacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    });

    // tryConsume returns false immediately if no tokens available with the bucket
    ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(10);
    if (probe.isConsumed()) {
        // the limit is not exceeded
        return Mono.just(new Response(true, Collections.emptyMap()));
    } else {
        log.warn("request rate limited");
        // limit is exceeded
        return Mono.just(new Response(false, getHeaders(routeConfig, 0)));
    }
}
 

private void calcAndSetRemainingBucket(Long consume, Rate rate, Bucket bucket, boolean isQuota) {
    ConsumptionProbe consumptionProbe = bucket.tryConsumeAndReturnRemaining(consume);
    long nanosToWaitForRefill = consumptionProbe.getNanosToWaitForRefill();
    rate.setReset(NANOSECONDS.toMillis(nanosToWaitForRefill));
    if (consumptionProbe.isConsumed()) {
        long remainingTokens = consumptionProbe.getRemainingTokens();
        setRemaining(rate, remainingTokens, isQuota);
    } else {
        setRemaining(rate, -1L, isQuota);
        bucket.tryConsumeAsMuchAsPossible(consume);
    }
}
 

@Override
protected void calcRemainingLimit(final Long limit, final Duration refreshInterval, final Long requestTime,
                                  final String key, final Rate rate) {
    if (limit == null) {
        return;
    }
    Bucket bucket = getLimitBucket(key, limit, refreshInterval);
    if (requestTime == null) {
        calcAndSetRemainingBucket(1L, rate, bucket, false);
    } else {
        calcAndSetRemainingBucket(bucket, rate, false);
    }
}
 

@Override
protected void calcRemainingQuota(final Long quota, final  Duration refreshInterval, final Long requestTime,
                                  final String key, final Rate rate) {
    if (quota == null) {
        return;
    }
    Bucket bucket = getQuotaBucket(key, quota, refreshInterval);
    if (requestTime != null) {
        calcAndSetRemainingBucket(requestTime, rate, bucket, true);
    } else {
        calcAndSetRemainingBucket(bucket, rate, true);
    }
}
 

@Override
public Bucket getProxy(K key, Supplier<BucketConfiguration> supplier) {
    return GridBucket.createLazyBucket(key, supplier, gridProxy);
}
 

public ConsumerThread(CountDownLatch startLatch, CountDownLatch endLatch, Bucket bucket, long workTimeNanos, Function<Bucket, Long> action) {
    this.startLatch = startLatch;
    this.endLatch = endLatch;
    this.bucket = bucket;
    this.workTimeNanos = workTimeNanos;
    this.action = action;
}
 

public ConsumptionScenario(int threadCount, long workTimeNanos, Supplier<Bucket> bucketSupplier, Function<Bucket, Long> action, double permittedRatePerSecond) {
    this.startLatch = new CountDownLatch(threadCount);
    this.endLatch = new CountDownLatch(threadCount);
    this.consumers = new ConsumerThread[threadCount];
    this.initializationTimeMillis = System.currentTimeMillis();
    this.permittedRatePerSecond = permittedRatePerSecond;
    Bucket bucket = bucketSupplier.get();
    for (int i = 0; i < threadCount; i++) {
        this.consumers[i] = new ConsumerThread(startLatch, endLatch, bucket, workTimeNanos, action);
    }
}
 

@Override
public Optional<Bucket> getProxy(K key) {
    return getProxyConfiguration(key)
            .map(configuration -> GridBucket.createLazyBucket(key, () -> configuration, gridProxy));
}
 

@Override
protected Bucket build(HazelcastBucketBuilder builder, String key, RecoveryStrategy recoveryStrategy) {
    return builder.build(map, key, recoveryStrategy);
}
 

@Override
public Bucket getProxy(K key, Supplier<BucketConfiguration> supplier) {
    return GridBucket.createLazyBucket(key, supplier, gridProxy);
}
 

@Override
protected Bucket build(InfinispanBucketBuilder builder, String key, RecoveryStrategy recoveryStrategy) {
    return builder.build(readWriteMap, key, recoveryStrategy);
}
 

@Override
public Optional<Bucket> getProxy(K key) {
    return getProxyConfiguration(key)
            .map(configuration -> GridBucket.createLazyBucket(key, () -> configuration, gridProxy));
}
 

@Override
public Bucket getProxy(K key, Supplier<BucketConfiguration> supplier) {
    return GridBucket.createLazyBucket(key, supplier, gridProxy);
}
 

@Override
protected Bucket build(InfinispanBucketBuilder builder, String key, RecoveryStrategy recoveryStrategy) {
    return builder.build(readWriteMap, key, recoveryStrategy);
}
 

@Override
public Optional<Bucket> getProxy(K key) {
    return getProxyConfiguration(key)
            .map(configuration -> GridBucket.createLazyBucket(key, () -> configuration, gridProxy));
}
 

@Override
public Bucket getProxy(K key, Supplier<BucketConfiguration> supplier) {
    return GridBucket.createLazyBucket(key, supplier, gridProxy);
}
 

@Override
public Optional<Bucket> getProxy(K key) {
    return getProxyConfiguration(key)
            .map(configuration -> GridBucket.createLazyBucket(key, () -> configuration, gridProxy));
}
 

@Override
public Bucket getProxy(K key, Supplier<BucketConfiguration> supplier) {
    return GridBucket.createLazyBucket(key, supplier, gridProxy);
}
 

@Override
public Optional<Bucket> getProxy(K key) {
    return getProxyConfiguration(key)
            .map(configuration -> GridBucket.createLazyBucket(key, () -> configuration, gridProxy));
}
 

@Override
protected Bucket build(IgniteBucketBuilder builder, String key, RecoveryStrategy recoveryStrategy) {
    return builder.build(cache, key, recoveryStrategy);
}
 

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

private Bucket newBucket(String apiKey) {
    PricingPlan pricingPlan = PricingPlan.resolvePlanFromApiKey(apiKey);
    return bucket(pricingPlan.getLimit());
}
 

private Bucket bucket(Bandwidth limit) {
    return Bucket4j.builder()
        .addLimit(limit)
        .build();
}
 

@Test
public void givenAPIKey_whenFreePlan_thenReturnFreePlanBucket() {
    Bucket bucket = service.resolveBucket("FX001-UBSZ5YRYQ");

    assertEquals(PricingPlan.FREE.bucketCapacity(), bucket.getAvailableTokens());
}
 

@Test
public void givenAPIKey_whenBasiclan_thenReturnBasicPlanBucket() {
    Bucket bucket = service.resolveBucket("BX001-MBSZ5YRYP");

    assertEquals(PricingPlan.BASIC.bucketCapacity(), bucket.getAvailableTokens());
}
 

@Test
public void givenAPIKey_whenProfessionalPlan_thenReturnProfessionalPlanBucket() {
    Bucket bucket = service.resolveBucket("PX001-NBSZ5YRYY");

    assertEquals(PricingPlan.PROFESSIONAL.bucketCapacity(), bucket.getAvailableTokens());
}