io.github.bucket4j.Bucket4j 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();
}
 

/**
 * Creates a new named strategy with specified {@link TokenBucket} configuration,
 * with minimum backoff period and with specific throttling headers scheme.
 *
 * @param tokenBucket {@link TokenBucket} configuration.
 * @param minimumBackoff optional {@link Duration} that defines a minimum backoff period
 *                       for throttled requests. By default, it will be set to 0 seconds.
 * @param headersScheme optional {@link ThrottlingHeaders} to define specific RateLimit Header Scheme
 *                      for HTTP. By default, no throttling headers will be used. The strategy will only use
 *                      standard HTTP Retry-After header.
 * @param sendQuota indicates whether to use quota header for the scheme.
 * @param name optional name of the strategy. By default, it will be assigned with a predefined name.
 */
TokenBucketThrottlingStrategy(TokenBucket tokenBucket,
                              @Nullable Duration minimumBackoff,
                              @Nullable ThrottlingHeaders headersScheme,
                              boolean sendQuota,
                              @Nullable String name) {
    super(name);
    // construct the bucket builder
    final LocalBucketBuilder builder = Bucket4j.builder().withNanosecondPrecision();
    for (BandwidthLimit limit : tokenBucket.limits()) {
        builder.addLimit(limit.bandwidth());
    }
    // build the bucket
    asyncBucket = builder.build().asAsync();
    minimumBackoffSeconds = (minimumBackoff == null) ? 0L : minimumBackoff.getSeconds();
    this.headersScheme = headersScheme;
    this.sendQuota = sendQuota;
    quota = sendQuota ? tokenBucket.toSpecString() : null;
}
 

/**
 * @param capacity Capacity is the maximum number of tokens can be consumed.
 * @param tokensGeneratedInPeriod Tokens generated in period.
 * @param period Period that generating specific number of tokens.
 */
public BucketRateLimiter(long capacity, long tokensGeneratedInPeriod, Duration period) {
  Bandwidth bandwidth =
      Bandwidth.classic(capacity, Refill.greedy(tokensGeneratedInPeriod, period));
  this.bucket =
      Bucket4j.builder()
          .addLimit(bandwidth)
          .withSynchronizationStrategy(SynchronizationStrategy.SYNCHRONIZED)
          .build();
}
 

@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));
}
 

@Test
public void tetsSerializationOfEntryProcessors() {
    BucketConfiguration configuration = Bucket4j.configurationBuilder()
            .addLimit(simple(10, ofSeconds(1)))
            .build();
    GridCommand command = new AddTokensCommand(42);

    testSerialization(new InitStateProcessor<>(configuration));
    testSerialization(new ExecuteProcessor<>(command));
    testSerialization(new InitStateAndExecuteProcessor(command, configuration));

    testSerialization(new SerializableFunctionAdapter<>(new ExecuteProcessor<>(command)));
}
 

@Before
public void setup() {
    for (SerializationHandle<?> serializationHandle : Bucket4j.getSerializationHandles()) {
        String protoTypeId = "bucket4j.Bucket4jType_" + serializationHandle.getTypeId();
        serializerByClass.put(serializationHandle.getSerializedType(), new ProtobufMessageMarshaller<>(serializationHandle, protoTypeId));
    }
}
 

@Override
public void registerMarshallers(SerializationContext serCtx) {
    for (SerializationHandle<?> serializationHandle : Bucket4j.getSerializationHandles()) {
        String protoTypeId = "bucket4j.Bucket4jType_" + serializationHandle.getTypeId();
        serCtx.registerMarshaller(new ProtobufMessageMarshaller<>(serializationHandle, protoTypeId));
    }
}
 

@Override
public void registerSchema(SerializationContext serCtx) {
    StringBuilder protoBuilder = new StringBuilder(FOOTER);

    for (SerializationHandle<?> serializationHandle : Bucket4j.getSerializationHandles()) {
        String typeName = "Bucket4jType_" + serializationHandle.getTypeId();
        String typeDefinition = TYPE_TEMPLATE.replace("[type_name]", typeName);
        protoBuilder.append(typeDefinition);
    }

    String generatedProtoFile = protoBuilder.toString();
    FileDescriptorSource protoSource = FileDescriptorSource.fromString(getProtoFileName(), generatedProtoFile);
    serCtx.registerProtoFiles(protoSource);
}
 

@Test
public void tetsSerializationOfEntryProcessors() {
    BucketConfiguration configuration = Bucket4j.configurationBuilder()
            .addLimit(simple(10, ofSeconds(1)))
            .build();
    GridCommand command = new AddTokensCommand(42);

    testSerialization(new InitStateProcessor<>(configuration));
    testSerialization(new ExecuteProcessor<>(command));
    testSerialization(new InitStateAndExecuteProcessor(command, configuration));
}
 

@Test
public void tetsSerializationOfEntryProcessors() {
    BucketConfiguration configuration = Bucket4j.configurationBuilder()
            .addLimit(simple(10, ofSeconds(1)))
            .build();
    GridCommand command = new AddTokensCommand(42);

    testSerialization(new InitStateProcessor<>(configuration));
    testSerialization(new ExecuteProcessor<>(command));
    testSerialization(new InitStateAndExecuteProcessor(command, configuration));
}
 

