Bucket4j를 사용해서 스프링부트 트래픽 제한하기

 

 

트래픽 제한을 고려하게 된 이유

트래픽 제한에 대한 고민은 gpt-3.5-turbo 모델을 Fine-Tuning하여 사용하는 졸업작품 코리를 개발하면서 시작되었다.

코리는 초기 유저들을 모으기 위해 한동안은 유료인 GPT API 사용 비용을 운영측에서 부담하는 구조로 운영하게 되었고, 아직은 BM(Business Model)이 딱히 존재하지 않았기 때문에 당장은 손해로 지어질 수 밖에 없는 구조였다.

결론적으로 사용자들의 GPT 모델 사용 비용을 운영 측에서 부담하게 되었으니, 유저들의 요청 횟수에 제한을 두어야한다는 생각이 들게 되었다.

 

 

---

 

Java/SpringBoot 환경에서 트래픽 제한 구현하기

 

토큰 형태의 트래픽 제한 라이브러리 Bucket4j

회원들의 트래픽을 제한할 수 있는 방법을 찾던 와중에 Bucket4j라는 라이브러리를 발견하게 되었다.

• Token Bucket 알고리즘이 적용된 해당 라이브러리는 회원마다 일정 개수의 Token이 담긴 Bucket을 쥐어주고 Token을 사용하여 요청을 보낼 수 있도록 해준다.
• 만약 Bucket에 Token이 존재하지 않는다면 429 Too Many Requests를 return 하여 요청을 제한할 수 있다.
• 유저의 Bucket에 Token이 존재한다면 토큰을 사용하고 요청을 수행한다.
• Bucket에 담긴 Token은 설정한 시간마다 설정한 갯수만큼 충전된다.

 

환경

• Java 17
• SpringBoot 3.0.6

 

의존성 추가

// Bucket4j
implementation group: 'com.github.vladimir-bukhtoyarov', name: 'bucket4j-core', version: '7.0.0'

 

 

---

 

RateLimitBucketProvider

@Component
@RequiredArgsConstructor
public class RateLimitBucketProvider {
    private Map<Long, Bucket> buckets = new ConcurrentHashMap<>();

    public Bucket getBucket(Member member) {
        long memberId = member.getId();
        if (!buckets.containsKey(memberId)) {
            Bucket memberBucket = createBucketByRole(member.getRole());
            buckets.put(memberId, memberBucket);

            return memberBucket;
        }
        return buckets.get(memberId);
    }

    private Bucket createBucketByRole(MemberRole role) {
        Bandwidth memberLimit = BucketPlan.findByMemberRole(role);
        return Bucket.builder()
                .addLimit(memberLimit)
                .build();
    }
}

 

회원의 id를 통해 Map에서 userId에 적합한 Bucket을 찾아 return해주는 객체이다.

만약 id에 해당하는 Bucket이 존재하지 않는다면 유저의 권한을 담은 객체인 MemberRole에 따라 Bucket을 생성하여 return해준다.

 

 

concurrentHashMap을 통한 유저 Bucket 관리

버킷에 대한 정보는 메모리 내에서 Key(UserId), Value(Bucket) 기반의 ConcurrentHashMap에서 관리하도록 하였다.

처음 구현할 때는, 일반적인 HashMap을 사용하였는데 ConcurrentHashMap을 사용하면 Thread Safety하게 데이터에 접근 및 제어가 가능했기 때문에 변경하였다.

 

 

왜 ConcurrentHashMap을 사용했을까?

ConcurrentHashMap은 HashMap의 멀티스레드 환경에서의 안전한 버전이다.

Map 내에 있는 하나의 데이터에 대해 두 Thread에서 동시에 수정 요청을 보내면 먼저 요청한 측에 Lock을 주고 수정이 끝나면 Lock을 반환, 그리고 그 다음 순서로 요청을 보낸 측에서 Lock을 흭득하고 데이터를 수정을 하는 방식으로 동시에 데이터가 수정되어 문제가 발생되는 일을 방지한다.

또한 읽기 연산에서 Lock을 사용하지 않기 때문에 여러 쓰레드에서 동시에 읽기 연산이 가능하여 읽기 연산이 많은 환경의 경우 성능 향상을 기대할 수 있다.

