bucket4j

ehcache 를 사용하여 Bucket Sort Algorithm 을 사용하려 한다.

Refil

기본적으로 gready 는 정해진 시간 마다 지속적으로 리필을 해주는 형식이다.
반면에 intervally 는 한번에 리필을 해준다.

public class Bucket {
    Refill refill = Refill.greedy(10, Duration.ofSeconds(1));
    Refill refill = Refill.intervally(10, Duration.ofSeconds(1));
}

Bandwidth

버킷의 총량을 설정해 준다.
simple 의 경우 greedy 를 사용하게 구현되어있다.
classic 의 경우 Refill 의 구현체를 직접 사용한다.

import java.time.Duration;

public class Bucket {
    Bandwidth bandwidth = Bandwidth.simple(10, Duration.ofSeconds(1));
    Bandwidth bandwidth = Bandwidth.classic(10, Refill.greedy(10, Duration.ofSeconds(1)));
}

Bucket

컨트롤 하기 위해 사용할 객체이다.

import java.time.Duration;

public class Bucket {
    Bucket bucket = Bucket4j.builder().addLimit(Bandwidth.simple(10, Duration.ofSeconds(1))).build();
}

Example

ip 접근 제어의 경우 nginx 에서 설정을 제공하기도 하지만 spring 내에서 특정 API 접근에 대해서 처리하려고 한다.

아래의 라이브러리들을 추가해준다.

For java 11+

implementation(     'com.bucket4j:bucket4j-core:8.1.0',     'com.bucket4j:bucket4j-jcache:8.1.0',     'org.ehcache:ehcache:3.10.8')

CacheManager 를 생성한다.

@Configuration
@EnableCaching
@EnableScheduling
@Slf4j
public class CacheConfig {
    private static final CacheManager cacheManager = Caching.getCachingProvider().getCacheManager();

    @Bean
    public Cache<String, byte[]> ipRateLimitCache() {
        return cacheManager.createCache("rateLimitCache", new MutableConfiguration<String, byte[]>().setExpiryPolicyFactory(
                () -> new CreatedExpiryPolicy(Duration.ONE_DAY)));
    }      // 0 시 마다 모든 캐시 초기화

    @Scheduled(cron = "0 0 0 * * *")
    public void clearCache() {
        cacheManager.getCacheNames().forEach(cacheName -> cacheManager.getCache(cacheName).clear());
        log.info("======= Clearing cache. =======");
    }
}

interceptor 를 생성한다.

@Slf4j
@Component
public class IpThrottlingInterceptor implements HandlerInterceptor {
    private static final String LOG_FORMAT = "ip: {}, rest token: {}";

    private static BucketConfiguration ofBuilder(final ConfigurationBuilder builder) {
        return builder.addLimit(Bandwidth.classic(5, Refill.intervally(5, Duration.ofDays(1L)))).build();
    }      // cache for storing token buckets, where IP is key.

    private final ProxyManager<String> buckets;

    @Autowired
    public IpThrottlingInterceptor(final Cache<String, byte[]> ipRateLimitCache) {
        buckets = new JCacheProxyManager<>(ipRateLimitCache);
    }

    @Override
    public boolean preHandle(final @NotNull HttpServletRequest request, final @NotNull HttpServletResponse response, final @NotNull Object handler) throws IOException {
        final String ip = IpAddressUtil.getClientIp(request);
        final Bucket bucket = this.ofBucket(ip);
        if (bucket.tryConsume(1)) {
            log.info(LOG_FORMAT, ip, bucket.getAvailableTokens());
            return true;
        }
        return this.error(ip, bucket, response);
    }

    private BucketProxy ofBucket(final String name) {
        return buckets.builder().build(name, IpThrottlingInterceptor.ofBuilder(BucketConfiguration.builder().addLimit(Bandwidth.classic(1, Refill.greedy(1, Duration.ofMinutes(1))))));
    }

    private boolean error(final String ip, final @NotNull Bucket bucket, final @NotNull HttpServletResponse response) throws IOException {
        log.error(LOG_FORMAT, ip, bucket.getAvailableTokens());
        response.sendError(HttpStatus.TOO_MANY_REQUESTS.value(), "To many requests");
        return false;
    }
}

인터셉터에 등록을 해주면 /api 로 접근하는 모든 요청에 대해 bucket 로직이 실행 된다.

import org.springframework.context.annotation.Configuration;

@Configuration
@RequiredArgsConstructor
public class WebMvcConfig implements WebMvcConfigurer {
    private final IpThrottlingInterceptor ipThrottlingInterceptor;

    @Override
    public void addInterceptors(final InterceptorRegistry registry) {
        registry.addInterceptor(ipThrottlingInterceptor).addPathPatterns("/api");
    }
}

