Spring Boot 3 RateLimiter Bucket4j with Hazelcast

Coding

•  Config Annotation (RateLimiter and RateLimiters)

package com.tra22.ratelimiter.config.annotation;

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;
import java.util.concurrent.TimeUnit;

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Documented
public @interface RateLimiter {
        TimeUnit timeUnit() default TimeUnit.MINUTES;
        long timeValue() default 60;
        long restriction() default 500;
}

package com.tra22.ratelimiter.config.annotation;

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;

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Documented
public @interface RateLimiters {
    RateLimiter[] value();
}

• CacheConfig with Hazelcast

package com.tra22.ratelimiter.config.cache;

import com.hazelcast.client.config.ClientConfig;
import com.hazelcast.config.MapConfig;
import com.hazelcast.core.Hazelcast;
import com.hazelcast.core.HazelcastInstance;
import com.hazelcast.map.IMap;
import com.hazelcast.spring.cache.HazelcastCacheManager;
import com.tra22.ratelimiter.payloads.RateLimiterKey;
import io.github.bucket4j.grid.hazelcast.HazelcastProxyManager;
import org.springframework.cache.CacheManager;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class CacheConfig {
    @Bean
    public HazelcastInstance hazelcastInstance() {
        HazelcastInstance instance = Hazelcast.newHazelcastInstance();
        instance.getConfig().addMapConfig(new MapConfig("rate-limiters"));
        instance.getConfig().getNetworkConfig().getRestApiConfig().setEnabled(true);
        return instance;
    }
    @Bean
    public IMap<RateLimiterKey, byte[]>  rateLimitingMap(HazelcastInstance hazelcastInstance) {
        return hazelcastInstance.getMap("rate-limiters");
    }
    @Bean
    public HazelcastProxyManager<RateLimiterKey> hazelcastProxyManager() {
        return new HazelcastProxyManager<RateLimiterKey>(hazelcastInstance().getMap("rate-limiters"));
    }
    @Bean
    public ClientConfig clientConfig() {
        ClientConfig cfg = ClientConfig.load();
        cfg.setClusterName("statsCluster");
        return cfg;
    }
    @Bean
    CacheManager hazelcastCacheManager() {
        return new HazelcastCacheManager(hazelcastInstance());
    }

}

• Create RateLimiterUtil 

package com.tra22.ratelimiter.utils;

import com.tra22.ratelimiter.config.annotation.RateLimiter;
import com.tra22.ratelimiter.config.annotation.RateLimiters;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.BucketConfiguration;
import io.github.bucket4j.ConfigurationBuilder;
import jakarta.servlet.http.HttpServletRequest;
import org.slf4j.helpers.MessageFormatter;
import org.springframework.util.StringUtils;
import org.springframework.web.method.HandlerMethod;

import java.time.Duration;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

public class RateLimiterUtils {
    public static String getUniqueKeyName(HttpServletRequest httpServletRequest){
        if(httpServletRequest != null){
            return MessageFormatter.arrayFormat("{}:{}",
                    new Object[] {
                            StringUtils.hasText(httpServletRequest.getRemoteAddr()) ? httpServletRequest.getRemoteAddr() : "0.0.0.0",
                            StringUtils.hasText(httpServletRequest.getRemoteUser()) ? httpServletRequest.getRemoteUser() : "anonymous"
                    }
            ).getMessage();
        }
        return "0.0.0.0:anonymous:NONE_REQUEST_HOST:NONE_REQUEST_PORT";
    }
    public static BucketConfiguration rateLimiterAnnotationsToBucketConfiguration(List<RateLimiter> rateLimiters) {
        ConfigurationBuilder configBuilder  = new ConfigurationBuilder();
        rateLimiters.forEach(limiter -> configBuilder.addLimit(buildBandwidth(limiter)));
        return  configBuilder.build();
     }

    public  static Optional<List<RateLimiter>> getRateLimiters(HandlerMethod handlerMethod) {
        RateLimiters  rateLimitersAnnotation  = handlerMethod.getMethodAnnotation(RateLimiters.class);
        if(rateLimitersAnnotation !=  null) {
            return  Optional.of(Arrays.asList(rateLimitersAnnotation.value()));
        }
        RateLimiter  rateLimiterAnnotation  = handlerMethod.getMethodAnnotation(RateLimiter.class);
        if(rateLimiterAnnotation !=  null) {
            return  Optional.of(List.of(rateLimiterAnnotation));
        }
        return  Optional.empty();
    }

    private static Bandwidth buildBandwidth(RateLimiter rateLimiter) {
        TimeUnit timeUnit  = rateLimiter.timeUnit();
        long  timeValue  = rateLimiter.timeValue();
        long  restriction  = rateLimiter.restriction();
        return Bandwidth.builder()
                .capacity(restriction)
                .refillGreedy(restriction, convert(timeValue, timeUnit))
                .build();
    }
    private static Duration convert(long dur, TimeUnit timeUnit) {
        return Duration.of(dur, timeUnit.toChronoUnit());
    }
}