Map의 데이터를 수정할 일이 많진 않아 처음 말한 이점이 적용되어 이점을 얻을 일은 많이 없겠지만, 유저가 로직을 보낼 때마다 Map에서 Bucket 정보를 읽어와야 해야하기 때문에 읽기 연산이 많은 환경이라고 판단이 되었고, 때문에 읽기에 대한 Lock이 걸리지 않는 다는 점이 성능 향상에 도움을 줄 수 있을 것이라고 생각되어 ConcurrentHashMap을 적용하였다.

 

 

---

 

BucketPlan

public enum BucketPlan {
    // 총 5개의 토큰, 1분마다 5개의 토큰 충전
    MEMBER {
        public Bandwidth getLimit() {
            return Bandwidth.classic(5, Refill.intervally(5, Duration.ofMinutes(1)));
        }
    };

    public abstract Bandwidth getLimit();

    public static Bandwidth findByMemberRole(MemberRole memberRole) {
        return Arrays.stream(values())
                .filter(v -> v.name().equals(memberRole.name()))
                .findFirst()
                .orElseThrow(() -> new IllegalArgumentException(ErrorCode.NOT_EXIST_MEMBER_ROLE))
                .getLimit();
    }
}

 

BucketPlan객체는 회원에 권한에 다른 버킷/토큰 정보를 담고 있는 Enum객체이다.

MemberRole과 똑같이 BucketPlan Enum의 열거 상수명을 네이밍하여, MemberRole을 통해 Role에 따른 버킷을 조회할 수 있도록 하였다.

 

Bandwidth

enum에 담겨있는 Bandwidth는 버킷에 대한 설정 객체이며, 다음과 같이 2가지 매개변수를 가진다.

• Bucket이 가질 수 있는 최대 토큰 개수
• Token Refill 규칙

 

현재 설정을 통해 role이 MEMBER이면 최대 5개의 토큰을 가질 수 있으며, 1분마다 5개의 토큰이 채워진다는 것을 알 수 있다.

 

 

---

 

RateLimitFilter

@Component
@Slf4j
@RequiredArgsConstructor
public class RateLimitFilter implements Filter {
    private final JwtUtils jwtUtils;
    private final RateLimitBucketProvider bucketProvider;
    private final MemberRepository memberRepository;

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        String authorizationHeader = httpRequest.getHeader("Authorization");

        // authorizationHeader가 없는 경우 Filter 동작 생략
        if(authorizationHeader == null || !authorizationHeader.startsWith("Bearer ")) {
            chain.doFilter(request, response);
            return;
        }

        //authorizationHeader가 있는 경우 JWT AccessToken이 정상적인 경우 token을 사용
        String accessToken = authorizationHeader.substring("Bearer ".length());
        if (jwtUtils.validateToken(accessToken)) {
            Long memberId = jwtUtils.getMemberIdFromJwt(accessToken);
            Member member = MemberServiceUtils.findMemberById(memberRepository, memberId);
            //관리자 계정이면 Token 사용X
            if(member.getRole() == MemberRole.ADMIN) {
                log.info("MemberRole이 ADMIN이기 때문에 토큰이 소모되지 않았습니다.");
                chain.doFilter(request, response);
                return;
            }

            Bucket memberBucket = bucketProvider.getBucket(member);
            if (!memberBucket.tryConsume(1)) {
                log.error(String.format("회원ID(%d)의 잔여 토큰이 존재하지 않습니다.", memberId));
                throw new RateLimitException(ErrorCode.TOKEN_NOT_EXIST_EXCEPTION);
            }
            log.info(String.format("회원ID(%d)의 토큰 사용 : 잔여 토큰(%d)", memberId, memberBucket.getAvailableTokens()));
        }

        chain.doFilter(request, response);
    }
}

 

해당 클래스가 지금까지 만든 버킷 로직을 통해 트래픽 제한을 적용시켜줄 Filter이며, Filter 인터페이스를 implements 하여 구현할 수 있다.

 

들어온 요청에 대해 Authorization Header를 검사하여 만약 Header가 존재하지 않으면 doFilter를 통해 다음 필터로 요청을 넘긴다.

