Rest API 호출 제한 - Bucket4J

배경

---

특정 사용자가 서비스에서 제공되는 검색 기능을 과도하게 이용해서 DDoS 적인 현상을 발생시켰고, 이를 방지하고자 몇 가지 방안을 마련하는 과정에서, 몇몇 Rest API에 대해 사용 제한 설정을 위해 알아본 내용을 정리한 내용이다.

Rate Limit Algorithm

---

위와 같이 클라이언트로부터의 과도한 사용에 대해 서비스의 가용성을 안정적으로 유지하기 위해 트래픽을 일정 수준으로 조정하는 수단으로 언급되는 것이 Rate Limit Algorithm 이다. 특정 IP, 특정 Access-key 등에 대해 일정 횟수로 사용 제한을 하거나, 사용 시간 제한 등을 적용해서 사용 제한을 넘는 요청에 대해 거부를 하거나, 혹은 일정 시간 대기, 또는 요금을 받고 제한을 풀어주는 Business Model을 적용할 경우에 사용하게 된다.

Bucket4j

---

Bucket4j는 Token Bucket 방식으로 구현된 Rate Limit Algorithm중 하나이다. 간단하게 아래와 같이 관련 내용을 정리할 수 있다.

• Token Bucket Algorithm
  • Token이 담겨진 Bucket을 가정하고, 요청을 처리하는 비용으로 Token을 지불하는 방식으로 처리에 제한을 설정한 Algorithm
  • Bucket에 남겨진 Token이 부족하면 요청은 Reject 처리
  • Bucket에 Token은 시간을 기반으로 다시 채워지게 설계되어 있다.
  • 일정 수준의 부하 처리도 Bucket에 남은 Token을 기반으로 해서 일정 수준으로 처리 가능
• Bucket4j
  • Java 기반의 Token Bucket Algorithm 구현 Library
  • Standalone, Cluster 환경 모두에서 사용 가능
  • in-memory, 분산 Cache 모두 지원

기본 사용 방법

---

의존 설정

현재(2022.02.01) 시점으로 7.1.0 버전이 최신 버전으로 Release 되어 있고, Gradle Project의 경우 아래와 같이 의존을 추가하면 된다.

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

주요 Class 및 Interface

• Bucket Interface : 트래픽 처리를 위해 사용하는 Token 보관용 Bucket을 제어하기 위한 Interface.
• Bandwidth class : Bucket을 생성하는데 필요한 핵심 개념. Bucket이 보유할 수 있는 최대한의 Toekn갯수와 토큰이 생성되는 비율을 정의하는 class.
• Refill class : Bucket에 다시 Token이 Refill되는 주기와 갯수를 정의하는 Class

Refill class 생성 예

Refill class에서 제공하는 static function인 intervally(), greedy()을 통해 Instance를 생성하고 있고 아래가 그 예이다.

...
    private static final int TOKEN_REFILL_INTERVAL_SECONDS = 10;
    private static final int TOKEN_REFILL_COUNT_AT_ONCE = 3;

    private static final int TOKEN_REFILL_DURATION_MINUTES = 1;
    private static final int TOKEN_REFILL_COUNT = 5;

    // 1. intervally() 사용 Case
    private Refill getIntervalRefill() {
        return Refill.intervally(TOKEN_REFILL_COUNT_AT_ONCE, Duration.ofSeconds(TOKEN_REFILL_INTERVAL_SECONDS));
    }

    // 2. greedy() 사용 Case
    private Refill getGreedyRefill() {
        return Refill.greedy(TOKEN_REFILL_COUNT, Duration.ofMinutes(TOKEN_REFILL_DURATION_MINUTES));
    }
...

1. Intervally() : 설정된 시간에 설정한 Toekn의 갯수만큼 새로 생성된다.
   • 위의 예에서는 20초 지난 후 3개의 Token이 한 번에 생성된다.
2. greedy() : (설정된 시간 / 설정한 Token의 갯수) 만큼 시간이 지날 때 마다 Token이 1개씩 생성된다.
   • 위의 예에서는 (1분 / 5) 이기 때문에 12초마다 Token이 하나씩 생성되게 된다.

Bandwidth class 생성 예

Bandwidth class에서 제공하는 static function인 simple(), classic()을 통해 Instance를 생성하고 있고 아래가 그 예이다.