/**
 * Resets the Token-Bucket configuration at runtime by providing a new set of limits.
 * Empty set of limits will remove previously set limits. Reconfiguration will take place asynchronously.
 * @param tokenBucket Token-Bucket configuration
 * @return A {@link CompletableFuture} to handle asynchronous result
 */
public CompletableFuture<Void> reconfigure(TokenBucket tokenBucket) {
    // construct the configuration builder
    final ConfigurationBuilder builder = Bucket4j.configurationBuilder();
    for (BandwidthLimit limit : tokenBucket.limits()) {
        builder.addLimit(limit.bandwidth());
    }
    // reconfigure the bucket
    return asyncBucket.replaceConfiguration(builder.build())
                      .thenRun(() -> quota = sendQuota ? tokenBucket.toSpecString() : null);
}
 

@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 ConfigurationBuilder prepareBucket4jConfigurationBuilder(RateLimit rl) {
	ConfigurationBuilder configBuilder = Bucket4j.configurationBuilder();
	for (BandWidth bandWidth : rl.getBandwidths()) {
			Bandwidth bucket4jBandWidth = Bandwidth.simple(bandWidth.getCapacity(), Duration.of(bandWidth.getTime(), bandWidth.getUnit()));
			if(bandWidth.getFixedRefillInterval() > 0) {
				bucket4jBandWidth = Bandwidth.classic(bandWidth.getCapacity(), Refill.intervally(bandWidth.getCapacity(), Duration.of(bandWidth.getFixedRefillInterval(), bandWidth.getFixedRefillIntervalUnit())));
			}
			configBuilder = configBuilder.addLimit(bucket4jBandWidth);
	};
	return configBuilder;
}
 

public ProxyManager<String>  resolve(String cacheName) {
	Cache springCache = cacheManager.getCache(cacheName);
	if (springCache == null) {
		throw new JCacheNotFoundException(cacheName);
	}

	return Bucket4j.extension(JCache.class).proxyManagerForCache(springCache);
}
 

public ProxyManager<String> resolve(String cacheName) {
	Cache<Object, Object> cache = cacheContainer.getCache(cacheName);
	if (cache == null) {
		throw new JCacheNotFoundException(cacheName);
	}

	FunctionalMapImpl functionalMap = FunctionalMapImpl.create(cache.getAdvancedCache());
	return Bucket4j.extension(Infinispan.class).proxyManagerForMap(ReadWriteMapImpl.create(functionalMap));
}
 

@Override
protected ProxyManager<String> newProxyManager() {
    return Bucket4j.extension(getExtensionClass()).proxyManagerForMap(map);
}
 

@Test(expected = IllegalArgumentException.class)
@Override
public void testThatImpossibleToPassNullCacheToProxyManagerConstructor() {
    Bucket4j.extension(getExtensionClass()).proxyManagerForMap(null);
}
 

@Test(expected = IllegalArgumentException.class)
@Override
public void testThatImpossibleToPassNullCacheToProxyManagerConstructor() {
    Bucket4j.extension(getExtensionClass()).proxyManagerForMap(null);
}
 

@Override
protected ProxyManager<String> newProxyManager() {
    return Bucket4j.extension(Infinispan.class).proxyManagerForMap(readWriteMap);
}
 

@Override
public ProxyManager<String> resolve(String cacheName) {
	com.hazelcast.core.IMap<String, GridBucketState> map = hazelcastInstance.getMap(cacheName);
	return Bucket4j.extension(Hazelcast.class).proxyManagerForMap(map);
}
 

@Test(expected = IllegalArgumentException.class)
@Override
public void testThatImpossibleToPassNullCacheToProxyManagerConstructor() {
    Bucket4j.extension(getExtensionClass()).proxyManagerForMap(null);
}
 

@Override
protected ProxyManager<String> newProxyManager() {
    return Bucket4j.extension(Infinispan.class).proxyManagerForMap(readWriteMap);
}
 

@Test(expected = UnsupportedOperationException.class)
public void checkThatInfinispanProviderUnsupported() {
    Bucket4j.extension(JCache.class).builder().
            addLimit(Bandwidth.simple(10, Duration.ofSeconds(60))).
            build(cache, "42", RecoveryStrategy.RECONSTRUCT);
}
 

@Test(expected = IllegalArgumentException.class)
@Override
public void testThatImpossibleToPassNullCacheToProxyManagerConstructor() {
    Bucket4j.extension(getExtensionClass()).proxyManagerForCache(null);
}
 

@Override
protected ProxyManager<String> newProxyManager() {
    return Bucket4j.extension(getExtensionClass()).proxyManagerForCache(cache);
}
 

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

@Override
public ProxyManager<String> resolve(String cacheName) {
	Cache<Object, Object> cache = cacheContainer.getCache(cacheName);
	// TODO how to create an instance of ReadWriteMap
	return Bucket4j.extension(io.github.bucket4j.grid.infinispan.Infinispan.class).proxyManagerForMap(null);
}
 

@Override
public ProxyManager<String> resolve(String cacheName) {
	org.apache.ignite.IgniteCache<String, GridBucketState> cache = ignite.cache(cacheName);
	return Bucket4j.extension(io.github.bucket4j.grid.ignite.Ignite.class).proxyManagerForCache(cache);
}
 

private Supplier<BucketConfiguration> getBucketConfiguration(Long capacity, Duration period) {
    return () -> Bucket4j.configurationBuilder().addLimit(Bandwidth.simple(capacity, period)).build();
}