권한에 대한 검사는 이후 Intercepter에서 처리하도록 구현했기 때문에 여기서는 넘겨도 괜찮다.

 

만약 Authorization 헤더가 존재하면 Authorization 헤더에 담겨있는 JWT 토큰을 검사 후, 토큰을 통해 memberId를 흭득하여 member에 대한 정보를 얻는다.

• 여기서 만약 관리자 권한(ADMIN)의 유저이면 토큰을 사용하지 않는다.

 

그리고 bucketProvider를 통해 Bucket을 흭득하여 토큰을 소모하고 다음 Filter로 요청을 전달하는데, 사용 가능한 token이 존재하지 않으면 잔여 토큰이 존재하지 않는다는 예외를 발생시킨다.

 

 

---

 

ExceptionHandlerFilter

@Component
@Order(Integer.MIN_VALUE)
@Slf4j
public class ExceptionHandlerFilter implements Filter {
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        try {
            chain.doFilter(request, response);
        } catch (RateLimitException | IllegalArgumentException exception) {
            setErrorResponse((HttpServletResponse) response, exception);
        }
    }

    private void setErrorResponse(HttpServletResponse response, CustomException exception) throws IOException {
        ApiResponse<Object> errorResponse = ApiResponse.error(exception.getErrorCode());

        response.getWriter().write(objectMapper.writeValueAsString(errorResponse));
        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        response.setStatus(exception.getErrorCode().getHttpStatusCode());
    }
}

 

보통 Rest API 서버를 개발할 때 예외를 핸들링은 @RestControllerAdvice 애노테이션을 적용한 ExceptionController를 개발하여 수행한다.

 

하지만 @RestControllerAdvice 애노테이션을 통해 캐치할 수 있는 예외는 컨트롤러 단에서 발생한 예외뿐이다. Filter는 DispatcherServlet에 요청이 도달하기 전 실행되는 영역이므로 ExceptionController에서 처리가 불가능하다.

때문에 Filter에서 발생한 예외처리를 담당하는 ExceptionHandlerFilter를 작성해주었다.

 

 

@Order 애노테이션을 통한 Filter를 최상단에 위치시키기

해당 필터를 최상단에 위치시키고 try/catch 를 통해 다음 필터들에서 발생하는 예외를 캐치하도록 하였다.

그리고 캐치된 예외는 setErrorResponse메서드를 통해 response되어 사용자에게 전달되도록 구현하였다.

 

 

---

 

작동 결과

 

유저가 1분 내에 5번 연속으로 API를 호출하자 예측한대로 429 Too Many Requests가 발생하여 23ms라는 빠른 속도로 요청을 거부하였다.

 

 

---

 

성능 개선

사람들마다 구현하는 방식이 다르겠지만 본인은 SpringBoot의 Filter단에서 요청에 대한 Token 개수를 검사하도록 구현하였다.

처음에 인터넷에서 찾은 예시에는 Interceptor에서 토큰 개수 검사를 수행하도록 구현하고 있었고 본인도 처음에는 그 예제를 따라 코드를 작성했다.

하지만 현재 프로젝트는 전체적인 트래픽 조절을 위해 모든 API에 대해서 토큰 검사를 수행하고 있기 때문에 Handler의 정보가 필요하지 않았고 때문에 Interceptor가 아닌, Filter에서 검사해야 불필요한 Controller Mapping등의 작업을 수행하지 않기 때문에 성능 면에서 가장 최적화된 방법이라고 생각했다.

 

개선 결과

 

실제로 비교를 해본 결과 예측한대로 Filter와 Interceptor의 차이는 386msec만큼 차이가 났다. 작은 수 같지만 이런 작은 차이가 쌓여 성능을 낮추기 때문에, 이런 것들을 하나하나 줄여나가는 것이 개발의 재미이자 덕목이라고 생각한다.

 

✅ 만약 이후에 모든 API가 아닌 특정 API에 대해서만 요청을 제한하도록 변경한다면 다시 Interceptor에서 토큰 검사를 수행하거나, 다른 방안을 생각해볼 것이다.