토큰 버킷 알고리즘 기반 처리율 제한기 구현

Author: gunkim

build.gradle.kts

@@ -46,5 +46,8 @@ subprojects {
         testImplementation("org.junit.jupiter:junit-jupiter")
         testImplementation("org.mockito:mockito-core:5.14.2")
         testImplementation("org.assertj:assertj-core:3.26.3")
+
+        implementation("org.slf4j:slf4j-api:2.0.16")
+        implementation("org.slf4j:slf4j-simple:2.0.16")
     }
 }

rate-limiter/build.gradle.kts

@@ -0,0 +1,29 @@
+package io.github.gunkim.ratelimiter;
+
+import io.github.gunkim.ratelimiter.tokenbucket.TokenBucket;
+
+public class Application {
+    public static void main(String[] args) {
+        //4초에 2개 요청에 대한 처리율 제한 설정
+        final int bucketSize = 20;
+        final long refillRate = 60_000L;
+        var bucket = new TokenBucket(bucketSize, refillRate);
+
+        while (true) {
+            processRequest(bucket);
+            sleep(1_000);
+        }
+    }
+
+    private static void processRequest(TokenBucket bucket) {
+        bucket.request(() -> System.out.println("Request"));
+    }
+
+    private static void sleep(long millis) {
+        try {
+            Thread.sleep(millis);
+        } catch (InterruptedException ex) {
+            Thread.currentThread().interrupt();
+        }
+    }
+}

rate-limiter/src/main/java/io/github/gunkim/ratelimiter/Application.java

@@ -0,0 +1,67 @@
+package io.github.gunkim.ratelimiter.tokenbucket;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class TokenBucket implements AutoCloseable {
+    private static final Logger logger = LoggerFactory.getLogger(TokenBucket.class);
+
+    private final ScheduledExecutorService executorService;
+    private final AtomicInteger tokens;
+
+    public TokenBucket(int bucketSize, long refillRate) {
+        this.executorService = Executors.newSingleThreadScheduledExecutor();
+        this.tokens = new AtomicInteger(bucketSize);
+
+        this.executorService.scheduleAtFixedRate(() -> refillTokens(bucketSize), refillRate, refillRate, TimeUnit.MILLISECONDS);
+
+        logger.info("Token bucket created with size: {} and refill rate: {}", bucketSize, refillRate);
+    }
+
+    public void request(Runnable request) {
+        int currentTokens;
+        do {
+            currentTokens = tokens.get();
+            if (currentTokens <= 0) {
+                logger.debug("Dropped request");
+                return;
+            }
+        } while (!tokens.compareAndSet(currentTokens, currentTokens - 1));
+
+        request.run();
+        logger.debug("Request processed. Remaining tokens: {}", tokens.get());
+    }
+
+    public void shutdown() {
+        executorService.shutdown();
+        try {
+            if (!executorService.awaitTermination(1, TimeUnit.SECONDS)) {
+                executorService.shutdownNow();
+            }
+        } catch (InterruptedException ex) {
+            executorService.shutdownNow();
+            Thread.currentThread().interrupt();
+        }
+    }
+
+    @Override
+    public void close() {
+        shutdown();
+    }
+
+    private void refillTokens(int bucketSize) {
+        int currentTokens;
+        do {
+            currentTokens = tokens.get();
+            if (currentTokens >= bucketSize) {
+                return;
+            }
+        } while (!tokens.compareAndSet(currentTokens, bucketSize));
+        logger.debug("Refilled tokens: {}", tokens.get());
+    }
+}

rate-limiter/src/main/java/io/github/gunkim/ratelimiter/tokenbucket/TokenBucket.java

@@ -0,0 +1 @@
+org.slf4j.simpleLogger.defaultLogLevel=debug

rate-limiter/src/main/resources/simplelogger.properties

@@ -0,0 +1,66 @@
+package io.github.gunkim.ratelimiter.tokenbucket;
+
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+
+@DisplayName("TokenBucket은")
+class TokenBucketTest {
+    @Test
+    void 토큰_수가_0이면_요청_메서드가_요청을_삭제한다() {
+        var tokenBucket = new TokenBucket(0, 20_000);
+        Runnable request = Mockito.mock(Runnable.class);
+        tokenBucket.request(request);
+        Mockito.verify(request, never()).run();
+    }
+
+    @Test
+    void 토큰_수가_최대치일_때_요청_메소드는_요청을_처리한다() {
+        var tokenBucket = new TokenBucket(10, 20_000);
+        Runnable request = Mockito.mock(Runnable.class);
+        for (int i = 0; i < 10; i++) {
+            tokenBucket.request(request);
+        }
+        Mockito.verify(request, times(10)).run();
+    }
+
+    @Test
+    void 토큰_수가_3이고_요청이_6개_들어오면_3개만_처리한다() {
+        var tokenBucket = new TokenBucket(3, 20_000);
+        Runnable request = Mockito.mock(Runnable.class);
+        for (int i = 0; i < 6; i++) {
+            tokenBucket.request(request);
+        }
+        Mockito.verify(request, times(3)).run();
+    }
+
+    @Test
+    void 단일_스레드_환경에서_토큰_버킷이_올바르게_작동한다() {
+        var tokenBucket = new TokenBucket(5, 1000);
+        Runnable request = Mockito.mock(Runnable.class);
+        for (int i = 0; i < 7; i++) {
+            tokenBucket.request(request);
+        }
+        Mockito.verify(request, times(5)).run();
+    }
+
+    @Test
+    void 다중_스레드_환경에서_토큰_버킷이_올바르게_작동한다() throws InterruptedException {
+        var tokenBucket = new TokenBucket(100, 5_000);
+        Runnable request = Mockito.mock(Runnable.class);
+        ExecutorService taskExecutor = Executors.newFixedThreadPool(10);
+        for (int i = 0; i < 150; i++) {
+            taskExecutor.execute(() -> tokenBucket.request(request));
+        }
+        taskExecutor.shutdown();
+        taskExecutor.awaitTermination(1000, TimeUnit.MILLISECONDS);
+        Mockito.verify(request, times(100)).run();
+    }
+}

rate-limiter/src/test/java/io/github/gunkim/ratelimiter/tokenbucket/TokenBucketTest.java

@@ -23,3 +23,4 @@ include("string-calculator")
 include("blackjack")
 include("banking")
 include("number-baseball")
+include("rate-limiter")