perf(s3stream): optimize limiter lock range (#2998)

Signed-off-by: Robin Han <[email protected]>

Author: superhx

s3stream/src/main/java/com/automq/stream/s3/metrics/stats/NetworkStats.java

@@ -35,7 +35,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 public class NetworkStats {
-    private static volatile NetworkStats instance = null;
+    private static final NetworkStats INSTANCE = new NetworkStats();
     // <StreamId, <FastReadBytes, SlowReadBytes>>
     private final Map<Long, Pair<Counter, Counter>> streamReadBytesStats = new ConcurrentHashMap<>();
     private final Counter networkInboundUsageTotal = new Counter();
@@ -49,20 +49,20 @@ private NetworkStats() {
     }
 
     public static NetworkStats getInstance() {
-        if (instance == null) {
-            synchronized (NetworkStats.class) {
-                if (instance == null) {
-                    instance = new NetworkStats();
-                }
-            }
-        }
-        return instance;
+        return INSTANCE;
     }
 
     public CounterMetric networkUsageTotalStats(AsyncNetworkBandwidthLimiter.Type type, ThrottleStrategy strategy) {
-        return type == AsyncNetworkBandwidthLimiter.Type.INBOUND
-                ? networkInboundUsageTotalStats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkInboundUsageMetric(strategy, networkInboundUsageTotal::inc))
-                : networkOutboundUsageTotalStats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkOutboundUsageMetric(strategy, networkOutboundUsageTotal::inc));
+        Map<ThrottleStrategy, CounterMetric> stats = type == AsyncNetworkBandwidthLimiter.Type.INBOUND ? networkInboundUsageTotalStats : networkOutboundUsageTotalStats;
+        CounterMetric metric = stats.get(strategy);
+        if (metric == null) {
+            if (type == AsyncNetworkBandwidthLimiter.Type.INBOUND) {
+                metric = stats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkInboundUsageMetric(strategy, networkInboundUsageTotal::inc));
+            } else {
+                metric = stats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkOutboundUsageMetric(strategy, networkOutboundUsageTotal::inc));
+            }
+        }
+        return metric;
     }
 
     public Optional<Counter> fastReadBytesStats(long streamId) {

s3stream/src/main/java/com/automq/stream/s3/network/AsyncNetworkBandwidthLimiter.java

@@ -34,6 +34,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
@@ -51,7 +52,7 @@ public class AsyncNetworkBandwidthLimiter implements NetworkBandwidthLimiter {
     private final Queue<BucketItem> queuedCallbacks;
     private final Type type;
     private final long tokenSize;
-    private long availableTokens;
+    private final AtomicLong availableTokens;
 
     public AsyncNetworkBandwidthLimiter(Type type, long tokenSize, int refillIntervalMs) {
         this(type, tokenSize, refillIntervalMs, tokenSize);
@@ -61,12 +62,13 @@ public AsyncNetworkBandwidthLimiter(Type type, long tokenSize, int refillInterva
     public AsyncNetworkBandwidthLimiter(Type type, long tokenSize, int refillIntervalMs, long maxTokens) {
         this.type = type;
         this.tokenSize = tokenSize;
-        this.availableTokens = this.tokenSize;
+        this.availableTokens = new AtomicLong(this.tokenSize);
         this.maxTokens = maxTokens;
         this.queuedCallbacks = new PriorityQueue<>();
         this.refillThreadPool =
             Threads.newSingleThreadScheduledExecutor(new DefaultThreadFactory("refill-bucket-thread"), LOGGER);
-        this.callbackThreadPool = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(1, "callback-thread", true, LOGGER);
+        // The threads number must be larger than 1 because the #run will occupy one thread.
+        this.callbackThreadPool = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(2, "callback-thread", true, LOGGER);
         this.callbackThreadPool.execute(this::run);
         this.refillThreadPool.scheduleAtFixedRate(this::refillToken, refillIntervalMs, refillIntervalMs, TimeUnit.MILLISECONDS);
         S3StreamMetricsManager.registerNetworkLimiterQueueSizeSupplier(type, this::getQueueSize);
@@ -88,7 +90,7 @@ private void run() {
                     }
                     long size = Math.min(head.size, MAX_TOKEN_PART_SIZE);
                     reduceToken(size);
-                    if (head.complete(size)) {
+                    if (head.complete(size, callbackThreadPool)) {
                         queuedCallbacks.poll();
                     }
                 }
@@ -103,7 +105,7 @@ private void run() {
     private void refillToken() {
         lock.lock();
         try {
-            availableTokens = Math.min(availableTokens + this.tokenSize, this.maxTokens);
+            this.availableTokens.getAndUpdate(old -> Math.min(old + this.tokenSize, this.maxTokens));
             condition.signalAll();
         } finally {
             lock.unlock();
@@ -114,7 +116,7 @@ private boolean ableToConsume() {
         if (queuedCallbacks.isEmpty()) {
             return false;
         }
-        return availableTokens > 0;
+        return availableTokens.get() > 0;
     }
 
     public void shutdown() {
@@ -127,12 +129,7 @@ public long getMaxTokens() {
     }
 
     public long getAvailableTokens() {
-        lock.lock();
-        try {
-            return availableTokens;
-        } finally {
-            lock.unlock();
-        }
+        return availableTokens.get();
     }
 
     public int getQueueSize() {
@@ -145,12 +142,7 @@ public int getQueueSize() {
     }
 
     private void forceConsume(long size) {
-        lock.lock();
-        try {
-            reduceToken(size);
-        } finally {
-            lock.unlock();
-        }
+        reduceToken(size);
     }
 
     public CompletableFuture<Void> consume(ThrottleStrategy throttleStrategy, long size) {
@@ -162,7 +154,7 @@ public CompletableFuture<Void> consume(ThrottleStrategy throttleStrategy, long s
         } else {
             lock.lock();
             try {
-                if (availableTokens <= 0 || !queuedCallbacks.isEmpty()) {
+                if (availableTokens.get() <= 0 || !queuedCallbacks.isEmpty()) {
                     queuedCallbacks.offer(new BucketItem(throttleStrategy, size, cf));
                     condition.signalAll();
                 } else {
@@ -177,7 +169,7 @@ public CompletableFuture<Void> consume(ThrottleStrategy throttleStrategy, long s
     }
 
     private void reduceToken(long size) {
-        this.availableTokens = Math.max(-maxTokens, availableTokens - size);
+        this.availableTokens.getAndUpdate(old -> Math.max(-maxTokens, old - size));
     }
 
     public enum Type {
@@ -216,10 +208,10 @@ public int compareTo(BucketItem o) {
             return Long.compare(strategy.priority(), o.strategy.priority());
         }
 
-        public boolean complete(long completeSize) {
+        public boolean complete(long completeSize, ExecutorService executor) {
             size -= completeSize;
             if (size <= 0) {
-                cf.complete(null);
+                executor.submit(() -> cf.complete(null));
                 return true;
             }
             return false;

s3stream/src/test/java/com/automq/stream/s3/AsyncNetworkBandwidthLimiterTest.java

@@ -118,13 +118,10 @@ public void testThrottleConsume5() throws InterruptedException {
         Assertions.assertEquals(-200, bucket.getAvailableTokens());
         Thread.sleep(500);
         bucket.consume(ThrottleStrategy.BYPASS, 500);
-        CompletableFuture<Void> cf = bucket.consume(ThrottleStrategy.CATCH_UP, 5);
-        bucket.consume(ThrottleStrategy.CATCH_UP, 10);
-        CompletableFuture<Void> result = cf.whenComplete((v, e) -> {
-            Assertions.assertNull(e);
-            Assertions.assertTrue(bucket.getAvailableTokens() >= 0);
-        });
-        result.join();
+        bucket.consume(ThrottleStrategy.CATCH_UP, 5);
+        CompletableFuture<Void> cf = bucket.consume(ThrottleStrategy.CATCH_UP, 10);
+        cf.join();
+        Assertions.assertEquals(-5, bucket.getAvailableTokens());
     }
 
     @Test