...
    private static final int MAX_BANDWIDTH = 10;
    private static final int TOKEN_REFILL_DURATION_MINUTES = 1;

    // 1. simple() 사용 case
    private Bandwidth getSimpleBandwidth() {
        return Bandwidth.simple(MAX_BANDWIDTH, Duration.ofMinutes(TOKEN_REFILL_DURATION_MINUTES));
    }

    // 2. classic() 사용 case
    private Bandwidth getClassicBandwidth(Refill refill) {
        return Bandwidth.classic(MAX_BANDWIDTH, refill);
    }
}

1. simple() : (설정된 시간 / 설정된 Bandwidth의 크기) 만큼 시간이 지날 때 마다 Token이 1개씩 생성된다. - 위의 예에서는 (1분 / 10) 이기 때문에 6초마다 Token이 하나씩 생성되게 된다.
2. classic() : 인자로 정의된 Refill 방식에 따라 설정된 Bandwidth 만큼 Token이 생성되게 된다.

Bucket 생성 예

Bucket은 static function으로 정의된 builder()에 의해서 Instance가 생성되어지고, 이 때 Bandwidth를 추가해서 Token 갯수와 Token이 채워지는 방식을 정의하게 된다. (7.0.0 이전까지는 Bucket4j class에서 제공하는 builder() 함수에 의해 Instance를 생성했지만, 7.0.0 부터 Deprecated 되었고, Bucket interface에서 static function을 제공한다.)

... 
    // 단일 Bandwidth를 적용해서 Bucket instance를 생성
    private Bucket generateSimpleBucket() {
        return Bucket.builder()
                .addLimit(getSimpleBandwidth())
                .build();
    }

    // 복수의 Bandwidth를 적용해서 Bucket instance를 생성
    private Bucket generateComplexBucket(List<Bandwidth> bandwidthList) {
        LocalBucketBuilder bucketBuilder = Bucket.builder();

        for(Bandwidth bandwidth : bandwidthList) {
            bucketBuilder.addLimit(bandwidth);
        }

        return bucketBuilder.build();
    }
...

사용 예

---

1. Simple Bandwith - Bucket 사용 예

GET /member/count API 추가

현재 회원 수를 조회하는 Rest API를 추가하고, 위에서 정의한 SImple Bucket을 적용해서 분당 10회 이상 호출 될 때 HTTP Status 429 (Too Many Request)를 Return 하게 설정했다.

// Service
public class MemberService {
...
    public int getAllFriendsCount() {
        // 회원 수를 Return 한다
    }
}

--------------------------------

// Controller
@RestController
@Slf4j
public class MemberController {

    private Bucket    simpleBucket;

    public MemberController() {
        simpleBucket = generateSimpleBucket();
    }

    @GetMapping("/member/count")
    public    ResponseEntity<Integer>    getTotalMemberCount() {

        if (simpleBucket.tryConsume(TOKEN_CONSUME_COUNT)) {
            log.info(">>> Remain bucket Count : {}", simpleBucket.getAvailableTokens()); 
            return    ResponseEntity.ok(memberService.getAllFriendsCount());
        }

        log.warn(">>> Exhausted Limit in Simple Bucket");
        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).build();
    }
    ...
}

제일 주목해야 할 구문은 simpleBucket.tryConsume(TOKEN_CONSUME_COUNT) 이고, simpleBucket에서 TOKEN_BUCKET_COUNT 만큼 사용할 만큼 남아있으면 true가 되면서 정상적으로 memberService.getAllFriendsCount()가 호출되고, Buckeket에 Token이 남아있지 않는 경우 Http Status 429 (Too Many Request)가 Response되게 구현되었다.

Test 1. 10회 정상호출 후, 호출 제한에 걸리는 Case

아래의 Test Code는 10회 호출 할 때 까지는 회원 수가 정상적으로 Response로 오다, 11회째 호출할 때 Http Status 429가 오는 경우를 Test 한 것이다.

@Slf4j
@WebMvcTest(MemberController.class)
class MemberControllerTest {

    private final static int    MAX_BANDWIDTH = 10;
    private final static int    FAKE_MEMBER_COUNT = 10;

    @Autowired
    private MockMvc mockMvc;

    @MockBean
    private MemberService    memberService;


    @Test
    @Order(1)
    void testGetTotalMemberCount() {
        // Given
        when(memberService.getAllFriendsCount()).thenReturn(FAKE_MEMBER_COUNT);

        // When & Then        
        try {

            log.info("--- Call API with remained bucket");
            for(int i=0;i < MAX_BANDWIDTH;i++) {
                mockMvc.perform(get("/member/count"))        
                .andExpect(status().isOk())
                .andExpect(content().string(String.valueOf(FAKE_MEMBER_COUNT)));                
            }

            log.info("--- Call API with exausted bucket");
            mockMvc.perform(get("/member/count"))        
            .andExpect(status().isTooManyRequests());    

        } catch (Exception e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }            
    }
}

