Bucket4j로 트래픽 제한하기(Redis & MariaDB)

개요

최근 업무 프로젝트에서 특정(요금이 부가되는) 로직에 대해 월별 사용량을 제한하는 기능이 추가되어야 했습니다. 이와 관련하여 처리율 제한 기술을 알아보았는데 Bucket4j, Guava, RateLimitj, Resilience4j 등 다양한 라이브러리가 있는 걸 알게 되었고, 프로젝트 환경인 Spring Boot 2.7.x, MariaDB(nosql & in-momory부재), Java 11에 알맞은 Bucket4j를 선택하게 되었습니다.

다음은 고려한 라이브러리들의 특징과 그 선택 이유입니다.

• Guava: 다양한 핵심 Java 라이브러리를 제공하지만, 단순히 Rate Limiting 기능만을 위해 사용하기에는 다소 무거운 느낌이다.
• Resilience4j: 서킷 브레이커를 제공하는 라이브러리로 Rate Limiter 기능도 포함되어 있지만, 마찬가지로 단순히 Rate Limiting 기능만을 위해 사용하기에는 다소 무거운 느낌이다.
• RateLimitj: 슬라이딩 윈도우 알고리즘을 기반으로 하는 Rate Limiter 구현체이다. 그러나 공식 문서에 더 이상 업데이트가 없을 것이라고 명시되어 있고, Bucket4j를 대체제로 권장하고 있다.
• Bucket4j: 토큰 버킷 알고리즘을 기반으로 하며, lock-free한 구현으로 멀티 스레딩 환경에서의 확장성이 우수하며, Rate Limiting만을 목적으로 제공되는 라이브러리이다.

 따라서 이 글에서는 Bucket4j를 사용하여 트래픽 제한을 구현하는 방법에 대해 정리합니다. 

 

 

토큰 버킷 알고리즘 이란?

토큰 버킷 알고리즘(Token Bucket Algorithm)은 네트워크 트래픽 및 데이터 전송 속도를 제어하는 데 널리 사용되는 기법입니다. 고정된 속도로 토큰이 버킷에 채워지며, 각 요청은 버키에서 하나 이상의 토큰을 소비합니다. 버킷에 충분한 토큰이 없는 경우, 요청은 대기하거나 거부되며 이 방식은 트래픽을 유연하게 제어하고 피크 시간에 대한 부하를 관리하는데 효과적입니다. 주요 특징과 동작 원리는 다음과 같습니다.

• 토큰의 생성과 소비: 버킷에는 일정한 속도로 토큰이 추가된다. 예를 들어, 1분에 100개의 토큰이 생성될 수 있으며 요청이 들어올 때마다 버킷에서 토큰을 하나 또는 그 이상 소비한다.
• 요청 제한: 버킷에 토큰이 충분하지 않으면 새로운 요청은 처리되지 않고 버려진다.
• 토큰의 재충전: 일정 시간이 지나면 토큰이 다시 생성되어 버킷에 추가된다. 예를 들어, 초기에 120개의 토큰으로 시작하여 1분에 100개씩 재충전되는 방식이다.

 

 

Bucket4j 적용

Bucket4j는 로컬 메모리 외에도, JDBC(MySQL, PostgreSQL, Oracle, MSSQL)와 Redis와 같이 다양한 분산 환경의 데이터베이스를 지원합니다.

아래는 MariaDB와 Redis에 Bucket4j를 적용한 예제 코드들입니다. 로컬 메모리(캐시)를 사용한 코드에 대해서는 인터넷에 다양한 예제가 있으므로, 필요하다면 해당 자료들을 참조하시면 좋을 것 같습니다. 실제 서비스 환경에서는 대개 여러 인스턴스가 동시에 운영되므로, 이처럼 분산 데이터베이스에 Bucket4j를 적용하는 것이 필요합니다.

• Bucket4j를 MariaDB에 적용한 예제: Bucket4j-mariadb
• Bucket4j를 Redis에 적용한 예제: Bucket4j-redis

 

