Repository: skiptests/astraea
Pull request #489 (2022-07-18 13:06:31)

QuotaCreator should use DataRate rather than long
感謝 @garyparrot  在 #481 中重構了 `DataRate`，因此可以讓其作為public APIs的一部分了


Patch:
diff --git a/app/src/main/java/org/astraea/app/admin/Builder.java b/app/src/main/java/org/astraea/app/admin/Builder.java
index 1fa30c99..e588d7ca 100644
--- a/app/src/main/java/org/astraea/app/admin/Builder.java
+++ b/app/src/main/java/org/astraea/app/admin/Builder.java
@@ -52,6 +52,7 @@ import org.apache.kafka.common.quota.ClientQuotaAlteration;
 import org.apache.kafka.common.quota.ClientQuotaEntity;
 import org.apache.kafka.common.quota.ClientQuotaFilter;
 import org.apache.kafka.common.quota.ClientQuotaFilterComponent;
+import org.astraea.app.common.DataRate;
 import org.astraea.app.common.Utils;
 
 public class Builder {
@@ -936,17 +937,17 @@ public class Builder {
     @Override
     public Client clientId(String id) {
       return new Client() {
-        private int produceRate = Integer.MAX_VALUE;
-        private int consumeRate = Integer.MAX_VALUE;
+        private DataRate produceRate = null;
+        private DataRate consumeRate = null;
 
         @Override
-        public Client produceRate(int value) {
+        public Client produceRate(DataRate value) {
           this.produceRate = value;
           return this;
         }
 
         @Override
-        public Client consumeRate(int value) {
+        public Client consumeRate(DataRate value) {
           this.consumeRate = value;
           return this;
         }
@@ -954,14 +955,14 @@ public class Builder {
         @Override
         public void create() {
           var q = new ArrayList<ClientQuotaAlteration.Op>();
-          if (produceRate != Integer.MAX_VALUE)
+          if (produceRate != null)
             q.add(
                 new ClientQuotaAlteration.Op(
-                    Quota.Limit.PRODUCER_BYTE_RATE.nameOfKafka(), (double) produceRate));
-          if (consumeRate != Integer.MAX_VALUE)
+                    Quota.Limit.PRODUCER_BYTE_RATE.nameOfKafka(), produceRate.byteRate()));
+          if (consumeRate != null)
             q.add(
                 new ClientQuotaAlteration.Op(
-                    Quota.Limit.CONSUMER_BYTE_RATE.nameOfKafka(), (double) consumeRate));
+                    Quota.Limit.CONSUMER_BYTE_RATE.nameOfKafka(), consumeRate.byteRate()));
           if (!q.isEmpty())
             Utils.packException(
                 () ->
diff --git a/app/src/main/java/org/astraea/app/admin/QuotaCreator.java b/app/src/main/java/org/astraea/app/admin/QuotaCreator.java
index ccac40e0..f9de4e7c 100644
--- a/app/src/main/java/org/astraea/app/admin/QuotaCreator.java
+++ b/app/src/main/java/org/astraea/app/admin/QuotaCreator.java
@@ -16,6 +16,8 @@
  */
 package org.astraea.app.admin;
 
+import org.astraea.app.common.DataRate;
+
 /**
  * Kafka quota APIs are too incomprehensible to use, so we re-design quota APIs to builder pattern.
  */
@@ -38,16 +40,16 @@ public interface QuotaCreator {
    */
   interface Client {
     /**
-     * @param value A rate representing the upper bound (bytes/sec) for producer traffic
+     * @param value A rate representing the upper bound for producer traffic
      * @return this object
      */
-    Client produceRate(int value);
+    Client produceRate(DataRate value);
 
     /**
-     * @param value A rate representing the upper bound (bytes/sec) for consumer traffic
+     * @param value A rate representing the upper bound for consumer traffic
      * @return this object
      */
-    Client consumeRate(int value);
+    Client consumeRate(DataRate value);
 
     void create();
   }
diff --git a/app/src/main/java/org/astraea/app/performance/DataSupplier.java b/app/src/main/java/org/astraea/app/performance/DataSupplier.java
index b20caf48..b3e62ab1 100644
--- a/app/src/main/java/org/astraea/app/performance/DataSupplier.java
+++ b/app/src/main/java/org/astraea/app/performance/DataSupplier.java
@@ -152,27 +152,13 @@ interface DataSupplier extends Supplier<DataSupplier.Data> {
       Supplier<Long> valueSizeDistribution,
       DataSize throughput) {
     return new DataSupplier() {
+      private final Throttler throttler = new Throttler(throughput);
       private final long start = System.currentTimeMillis();
       private final Random rand = new Random();
       private final AtomicLong dataCount = new AtomicLong(0);
-      private long intervalStart = 0;
-      private long payloadBytes;
       private final Map<Long, byte[]> recordKeyTable = new ConcurrentHashMap<>();
       private final Map<Long, byte[]> recordValueTable = new ConcurrentHashMap<>();
 
-      synchronized boolean checkAndAdd(int payloadLength) {
-        if (System.currentTimeMillis() - intervalStart > 1000) {
-          intervalStart = System.currentTimeMillis();
-          payloadBytes = payloadLength;
-          return true;
-        } else if (payloadBytes < throughput.measurement(DataUnit.Byte).longValue()) {
-          payloadBytes += payloadLength;
-          return true;
-        } else {
-          return false;
-        }
-      }
-
       byte[] value() {
         return getOrNew(
             recordValueTable, valueDistribution, valueSizeDistribution.get().intValue(), rand);
@@ -189,9 +175,10 @@ interface DataSupplier extends Supplier<DataSupplier.Data> {
             >= 100D) return NO_MORE_DATA;
         var key = key();
         var value = value();
-        if (checkAndAdd((value != null ? value.length : 0) + (key != null ? key.length : 0)))
-          return data(key, value);
-        return THROTTLED_DATA;
+        if (throttler.throttled(
+            (value != null ? value.length : 0) + (key != null ? key.length : 0)))
+          return THROTTLED_DATA;
+        return data(key, value);
       }
     };
   }
@@ -210,4 +197,35 @@ interface DataSupplier extends Supplier<DataSupplier.Data> {
           return value;
         });
   }
+
+  class Throttler {
+    private final long start = System.currentTimeMillis();
+    private final long throughput;
+    private final AtomicLong totalBytes = new AtomicLong();
+
+    Throttler(DataSize max) {
+      throughput = max.measurement(DataUnit.Byte).longValue();
+    }
+
+    /**
+     * @param payloadLength of new data
+     * @return true if the data need to be throttled. Otherwise, false
+     */
+    boolean throttled(long payloadLength) {
+      var duration = durationInSeconds();
+      if (duration <= 0) return false;
+      var current = totalBytes.addAndGet(payloadLength);
+      // too much -> slow down
+      if ((current / duration) > throughput) {
+        totalBytes.addAndGet(-payloadLength);
+        return true;
+      }
+      return false;
+    }
+
+    // visible for testing
+    long durationInSeconds() {
+      return (System.currentTimeMillis() - start) / 1000;
+    }
+  }
 }
diff --git a/app/src/main/java/org/astraea/app/performance/Performance.java b/app/src/main/java/org/astraea/app/performance/Performance.java
index 742cd681..be6933c9 100644
--- a/app/src/main/java/org/astraea/app/performance/Performance.java
+++ b/app/src/main/java/org/astraea/app/performance/Performance.java
@@ -365,6 +365,7 @@ public class Performance {
         validateWith = NonEmptyStringField.class)
     List<Integer> specifyBroker = List.of(-1);
 
+    // replace DataSize by DataRate (see https://github.com/skiptests/astraea/issues/488)
     @Parameter(
         names = {"--throughput"},
         description = "dataSize: size output per second. e.g. \"500KiB\"",
diff --git a/app/src/main/java/org/astraea/app/web/QuotaHandler.java b/app/src/main/java/org/astraea/app/web/QuotaHandler.java
index 0413e920..d212e56e 100644
--- a/app/src/main/java/org/astraea/app/web/QuotaHandler.java
+++ b/app/src/main/java/org/astraea/app/web/QuotaHandler.java
@@ -22,6 +22,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import org.astraea.app.admin.Admin;
+import org.astraea.app.common.DataRate;
 
 public class QuotaHandler implements Handler {
 
@@ -64,8 +65,20 @@ public class QuotaHandler implements Handler {
       admin
           .quotaCreator()
           .clientId(request.value(CLIENT_ID_KEY))
-          .produceRate(request.intValue(PRODUCE_RATE_KEY, Integer.MAX_VALUE))
-          .consumeRate(request.intValue(CONSUME_RATE_KEY, Integer.MAX_VALUE))
+          // TODO: use DataRate#Field (traced https://github.com/skiptests/astraea/issues/488)
+          // see https://github.com/skiptests/astraea/issues/490
+          .produceRate(
+              request
+                  .get(PRODUCE_RATE_KEY)
+                  .map(Long::parseLong)
+                  .map(v -> DataRate.Byte.of(v).perSecond())
+                  .orElse(null))
+          .consumeRate(
+              request
+                  .get(CONSUME_RATE_KEY)
+                  .map(Long::parseLong)
+                  .map(v -> DataRate.Byte.of(v).perSecond())
+                  .orElse(null))
           .create();
       return new Quotas(
           admin.quotas(org.astraea.app.admin.Quota.Target.CLIENT_ID, request.value(CLIENT_ID_KEY)));


Test patch:
diff --git a/app/src/test/java/org/astraea/app/admin/AdminTest.java b/app/src/test/java/org/astraea/app/admin/AdminTest.java
index d630c601..f40a4efd 100644
--- a/app/src/test/java/org/astraea/app/admin/AdminTest.java
+++ b/app/src/test/java/org/astraea/app/admin/AdminTest.java
@@ -35,6 +35,7 @@ import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 import org.apache.kafka.common.config.TopicConfig;
 import org.apache.kafka.common.errors.GroupNotEmptyException;
+import org.astraea.app.common.DataRate;
 import org.astraea.app.common.Utils;
 import org.astraea.app.concurrent.State;
 import org.astraea.app.concurrent.ThreadPool;
@@ -585,7 +586,12 @@ public class AdminTest extends RequireBrokerCluster {
   @Test
   void testClientQuota() {
     try (var admin = Admin.of(bootstrapServers())) {
-      admin.quotaCreator().clientId("my-id").produceRate(10).consumeRate(100).create();
+      admin
+          .quotaCreator()
+          .clientId("my-id")
+          .produceRate(DataRate.Byte.of(10L).perSecond())
+          .consumeRate(DataRate.Byte.of(100L).perSecond())
+          .create();
       Utils.sleep(Duration.ofSeconds(2));
 
       java.util.function.Consumer<List<Quota>> checker =
@@ -641,8 +647,16 @@ public class AdminTest extends RequireBrokerCluster {
   @Test
   void testMultipleClientQuota() {
     try (var admin = Admin.of(bootstrapServers())) {
-      admin.quotaCreator().clientId("my-id").consumeRate(100).create();
-      admin.quotaCreator().clientId("my-id").produceRate(999).create();
+      admin
+          .quotaCreator()
+          .clientId("my-id")
+          .consumeRate(DataRate.Byte.of(100L).perSecond())
+          .create();
+      admin
+          .quotaCreator()
+          .clientId("my-id")
+          .produceRate(DataRate.Byte.of(1000L).perSecond())
+          .create();
       Utils.sleep(Duration.ofSeconds(2));
       Assertions.assertEquals(2, admin.quotas(Quota.Target.CLIENT_ID, "my-id").size());
     }
diff --git a/app/src/test/java/org/astraea/app/performance/DataSupplierTest.java b/app/src/test/java/org/astraea/app/performance/DataSupplierTest.java
index 7e85482e..18774a3b 100644
--- a/app/src/test/java/org/astraea/app/performance/DataSupplierTest.java
+++ b/app/src/test/java/org/astraea/app/performance/DataSupplierTest.java
@@ -180,20 +180,20 @@ public class DataSupplierTest {
 
   @Test
   void testThrottle() {
-    var dataSupplier =
-        DataSupplier.of(
-            ExeTime.of("10s"),
-            () -> 10L,
-            () -> DataSize.KiB.of(50).measurement(DataUnit.Byte).longValue(),
-            () -> 10L,
-            () -> DataSize.KiB.of(50).measurement(DataUnit.Byte).longValue(),
-            DataSize.KiB.of(150));
+    var durationInSeconds = new AtomicLong(1);
+    var throttler =
+        new DataSupplier.Throttler(DataSize.KiB.of(150)) {
+          @Override
+          long durationInSeconds() {
+            return durationInSeconds.get();
+          }
+        };
     // total: 100KB, limit: 150KB -> no throttle
-    Assertions.assertTrue(dataSupplier.get().hasData());
-    // total: 200KB, limit: 150KB -> will throttle next data
-    Assertions.assertTrue(dataSupplier.get().hasData());
-    // throttled
-    Assertions.assertFalse(dataSupplier.get().hasData());
+    Assertions.assertFalse(
+        throttler.throttled(DataSize.KiB.of(100).measurement(DataUnit.Byte).longValue()));
+    // total: 200KB, limit: 150KB -> throttled
+    Assertions.assertTrue(
+        throttler.throttled(DataSize.KiB.of(100).measurement(DataUnit.Byte).longValue()));
   }
 
   @Test