# Result
Test Case 1
MemberControllerTest    : --- Call API with remained bucket
MemberController        : >>> Remain bucket Count : 9
MemberController        : >>> Remain bucket Count : 8
MemberController        : >>> Remain bucket Count : 7
MemberController        : >>> Remain bucket Count : 6
MemberController        : >>> Remain bucket Count : 5
MemberController        : >>> Remain bucket Count : 4
MemberController        : >>> Remain bucket Count : 3
MemberController        : >>> Remain bucket Count : 2
MemberController        : >>> Remain bucket Count : 1
MemberController        : >>> Remain bucket Count : 0
MemberControllerTest    : --- Call API with exausted bucket
MemberController        : >>> Exhausted Limit in Simple Bucket

Test 2. Test 1 이후 Bucket에 Token이 1개 추가되길 기다린 후에 재호출하는 경우

    @Test
    @Order(2)
    void testGetTotalMemberCount_WithRefilledSimpleBucket() {
        // Given
        when(memberService.getAllFriendsCount()).thenReturn(FAKE_MEMBER_COUNT);

        // When & Then        
        try {

            log.info("--- Call API again after refill 1 token");
            Thread.sleep(60 / MAX_BANDWIDTH * 1000 * 1);
            mockMvc.perform(get("/member/count"))        
            .andExpect(status().isOk())
            .andExpect(content().string(String.valueOf(FAKE_MEMBER_COUNT)));    


        } catch (Exception e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }

    }

# Result
Test Case 2
MemberControllerTest    : --- Call API with remained bucket
MemberController        : >>> Remain bucket Count : 9
MemberController        : >>> Remain bucket Count : 8
MemberController        : >>> Remain bucket Count : 7
MemberController        : >>> Remain bucket Count : 6
MemberController        : >>> Remain bucket Count : 5
MemberController        : >>> Remain bucket Count : 4
MemberController        : >>> Remain bucket Count : 3
MemberController        : >>> Remain bucket Count : 2
MemberController        : >>> Remain bucket Count : 1
MemberController        : >>> Remain bucket Count : 0
MemberControllerTest    : --- Call API with exausted bucket
MemberController        : >>> Exhausted Limit in Simple Bucket
<-- 약 6초 대기 -->
MemberControllerTest    : --- Call API again after refill 1 token
MemberController        : >>> Remain bucket Count : 0

2. Classic Bandwith - Bucket 사용 예

Classic Bandwidth를 이용해서 Bucket을 생성할 때는 2가지 Refill 방식을 적용할 수 있기 때문에, 2가지 방식을 각각 적용할 Rest API 2개를 추가하였다.

GET /member/name/{name} API 추가

회원명으로 회원 정보 List를 조회하는 Rest API를 추가하고, 위에서 1분 동안 5개의 Toekn이 추가되는 Classic Bandwith로 생성한 Bucket을 적용했고, Token이 다 소모된 상태에서 호출 될 때 HTTP Status 429 (Too Many Request)를 Return 하게 설정했다.

GET /member/id/{id} API 추가

회원 ID로 회원 정보를 조회하는 Rest API를 추가하고, 위에서 20초마다 3개의 Toekn이 추가되는 Classic Bandwith로 생성한 Bucket을 적용했고, Token이 다 소모된 상태에서 호출 될 때 HTTP Status 429 (Too Many Request)를 Return 하게 설정했다.

// Service
public class MemberService {
...
    public List<MemberDto> getFriendInfoListByName(String name) {
        // 회원명으로 회원 정보 List를 Return 한다
        ...
    }

    public MemberDto getFriendInfoById(String id) {
        // 회원 ID로 회원 정보를 Return 한다
        ...
    }
}

--------------------------------

// Controller
@RestController
@Slf4j
public class MemberController {

    private static final int TOKEN_REFILL_INTERVAL_SECONDS = 20;
    private static final int TOKEN_REFILL_COUNT_AT_ONCE = 3;

    private static final int TOKEN_REFILL_DURATION_MINUTES = 1;
    private static final int TOKEN_REFILL_COUNT = 5;

    private Bucket    complexGreedyRefillBucket;
    private Bucket    complexIntervalRefillBucket;