Bucket4j-mariadb

1. table 생성

먼저 아래와 같이 'bucket'이라는 이름의 테이블을 생성합니다. 공식 문서에서는 테이블 이름을 'buckets'로 지정하였지만, 실제로 BucketTableSetting.getDefault() 메서드를 호출하면 테이블 이름이 'bucket'으로 반환되므로, 테이블 이름을 'bucket'으로 설정하는 것이 편리합니다. 

create table bucket
(
    id    bigint not null
        primary key,
    state blob   null
);

2. dependency

'bucket4j-mysql' 의존성을 추가합니다.

dependencies {
    ...    
    
    // bucket4j
    implementation 'com.bucket4j:bucket4j-mysql:8.7.0'
}

3. APIRateLimiter

@Slf4j
@Component
public class APIRateLimiter {
    // 버킷의 용량을 설정, 이는 버킷에 담길 수 있는 토큰의 최대 수를 의미
    private static final int CAPACITY = 3;
    // 토큰이 얼마나 빠르게 재충전될지 설정, 이는 지정된 시간 동안 버킷에 추가될 토큰의 수를 의미
    private static final int REFILL_AMOUNT = 3;
    // 토큰이 재충전되는 빈도를 설정
    private static final Duration REFILL_DURATION = Duration.ofSeconds(5);

    // MySQLSelectForUpdateBasedProxyManager 객체를 생성, 이 객체는 버킷의 생성 및 관리를 담당
    private final MySQLSelectForUpdateBasedProxyManager<Long> proxyManager;
    // 동일한 API 키에 대한 요청을 처리하기 위해 버킷을 재사용하기 위해 버킷을 저장하는 맵을 생성
    private final ConcurrentMap<String, BucketProxy> buckets = new ConcurrentHashMap<>();

    /**
     * APIRateLimiter 생성자.
     * 
     * @param dataSource 데이터 소스
     */
    public APIRateLimiter(DataSource dataSource) {
        // 버킷 테이블 설정을 가져옴
        var tableSettings = BucketTableSettings.getDefault();
        // SQL 프록시 설정을 생성
        var sqlProxyConfiguration = SQLProxyConfiguration.builder()
                .withTableSettings(tableSettings)
                .build(dataSource);
        // SQL 프록시 설정을 이용해 MySQLSelectForUpdateBasedProxyManager 객체를 생성
        proxyManager = new MySQLSelectForUpdateBasedProxyManager<>(sqlProxyConfiguration);
    }

    /**
     * API 키에 해당하는 버킷을 가져오거나, 없을 경우 새로 생성하는 메서드.
     * 
     * @param apiKey API 키
     * @return 해당 API 키에 대응하는 버킷
     */
    private BucketProxy getOrCreateBucket(String apiKey) {
        return buckets.computeIfAbsent(apiKey, key -> {
            // API 키의 해시 코드를 버킷 ID로 사용
            Long bucketId = (long) key.hashCode();
            // 버킷 설정을 생성
            var bucketConfiguration = createBucketConfiguration();
            // 버킷 ID와 버킷 설정을 이용해 버킷을 생성하고, 이를 반환
            return proxyManager.builder().build(bucketId, bucketConfiguration);
        });
    }

    /**
     * 버킷 설정을 생성하는 메서드.
     * 
     * @return 생성된 버킷 설정
     */
    private BucketConfiguration createBucketConfiguration() {
        return BucketConfiguration.builder()
                // 버킷에 대한 제한(용량과 재충전 속도)을 설정
                .addLimit(Bandwidth.builder().capacity(CAPACITY).refillIntervally(REFILL_AMOUNT, REFILL_DURATION).build())
                .build();
    }