caffeine 사용하는 경우

@Configuration
@EnableCaching
@Slf4j
public class CacheConfig {

    @Bean
    public CacheManager cacheManager() {
        final SimpleCacheManager cacheManager = new SimpleCacheManager();
        cacheManager.setCaches(
                Arrays.stream(CacheType.values())
                        .map(cache -> new CaffeineCache(
                                cache.getName(),
                                Caffeine.newBuilder()
                                        .expireAfterWrite(cache.getExpireAfterWrite(), TimeUnit.DAYS)
                                        .maximumSize(cache.getMaximumSize())
//                                        .evictionListener((Object key, Object value, RemovalCause cause) ->
//                                                log.info("Key {} was evicted ({}): {}", key, cause, value))
                                        .recordStats()
                                        .build()
                        ))
                        .toList());
        return cacheManager;
    }

    @Bean
    public ProxyManager<Object> proxyManager() {
        final CacheType cacheType = CacheType.BUCKET;
        return new CaffeineProxyManager<>((Caffeine) Caffeine.newBuilder()
                .maximumSize(cacheType.getMaximumSize())
                .recordStats(), Duration.ofDays(cacheType.getExpireAfterWrite()));
    }
}

@Slf4j
@Component
@RequiredArgsConstructor
public class IpThrottlingInterceptor implements HandlerInterceptor {

    private final ProxyManager<Object> proxyManager;

    private BucketProxy ofBucket(final String name, final BucketConfiguration configuration) {
        return proxyManager.builder().build(name, configuration);
    }

    private static final String LOG_FORMAT = "ip: {}, rest token: {}, today rest token: {}";

    private static final BucketConfiguration configuration = IpThrottlingInterceptor.ofBuilder(
            BucketConfiguration.builder().addLimit(Bandwidth.classic(Bucket4jConstant.MIN_CAPACITY,
                    Refill.greedy(Bucket4jConstant.REFILL, Duration.ofMinutes(Bucket4jConstant.MINUTES)))),
            Bucket4jConstant.DAY_CAPACITY);

    private static final BucketConfiguration totalConfiguration = IpThrottlingInterceptor.ofBuilder(
            BucketConfiguration.builder(), Bucket4jConstant.TOTAL_CAPACITY);

    private static BucketConfiguration ofBuilder(final @NotNull ConfigurationBuilder builder, final Integer capacity) {
        return builder.addLimit(Bandwidth.classic(capacity,
                Refill.intervally(capacity, Duration.ofDays(Bucket4jConstant.DAYS)))).build();
    }

    @Override
    public boolean preHandle(final @NotNull HttpServletRequest request, final @NotNull HttpServletResponse response,
                             final @NotNull Object handler) throws Exception {
        final String ip = IpAddressUtil.getClientIp(request);

        final Bucket bucket = this.ofBucket(ip, configuration);
        final Bucket totalBucket = this.ofBucket(Bucket4jConstant.TOTAL_BUCKET, totalConfiguration);
        if (totalBucket.tryConsume(1)) {
            final long totalRestToken = totalBucket.getAvailableTokens();
            if (bucket.tryConsume(1)) {
                log.info(LOG_FORMAT, ip, bucket.getAvailableTokens(), totalRestToken);
                return true;
            } else {
                totalBucket.addTokens(1);
                return this.error(ip, bucket, totalRestToken, response);
            }
        } else {
            return this.error(ip, bucket, totalBucket.getAvailableTokens(), response);
        }
    }


    /**
     * 이미 있는 메일 인 경우 전송 시도 되지 않아 소모된 토큰 복구
     */
    @Override
    public void afterCompletion(final @NotNull HttpServletRequest request, final @NotNull HttpServletResponse response,
                                final @NotNull Object handler, Exception exception) {
        if (response.getStatus() == 400) {
            final String ip = IpAddressUtil.getClientIp(request);

            final Bucket bucket = this.ofBucket(ip, configuration);
            bucket.addTokens(1L);

            final Bucket totalBucket = this.ofBucket(Bucket4jConstant.TOTAL_BUCKET, totalConfiguration);
            totalBucket.addTokens(1L);

            log.info("[Restore Bucket]" + LOG_FORMAT, ip, bucket.getAvailableTokens(), totalBucket.getAvailableTokens());

        }

    }


    private boolean error(final String ip, final @NotNull Bucket bucket, final long totalBucket,
                          final @NotNull HttpServletResponse response) throws IOException {
        log.error(LOG_FORMAT, ip, bucket.getAvailableTokens(), totalBucket);
        response.sendError(HttpStatus.TOO_MANY_REQUESTS.value(), "To many requests");
        return false;
    }