    public MemberController() {
        ...
        complexGreedyRefillBucket = generateComplexBucket(Arrays.asList(getClassicBandwidth(getGreedyRefill())));
        complexIntervalRefillBucket = generateComplexBucket(Arrays.asList(getClassicBandwidth(getIntervalRefill()))); 
    }
    ...
    // Greddy refill 방식으로 생성된 Bucket을 사용하는 API
    @GetMapping("/member/name/{name}")
    public ResponseEntity<List<MemberDto>> getMemberListByName(@PathVariable("name") String name) {

        if (complexGreedyRefillBucket.tryConsume(TOKEN_CONSUME_COUNT)) {
            log.info(">>> Remain bucket Count : {}", complexGreedyRefillBucket.getAvailableTokens()); 
            return    ResponseEntity.ok(memberService.getFriendInfoListByName(name));            
        }
        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).build();
    }

    // Interval Refill 방식으로 생성된 Bucket을 사용하는 API
    @GetMapping("/member/id/{id}")
    public ResponseEntity<MemberDto> getMemberById(@PathVariable("id") String id) {
        if (complexIntervalRefillBucket.tryConsume(TOKEN_CONSUME_COUNT)) {
            log.info(">>> Remain bucket Count : {}", complexIntervalRefillBucket.getAvailableTokens()); 
            return    ResponseEntity.ok(memberService.getFriendInfoById(id));            
        }
        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).build();
    }
    ...
    // Create Interval Refill Instance
    private Refill getIntervalRefill() {
        return Refill.intervally(TOKEN_REFILL_COUNT_AT_ONCE, Duration.ofSeconds(TOKEN_REFILL_INTERVAL_SECONDS));
    }

    // Create Greedy Refill Instance
    private Refill getGreedyRefill() {
        return Refill.greedy(TOKEN_REFILL_COUNT, Duration.ofMinutes(TOKEN_REFILL_DURATION_MINUTES));
    }

}

두 개의 API는 각각 다른 방식으로 만들어진 Token Bucket을 사용해서 API 호출을 제한하고 있고, Token이 충원되는 방식이 getIntervalRefill(), getGreedyRefill() 에서 다르게 정의되어 있다.

Test 3. 10회 호출 후 호출 제한이 걸리고, 1개의 Token이 충전된 후에 다시 호출해 성공하는 Test

@Slf4j
@TestMethodOrder(OrderAnnotation.class)
@WebMvcTest(MemberController.class)
class MemberControllerTest {

    private final static int    MAX_BANDWIDTH = 10;

    private static final int GREEDY_TOKEN_REFILL_DURATION_MINUTES = 1;
    private static final int GREEDY_TOKEN_REFILL_COUNT = 5;
    private static final String    COMMON_NAME = "Kim";

    @Autowired
    private MockMvc mockMvc;

    @MockBean
    private MemberService    memberService;

    private    List<MemberDto>    testMemberList = new ArrayList<MemberDto>();

    ...

    @Test
    @Order(3)
    @DisplayName("Case 3. Use Complex Bucket which refill token greedly.")
    void testGetMemberListByName_WithComplexBucket_RefillGrdeedly() {
        // Given
        when(memberService.getFriendInfoListByName(COMMON_NAME)).thenReturn(testMemberList);
        String    targetURL = new StringBuilder("/member/name/").append(COMMON_NAME).toString();

        // When & Then        
        try {

            log.info("--- Call API with Complex bucket - Refilling Bandwidth Greedly");
            for(int i=0;i < MAX_BANDWIDTH;i++) {
                mockMvc.perform(get(targetURL))        
                .andExpect(status().isOk());
                // .andExpect(jsonPath("$.errors", hasSize(3)))                
            }

            log.info("--- Call API with exausted bucket");
            mockMvc.perform(get(targetURL))        
            .andExpect(status().isTooManyRequests());    

            log.info("--- Call API again after refill 1 token");
            Thread.sleep(( GREEDY_TOKEN_REFILL_DURATION_MINUTES * 60 ) / GREEDY_TOKEN_REFILL_COUNT * 1000 * 1);
            mockMvc.perform(get(targetURL))        
            .andExpect(status().isOk())
            .andExpect(jsonPath("$").isArray())
            .andExpect(jsonPath("$", hasSize(3)))
            .andExpect(jsonPath("$.[0].memberName", is(COMMON_NAME)));


        } catch (Exception e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }

    }
}