• Handler Interceptor Adaper for filter request to check RateLimit

package com.tra22.ratelimiter.config.handlers;

import com.tra22.ratelimiter.config.annotation.RateLimiter;
import com.tra22.ratelimiter.exception.BucketRateLimiterException;
import com.tra22.ratelimiter.payloads.RateLimiterKey;
import com.tra22.ratelimiter.utils.RateLimiterUtils;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.ConsumptionProbe;
import io.github.bucket4j.grid.hazelcast.HazelcastProxyManager;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.NonNull;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import java.util.List;
import java.util.Optional;

@Component
public class RateLimiterAnnotationHandlerInterceptorAdapter implements HandlerInterceptor {
    private final HazelcastProxyManager<RateLimiterKey> proxyManager;
    public RateLimiterAnnotationHandlerInterceptorAdapter(HazelcastProxyManager<RateLimiterKey> proxyManager) {
        this.proxyManager = proxyManager;
    }
    @Override
    public boolean preHandle(@NonNull HttpServletRequest request,
                             @NonNull HttpServletResponse response,
                             @NonNull Object handler) throws Exception {

        if (handler instanceof HandlerMethod handlerMethod) {

            //if into handlerMethod is present RateLimiter or RateLimiters annotation, we get it, if not, we get empty Optional
            Optional<List<RateLimiter>> rateLimiters = RateLimiterUtils.getRateLimiters(handlerMethod);

            if (rateLimiters.isPresent()) {
                //Get path from RequestMapping annotation(respectively we can get annotations such: GetMapping, PostMapping, PutMapping, DeleteMapping, because all of than annotations are extended from RequestMapping)
                RequestMapping requestMapping = handlerMethod.getMethodAnnotation(RequestMapping.class);

                //To get unique key we use bundle of 2-x values: path from RequestMapping and user id
                assert requestMapping != null;
                RateLimiterKey key = new RateLimiterKey(RateLimiterUtils.getUniqueKeyName(request), requestMapping.value());
                //Further we set key in proxy to get Bucket from cache or create a new Bucket
                Bucket bucket = this.proxyManager.builder().build(key, () -> RateLimiterUtils.rateLimiterAnnotationsToBucketConfiguration(rateLimiters.get()));

                ConsumptionProbe consumptionProbe =  bucket.tryConsumeAndReturnRemaining(1);
                //Try to consume token, if we don’t do that, we throw custom exception
                if (!consumptionProbe.isConsumed()) {
                    throw new BucketRateLimiterException("RateLimiter is applied.",  consumptionProbe);
                }
            }
        }
        return true;
    }
}

• ‎WebConfigHazelcastRateLimiter for add Inceptor Adapter

package com.tra22.ratelimiter.config;

import com.tra22.ratelimiter.config.handlers.RateLimiterAnnotationHandlerInterceptorAdapter;
import com.tra22.ratelimiter.payloads.RateLimiterKey;
import io.github.bucket4j.grid.hazelcast.HazelcastProxyManager;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

@Configuration
public class WebConfigHazelcastRateLimiter implements WebMvcConfigurer {
    @Autowired
    private HazelcastProxyManager<RateLimiterKey> hazelcastProxyManager;
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(new RateLimiterAnnotationHandlerInterceptorAdapter(hazelcastProxyManager));
    }
}

• Create Exception for throw when faced rate limit from hazelcast cache

package com.tra22.ratelimiter.exception;

import com.fasterxml.jackson.core.JsonProcessingException;
import io.github.bucket4j.ConsumptionProbe;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

@Slf4j
@Getter
public class BucketRateLimiterException extends RuntimeException{
    private final ConsumptionProbe consumptionProbe;
    public BucketRateLimiterException(String message, ConsumptionProbe consumptionProbe) throws JsonProcessingException {
        super(message);
        this.consumptionProbe = consumptionProbe;
        log.error("getNanosToWaitForReset: {} , getNanosToWaitForRefill: {}, getRemainingTokens: {}", this.consumptionProbe.getNanosToWaitForReset(), this.consumptionProbe.getNanosToWaitForRefill(), this.consumptionProbe.getRemainingTokens());
    }
}

• Create ExampleController for apply rate limit for testing

package com.tra22.ratelimiter.controllers;

import com.tra22.ratelimiter.config.annotation.RateLimiter;
import com.tra22.ratelimiter.config.annotation.RateLimiters;
import lombok.RequiredArgsConstructor;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.concurrent.TimeUnit;

@RestController
@RequiredArgsConstructor
public class ExampleController   {
    @RateLimiters({
        @RateLimiter(timeUnit = TimeUnit.SECONDS, timeValue = 10, restriction = 2),
        @RateLimiter(timeValue = 10, restriction = 5)
    })
    @GetMapping("/example")
    public ResponseEntity<Object> example() {
        return ResponseEntity.ok("Ok");
    }
}