[Spring Boot] Caffeine 로컬 캐시로 Rate Limiter를 구현해보자

# Rate Limit이란?

일정 시간 동안 호출될 수 있는 API 횟수를 제한하는 걸 말한다.

 

예를 들어 파일 업로드는 10분에 최대 3번만 가능 이런 걸 구현할 때 사용하는 개념이다.

 

자세한 건 아래글 참조

https://gngsn.tistory.com/224

Rate Limiter, 제대로 이해하기

 

 

이번 글에서는 로컬 캐시 라이브러리로 유명한 Caffeine으로 Rate Limiter 기능을 구현해 볼 것이다.

 

# 왜 Caffeine인가?

Spring Cloud Gateway나 Bucket4j를 활용해서 손쉽게 구현할 수도 있다. 

 

다만, 이번에는 외부 저장소 (ex. Redis) 없이 로컬 캐시로 라이트하게 구현하기 위해서 Caffeine을 선택했다.

 

단순 InMemory 저장을 위함이면 Collections를 써도되지 않나 싶었지만, TTL을 설정할 수 없다.

 

 

여담으로 당연하겠지만 Caffeine은 thread-safe 하다고 한다.

https://github.com/ben-manes/caffeine/issues/392

Question: Thread safety? · Issue #392 · ben-manes/caffeine

 

# 가상 요구 사항

파일 업로드 API에 Rate Limit을 구현해볼 예정이다.

 

인증된 사용자마다 10초동안 최대 5번만 업로드가 가능하다.

 

# 구현

build.gradle.kts

    implementation("com.github.ben-manes.caffeine:caffeine:3.1.8")

 

 

Rate Limiter

interface RateLimiter {
    fun checkLimit(userId: Long)

    fun increaseCount(userId: Long)

    class CaffeineLocalCacheImpl(
        val limit: Int,
        val duration: Long,
        val timeUnit: TimeUnit,
        val maxSize: Long = 100,
    ) : RateLimiter {
        private val cache =
            Caffeine
                .newBuilder()
                .maximumSize(maxSize)
                .expireAfterWrite(duration, timeUnit)
                .build<String, Any>()

        override fun checkLimit(userId: Long) {
            val uploadCount = cache.getIfPresent(userId.toString()) as Int?
            if (uploadCount != null && uploadCount >= limit) {
                throw RateLimitExceededException()
            }
        }

        override fun increaseCount(userId: Long) {
            val uploadCount = cache.getIfPresent(userId.toString()) as Int?
            if (uploadCount == null) {
                cache.put(userId.toString(), 1)
            } else {
                cache.put(userId.toString(), uploadCount + 1)
            }
        }
    }
}

 

Rate Limiter 인터페이스를 만들고 각 구현체는 Bean으로 등록하여 필요한 곳에서 Qualifier로 주입받아 사용할 수 있게끔 할 것이다.

 

Caffeine의 key는 userId: Long이고 value는 현재까지의 카운트이다.

 

class RateLimitExceededException(
    message: String = ErrorType.TOO_MANY_REQUEST.defaultMessage,
) : RuntimeException(message)

 

커스텀 예외를 하나 만들어준다. 그리고 이를 핸들 할 ControllerAdvice도 추가한다.

 

    @ExceptionHandler(RateLimitExceededException::class)
    fun handleRateLimitException(e: RateLimitExceededException): ResponseEntity<ErrorResponse> {
        return ResponseEntity
            .status(HttpStatus.TOO_MANY_REQUESTS)
            .body(
                ErrorResponse.from(ErrorType.TOO_MANY_REQUEST),
            )
    }

 

ErrorResponse과 ErrorType은 커스텀 객체로 적당히 구현해 준다.

 

 

RateLimiterConfig

@Configuration
class RateLimiterConfig() {
    @Bean
    fun fileUploadRateLimiter(): RateLimiter {
        return RateLimiter.CaffeineLocalCacheImpl(
            limit = 5,
            duration = 10,
            timeUnit = TimeUnit.SECONDS,
            maxSize = 100,
        )
    }

	// 추후 다른 기능에 필요하다면 추가할 수 있도록 만들었다
    @Bean
    fun programApplicationRateLimiter(): RateLimiter {
        return RateLimiter.CaffeineLocalCacheImpl(
            limit = 5,
            duration = 24,
            timeUnit = TimeUnit.HOURS,
            maxSize = 100,
        )
    }
}

 

별도 Config 파일을 만들어서 필요한 리미터를 빈으로 등록한다.

 

 

fun <T> rateLimit(rateLimiter: RateLimiter, userId: Long, block: () -> T): T {
    rateLimiter.checkLimit(userId)
    val result = block()
    rateLimiter.increaseCount(userId)
    return result
}

 

코드를 간결하게 하기 위해서 리미터 호출 부분은 별도의 최상위 함수로, 람다를 통해 수행 블록을 받도록 만들었다.

 

 

FileController

@Tag(name = "파일")
@RestController
@RequestMapping("/common/files")
class FileController(
    private val fileManager: FileManager,
    @Qualifier("fileUploadRateLimiter") private val rateLimiter: RateLimiter
) : FileApiSpec {
    @AuthenticationRequired
    @PostMapping(
        "/upload",
        consumes = [MediaType.MULTIPART_FORM_DATA_VALUE],
        produces = [MediaType.APPLICATION_JSON_VALUE],
    )
    override fun uploadFile(
        @RequestPart file: MultipartFile,
    ): FileMetaResponse {
        val loginUser = AuthenticationUtils.getCurrentLoginUser()
        return rateLimit(rateLimiter, loginUser.id) {
            fileManager.uploadFile(loginUser.id, file)
        }
    }
}

 

컨트롤러에서 앞서 정의한 fileUploadRateLimiter를 주입받는다.

 

rate limit 체크는 비즈니스 로직과는 거리가 멀다고 생각해서 컨트롤러단 로직으로 분리했다.

 

앞서 정의한 rateLimit 최상위 함수를 통해 리미터와 userId를 넘겨서 체크한다.

 

이제 호출해 보면

 

정상적으로 동작하는 걸 확인할 수 있다.

 

전체 코드는 해당 커밋을 보면 확인 할 수 있다.

https://github.com/mopil/youth-moa-server/commit/ea49f8edf67cf7401d569ff050768af8fbffd2db

feat: file upload rate limit 추가 · mopil/youth-moa-server@ea49f8e

 

 

# 한계 & 후기

복수 인스턴스의 서버 환경에서는 userId별 5회 제한을 할 수 없는 구조여서 그런 경우는 외부 중앙 저장소 (ex. Redis)가 필요하다.

 

당연히 로컬 캐시로 관리하기 때문에 트래픽을 많이 받는 환경이면 캐시로 인한 OOM에 신경 써야 한다.

 

서버를 재기동하면 카운트가 초기화된다.

 

토큰 버킷처럼 특정 시간이 지날 때마다 초기화가 아닌 토큰을 1개 추가하는 기능 구현이 제약된다.

(사실 스케줄링 같은 거 쓰면 할 순 있을 것 같은데 그냥 라이브러리 쓰는 게 편할 것 같다)

 

사이드 프로젝트 할 때 혹시 모를 AWS 과금을 방지용으로, 혹은 단순 rate limit 기능을 빠르게 구현하기 위해서라면 해당 방법도 나쁘지 않은 것 같다.

 

다음에는 더 많은 기능이 있는 Bucket4j로 한 번 만들어봐야겠다.