# Result
Test Case 3
--- Call API with Complex bucket - refill a token gradually
>>> Remain bucket Count : 9
>>> Remain bucket Count : 8
>>> Remain bucket Count : 7
>>> Remain bucket Count : 6
>>> Remain bucket Count : 5
>>> Remain bucket Count : 4
>>> Remain bucket Count : 3
>>> Remain bucket Count : 2
>>> Remain bucket Count : 1
>>> Remain bucket Count : 0
--- Call API with exhausted bucket
<-- 12 초 후에 호출 -->
--- Call API again after refill 1 token
>>> Remain bucket Count : 0

API에 적용된 Bucket은 최초 10개의 Token을 보유하고 있고, 1분에 5개의 Token이 보충됨. 즉 12초마다 1개의 Token이 추가된다. Test에서는 10개의 Token을 다 소진한 상태에서 12초 후에 Token이 하나 보충되서 API 요청을 정상 처리하는 예이다.

Test 4. 10회 호출 후 호출 제한이 걸리고, 정해진 시간 후에 정해진 수만큼의 Token이 추가된 후에 호출하는 Test

@Slf4j
@TestMethodOrder(OrderAnnotation.class)
@WebMvcTest(MemberController.class)
class MemberControllerTest {
    ...
    private static final String    COMMON_NAME = "Kim";

    private static final int INTERVAL_TOKEN_REFILL_DURATION_SECONDS = 20;
    private static final int INTERVAL_TOKEN_REFILL_COUNT = 3;
    private static final String    TEST_ID = "ID001";

    @Autowired
    private MockMvc mockMvc;

    @MockBean
    private MemberService    memberService;

    private    List<MemberDto>    testMemberList = new ArrayList<MemberDto>();

    ...

    @Test
    @Order(4)
    void testgetMemberById_WithComplexBucket_RefillIntervally() {

        log.info("Test Case 4");

        // Given
        when(memberService.getFriendInfoById(TEST_ID)).thenReturn(testMemberList.get(0));
        String    targetURL = new StringBuilder("/member/id/").append(TEST_ID).toString();

        // When & Then        
        try {

            log.info("--- Call API with Complex bucket - Refilling Bandwidth Intervally");
            for(int i=0;i < MAX_BANDWIDTH;i++) {
                mockMvc.perform(get(targetURL))        
                .andExpect(status().isOk());
                // .andExpect(jsonPath("$.errors", hasSize(3)))                
            }

            log.info("--- Call API with exausted bucket");
            mockMvc.perform(get(targetURL))        
            .andExpect(status().isTooManyRequests());    

            log.info("--- Call API again after partial interval period");
            Thread.sleep(INTERVAL_TOKEN_REFILL_DURATION_SECONDS / INTERVAL_TOKEN_REFILL_COUNT * 1000);
            mockMvc.perform(get(targetURL))        
            .andExpect(status().isTooManyRequests());

            log.info("--- Call API again after full interval period");
            Thread.sleep(INTERVAL_TOKEN_REFILL_DURATION_SECONDS * 1000);
            mockMvc.perform(get(targetURL))        
            .andExpect(status().isOk())
            .andExpect(jsonPath("$.memberId", is(TEST_ID)))
            .andExpect(jsonPath("$.memberName", is(testMemberList.get(0).getMemberName())))    
            .andExpect(jsonPath("$.phoneNo", is(testMemberList.get(0).getPhoneNo())));


        } catch (Exception e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }        
    }
}

# Result
Test Case 4
--- Call API with Complex bucket - Refilling tokens at regular intervals
>>> Remain bucket Count : 9
>>> Remain bucket Count : 8
>>> Remain bucket Count : 7
>>> Remain bucket Count : 6

>>> Remain bucket Count : 5
>>> Remain bucket Count : 4
>>> Remain bucket Count : 3
>>> Remain bucket Count : 2
>>> Remain bucket Count : 1
>>> Remain bucket Count : 0
--- Call API with exhausted bucket
<-- 6초 경과 : Greedy Refill 방식으로 1개의 Token이 충전되는 시간 -->
--- Call API again after partial interval period
<-- 20초 경과 -->
--- Call API again after full interval period
>>> Remain bucket Count : 2

Greedy Refill 방식과는 다르게 부분적으로 Token이 충전되지 않고, 20초로 정해진 Interval이 지난 후에 처음 정의된 Token 3개가 충전되서 API 요청을 정상 처리하고 나머지 Token이 남아있는 예이다.

기타

---

• 예시 Code
• 참조
  • https://github.com/bucket4j/bucket4j
  • https://www.baeldung.com/spring-bucket4j
  • https://www.mimul.com/blog/about-rate-limit-algorithm/