Rate Limit Algorithm

정명진·2023년 2월 3일
0
post-thumbnail

이번에 사이드 프로젝트를 진행하며 간편로그인 기능을 구현하게 되었다. 해당 인증 과정을 프론트단에서 처리후 유저 정보만 받아서 회원가입을 진행할지와 클라이언트가 간편로그인을 요청하면 서버와 플랫폼 인증서버간 통신을 통해 회원가입을 진행할지 2가지 방안에 대해 고민을 하였다.

프론트 개발자분이 프론트에서 정보만 넘겨서 회원가입을 요청하면 엔드포인트가 만약 노출된다면 다른 사람이 악의적으로 자신의 정보로 회원가입을 할 수 있지 않냐는 이의제기를 통해 서버와 플랫폼 인증서버간 통신을 통해 구현하는 방안으로 결정하였다. 하지만 이 방식으로 구현한다면 서버와 플랫폼 인증서버간 통신 비용이 발생하기 때문에 누군가 악의적인 DDoS 공격을 한다면 방어책이 필요하다고 생각하였다. 그래서 API 요청수를 제한하는 방법을 찾다가 Rate Limit Algorithm을 알게 되었다. 요청을 제한하는 알고리즘으로 이해하면 되는데 그중 Token Bucket 방식을 쓰기로 결정하였다.

Token Bucket 방식이란?

순간적으로 많은 트래픽이 와도 토큰이 있다면 요청을 처리하고 토큰 손실 처리를 통해 평균 처리 속도를 제한합니다. 평균 유입 속도를 제한하고 처리 패킷의 손실없이 특정 수준의 BURST 요청을 허용할 수 있습니다.

원리는 다음과 같습니다.

  • 토큰을 정해진 비율로 토큰 버킷에 넣는다.
  • 버킷은 최대 n개의 토큰 보관소 사이즈를 갖습니다. 버킷이 다 차면 새로 추가된 토큰은 삭제되거나 거부됩니다.
  • 요청이 들어오면 큐에 들어가며 요청을 처리하기 전에 버킷의 토큰을 획득해야 하며, 토큰을 보유한 후에 요청이 처리되며 처리된 후에는 사용한 토큰을 삭제합니다.
  • 토큰 버킷은 토큰이 배치되는 속도를 기반으로 액세스 속도를 제어합니다.
  • 전송 횟수를 누적할 수 있으며, 버킷이 가득차면 패킷 손실 없이 토큰이 손실됩니다.

알고리즘 코드는 다음과 같습니다.

public class TokenBucket extends RateLimiter {
  private int tokens; // 토큰의 개수
  private int capacity; // 버킷 사이즈
  private long lastRefillTime; // 마지막 리필 시간

  public TokenBucket(int maxRequestPerSec) {
    super(maxRequestPerSec);
    this.tokens = maxRequestPerSec;
    this.capacity = maxRequestPerSec;
    this.lastRefillTime = scaledTime();
  }

  @Override
  public boolean allow() {
    synchronized (this) {
      refillTokens(); // 토큰 리필할 시간이 되면 리필
      if (this.tokens == 0) {
        return false;
      }
      this.tokens--; // 사용 토큰 제거
      return true;
    }
  }

  private void refillTokens() {
    final long now = scaledTime();
    if (now > this.lastRefillTime) { // 만약 리필할 시간이 되면 리필
      final double elapsedTime = (now - this.lastRefillTime);
      int refill = (int) (elapsedTime * this.maxRequestPerSec);
      this.tokens = Math.min(this.tokens + refill, this.capacity);
      this.lastRefillTime = now;
    }
  }

  private long scaledTime() {
    return System.currentTimeMillis() / 1000;
  }
}

하지만 이미 기존에 검증된 Bucket4j 라는 라이브러리가 존재하므로 해당 라이브러리를 사용해 실제 서비스에 적용해 보겠습니다.

build.gradle에 다음을 추가하고 프로젝트를 빌드합니다.

implementation 'com.github.vladimir-bukhtoyarov:bucket4j-core:7.1.0'

버킷은 다음과 같이 생성이 가능합니다.

Bandwidth minLimit = Bandwidth.classic(5, Refill.intervally(5, Duration.ofMinutes(1)));

Bucket bucket = Bucket.builder()
				.addLimit(minLimit)
				.build();

앞에서 부터 Capacity, tokens, minutes 값입니다.
Capacity란 버킷의 사이즈, token는 재생성할 토큰의 갯수 minutes는 1분을 의미합니다.

즉 해당 버킷의 의미는 1분에 5번 요청이 가능한 규칙을 갖고 리필시 5개의 토큰을 재생성하는 버킷을 의미합니다.

이제 이를 적용하여 실제 서비스에 적용하면 다음과 같이 사용이 가능합니다.

	@Hidden
    @GetMapping("/kakao")
    public ResponseEntity<TokenDto> loginByKakao(@RequestParam String code){

        Bucket bucket = rateLimitService.resolveBucket("OAUTH");
        ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);

        long remainingTokens = probe.getRemainingTokens();

        if (probe.isConsumed()) { // 토큰 사용했는지 확인 그래야 호출 가능
            TokenDto tokenDto = kaKaoUseCase.login(code);
            return ResponseEntity.ok(tokenDto);
        }

        long remainTimeForRefill = probe.getNanosToWaitForRefill() / 1000000000;

        log.error("TOO_MANY_REQUEST");
        log.error("Available Token : {}", remainingTokens);
        log.error("Wait Time {} Seconds ", remainTimeForRefill);

        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).build();
    }

OAUTH 규칙의 bucket을 가져온후 토큰을 소모합니다. 만약 토큰 사용이 가능하다면 유효 횟수가 남았다는 뜻이므로 카카오 로그인을 진행합니다. 하지만 더이상 사용 가능한 토큰이 없을 경우 토큰 재생성에 얼마나 시간이 남았는지 로그로 남기며 429 TOO_MANY_REQUEST를 리턴합니다. 오늘도 새로운 내용을 공부하게된 유익한 시간이었습니다.

profile
개발자로 입사했지만 정체성을 잃어가는중... 다시 준비 시작이다..

0개의 댓글