bucket4j 사용한 제어

Posted on | In |

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

Refil

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

1
2
3
4
5
public class Bucket {
    Refill refill = Refill.greedy(10, Duration.ofSeconds(1));

    Refill refill = Refill.intervally(10, Duration.ofSeconds(1));
}

Bandwidth

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

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

Bucket

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

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

Example

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

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

For java 11+

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

CacheManager 를 생성한다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

@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 를 생성한다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
@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 로직이 실행 된다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14

@Configuration
@RequiredArgsConstructor
public class WebMvcConfig implements WebMvcConfigurer {

    private final IpThrottlingInterceptor ipThrottlingInterceptor;

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