    /**
     * API 키에 해당하는 버킷에서 토큰을 소비하려고 시도하는 메서드.
     * 
     * @param apiKey API 키
     * @return 토큰 소비 성공 여부
     */
    public boolean tryConsume(String apiKey) {
        // API 키에 해당하는 버킷을 가져옴
        BucketProxy bucket = getOrCreateBucket(apiKey);
        // 버킷에서 토큰을 소비하려고 시도하고, 그 결과를 반환
        boolean consumed = bucket.tryConsume(1);
        log.info("API Key: {}, Consumed: {}, Time: {}", apiKey, consumed, LocalDateTime.now());
        return consumed;
    }
}

만약 여기서 커스텀한 테이블 혹은 컬럼을 지정하거나, 스키마를 지정하고 싶다면 다음과 같이 정의하면 됩니다.

  var tableSettings = BucketTableSettings.customSettings("test.bucket","id", "state");

 

Bucket4j-redis

Bucket4j를 Redis에 적용하는 방법은 MariaDB에 적용하는 방법과 거의 동일합니다. Bucket4j는 Redis의 Redisson, Jedis, Lettuce를 모두 지원하며, 아래 예제에서는 Lettuce를 사용하였습니다.

1. dependency

의존성을 추가합니다.

dependencies {
    ...    
    
    // bucket4j
    implementation 'com.bucket4j:bucket4j-redis:8.7.0'
}

2. APIRateLimiter

@Slf4j
@Component
public class APIRateLimiter {
    // 버킷의 용량을 설정, 이는 버킷에 담길 수 있는 토큰의 최대 수를 의미
    private static final int CAPACITY = 3;
    // 토큰이 얼마나 빠르게 재충전될지 설정, 이는 지정된 시간 동안 버킷에 추가될 토큰의 수를 의미
    private static final int REFILL_AMOUNT = 3;
    // 토큰이 재충전되는 빈도를 설정
    private static final Duration REFILL_DURATION = Duration.ofSeconds(5);

    // LettuceBasedProxyManager 객체를 생성, 이 객체는 버킷의 생성 및 관리를 담당
    private final LettuceBasedProxyManager<String> proxyManager;
    // 동일한 API 키에 대한 요청을 처리하기 위해 버킷을 재사용하기 위해 버킷을 저장하는 맵을 생성
    private final ConcurrentMap<String, Bucket> buckets = new ConcurrentHashMap<>();

    /**
     * APIRateLimiter 생성자.
     * 
     * @param redisClient Redis 클라이언트
     */
    public APIRateLimiter(RedisClient redisClient) {
        // Redis 연결을 생성
        StatefulRedisConnection<String, byte[]> connection = redisClient.connect(RedisCodec.of(StringCodec.UTF8, ByteArrayCodec.INSTANCE));
        // Redis 연결을 이용해 LettuceBasedProxyManager 객체를 생성
        this.proxyManager = LettuceBasedProxyManager.builderFor(connection)
                .withExpirationStrategy(ExpirationAfterWriteStrategy.basedOnTimeForRefillingBucketUpToMax(Duration.ofSeconds(100)))
                .build();
    }

    /**
     * API 키에 해당하는 버킷을 가져오거나, 없을 경우 새로 생성하는 메서드.
     * 
     * @param apiKey API 키
     * @return 해당 API 키에 대응하는 버킷
     */
    private Bucket getOrCreateBucket(String apiKey) {
        return buckets.computeIfAbsent(apiKey, key -> {
            // 버킷 설정을 생성
            BucketConfiguration configuration = createBucketConfiguration();
            // 버킷 ID와 버킷 설정을 이용해 버킷을 생성하고, 이를 반환
            return proxyManager.builder().build(key, configuration);
        });
    }

    /**
     * 버킷 설정을 생성하는 메서드.
     * 
     * @return 생성된 버킷 설정
     */
    private BucketConfiguration createBucketConfiguration() {
        return BucketConfiguration.builder()
                // 버킷에 대한 제한(용량과 재충전 속도)을 설정
                .addLimit(Bandwidth.simple(CAPACITY, REFILL_DURATION).withInitialTokens(REFILL_AMOUNT))
                .build();
    }

    /**
     * API 키에 해당하는 버킷에서 토큰을 소비하려고 시도하는 메서드.
     * 
     * @param apiKey API 키
     * @return 토큰 소비 성공 여부
     */
    public boolean tryConsume(String apiKey) {
        // API 키에 해당하는 버킷을 가져옴
        Bucket bucket = getOrCreateBucket(apiKey);
        // 버킷에서 토큰을 소비하려고 시도하고, 그 결과를 반환
        boolean consumed = bucket.tryConsume(1);
        log.info("API Key: {}, Consumed: {}, Time: {}", apiKey, consumed, LocalDateTime.now());
        return consumed;
    }
}

 

 

테스트

아래와 같이 AOP(Aspect-Oriented Programming)를 적용하여 코드의 가독성을 높이고 중복을 제거할 수 있습니다.

RateLimit

먼저, 각 메서드에 적용할 사용자 정의 애노테이션인 RateLimit을 생성합니다. 이 애노테이션은 API 키를 지정하는 'key'라는 속성을 가집니다.

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimit {
    String key();
}

RateLimitingAspect

다음으로, RateLimit 애노테이션을 처리할 Aspect 클래스인 RateLimitingAspect를 생성합니다. 이 클래스는 APIRateLimiter 객체를 사용하여 API 키에 해당하는 버킷에서 토큰을 소비하려고 시도합니다. 토큰이 충분하면 요청이 성공적으로 처리되고, 그렇지 않으면 예외가 발생합니다.

@Aspect
@Component
public class RateLimitingAspect {

    private final APIRateLimiter apiRateLimiter;

    @Autowired
    public RateLimitingAspect(APIRateLimiter apiRateLimiter) {
        this.apiRateLimiter = apiRateLimiter;
    }

    @Around("@annotation(rateLimit)")
    public Object rateLimit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        if (apiRateLimiter.tryConsume(rateLimit.key())) {
            return joinPoint.proceed();
        } else {
            throw new RuntimeException("Rate limit exceeded for key: " + rateLimit.key());
        }
    }
}

RateLimitingController

@RestController
@Slf4j
public class RateLimitingController {
    @Autowired
    private RateLimitingService rateLimitingService;

    @GetMapping("/test")
    public String test() {
        return rateLimitingService.run();
    }
}

RateLimitingService

@Service
public class RateLimitingService {
    @RateLimit(key = "someUniqueKey")
    public String run() {
        return "요청 성공";
    }
}

토큰 최대 갯수인 3개를 초과하면 실패, 5초후 토큰이 채워지면 성공

 

 

정리

Bucket4j 라이브러리를 활용하면, API 호출의 처리율을 손쉽게 제한할 수 있습니다. 또한 다양한 곳에서 적용하고 부가 기능을 넣을 수 있는데, Baeldung의 Bucket4j 가이드를 보면 인터셉터를 이용하여 처리율을 제한할 수도 있고, 아래처럼 사용자의 요금제에 따라 API 호출 제한을 다르게 적용하는 것도 가능합니다.

enum PricingPlan {
    // 한시간에 20번 사용 가능
    FREE {
        Bandwidth getLimit() {
            return Bandwidth.classic(20, Refill.intervally(20, Duration.ofHours(1)));
        }
    },
    // 한시간에 40번 사용 가능
    BASIC {
        Bandwidth getLimit() {
            return Bandwidth.classic(40, Refill.intervally(40, Duration.ofHours(1)));
        }
    },
    // 한시간에 100번 사용 가능
    PROFESSIONAL {
        Bandwidth getLimit() {
            return Bandwidth.classic(100, Refill.intervally(100, Duration.ofHours(1)));
        }
    };
    //..
}

Bucket4j는 이와 같이 다양한 환경에서 API 호출 제한을 구현하는 데 유용하게 사용될 수 있습니다. 따라서, 서비스의 성능을 향상시키고 서버의 부하를 관리하는 데 필요한 기능을 쉽고 빠르게 구현할 수 있습니다.