Repository: skiptests/astraea
Pull request #345 (2022-05-15 18:39:35)

Enable to write/read quota
這讓我們的admin更加豐富，並且在下一個階段可以加上restful apis


Patch:
diff --git a/app/src/main/java/org/astraea/admin/Admin.java b/app/src/main/java/org/astraea/admin/Admin.java
index af51509c..bfb03694 100644
--- a/app/src/main/java/org/astraea/admin/Admin.java
+++ b/app/src/main/java/org/astraea/admin/Admin.java
@@ -43,7 +43,7 @@ public interface Admin extends Closeable {
   Set<TopicPartition> partitions(Set<String> topics);
 
   /** @return a topic creator to set all topic configs and then run the procedure. */
-  Creator creator();
+  TopicCreator creator();
 
   /** @return offsets of all partitions */
   default Map<TopicPartition, Offset> offsets() {
@@ -114,7 +114,7 @@ public interface Admin extends Closeable {
   Map<Integer, Set<String>> brokerFolders(Set<Integer> brokers);
 
   /** @return a partition migrator used to move partitions to another broker or folder. */
-  Migrator migrator();
+  ReplicaMigrator migrator();
 
   /** @return producer states of all topic partitions */
   default Map<TopicPartition, Collection<ProducerState>> producerStates() {
@@ -127,6 +127,25 @@ public interface Admin extends Closeable {
    */
   Map<TopicPartition, Collection<ProducerState>> producerStates(Set<TopicPartition> partitions);
 
+  /** @return a progress to set quota */
+  QuotaCreator quotaCreator();
+
+  /**
+   * @param target to search
+   * @return quotas
+   */
+  Collection<Quota> quotas(Quota.Target target);
+
+  /**
+   * @param target to search
+   * @param value assoicated to target
+   * @return quotas
+   */
+  Collection<Quota> quotas(Quota.Target target, String value);
+
+  /** @return all quotas */
+  Collection<Quota> quotas();
+
   @Override
   void close();
 }
diff --git a/app/src/main/java/org/astraea/admin/Builder.java b/app/src/main/java/org/astraea/admin/Builder.java
index 05700d54..2471fc66 100644
--- a/app/src/main/java/org/astraea/admin/Builder.java
+++ b/app/src/main/java/org/astraea/admin/Builder.java
@@ -28,6 +28,10 @@ import org.apache.kafka.common.ElectionType;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartitionReplica;
 import org.apache.kafka.common.config.ConfigResource;
+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.Utils;
 
 public class Builder {
@@ -81,7 +85,7 @@ public class Builder {
     }
 
     @Override
-    public Migrator migrator() {
+    public ReplicaMigrator migrator() {
       return new MigratorImpl(admin, this::partitions);
     }
 
@@ -338,10 +342,39 @@ public class Builder {
     }
 
     @Override
-    public Creator creator() {
+    public TopicCreator creator() {
       return new CreatorImpl(
           admin, topic -> this.replicas(Set.of(topic)), topic -> topics().get(topic));
     }
+
+    @Override
+    public QuotaCreator quotaCreator() {
+      return new QuotaImpl(admin);
+    }
+
+    @Override
+    public Collection<Quota> quotas(Quota.Target target) {
+      return quotas(
+          ClientQuotaFilter.contains(
+              List.of(ClientQuotaFilterComponent.ofEntityType(target.nameOfKafka()))));
+    }
+
+    @Override
+    public Collection<Quota> quotas(Quota.Target target, String value) {
+      return quotas(
+          ClientQuotaFilter.contains(
+              List.of(ClientQuotaFilterComponent.ofEntity(target.nameOfKafka(), value))));
+    }
+
+    @Override
+    public Collection<Quota> quotas() {
+      return quotas(ClientQuotaFilter.all());
+    }
+
+    private Collection<Quota> quotas(ClientQuotaFilter filter) {
+      return Quota.of(
+          Utils.handleException(() -> admin.describeClientQuotas(filter).entities().get()));
+    }
   }
 
   private static class ConfigImpl implements Config {
@@ -379,7 +412,7 @@ public class Builder {
     }
   }
 
-  private static class CreatorImpl implements Creator {
+  private static class CreatorImpl implements TopicCreator {
     private final org.apache.kafka.clients.admin.Admin admin;
     private final Function<String, Map<TopicPartition, List<Replica>>> replicasGetter;
     private final Function<String, Config> configsGetter;
@@ -398,31 +431,31 @@ public class Builder {
     }
 
     @Override
-    public Creator topic(String topic) {
+    public TopicCreator topic(String topic) {
       this.topic = Objects.requireNonNull(topic);
       return this;
     }
 
     @Override
-    public Creator numberOfPartitions(int numberOfPartitions) {
+    public TopicCreator numberOfPartitions(int numberOfPartitions) {
       this.numberOfPartitions = numberOfPartitions;
       return this;
     }
 
     @Override
-    public Creator numberOfReplicas(short numberOfReplicas) {
+    public TopicCreator numberOfReplicas(short numberOfReplicas) {
       this.numberOfReplicas = numberOfReplicas;
       return this;
     }
 
     @Override
-    public Creator config(String key, String value) {
+    public TopicCreator config(String key, String value) {
       this.configs.put(key, value);
       return this;
     }
 
     @Override
-    public Creator configs(Map<String, String> configs) {
+    public TopicCreator configs(Map<String, String> configs) {
       this.configs.putAll(configs);
       return this;
     }
@@ -483,7 +516,7 @@ public class Builder {
     }
   }
 
-  private static class MigratorImpl implements Migrator {
+  private static class MigratorImpl implements ReplicaMigrator {
     private final org.apache.kafka.clients.admin.Admin admin;
     private final Function<Set<String>, Set<TopicPartition>> partitionGetter;
     private final Set<TopicPartition> partitions = new HashSet<>();
@@ -497,13 +530,13 @@ public class Builder {
     }
 
     @Override
-    public Migrator topic(String topic) {
+    public ReplicaMigrator topic(String topic) {
       partitions.addAll(partitionGetter.apply(Set.of(topic)));
       return this;
     }
 
     @Override
-    public Migrator partition(String topic, int partition) {
+    public ReplicaMigrator partition(String topic, int partition) {
       partitions.add(new TopicPartition(topic, partition));
       return this;
     }
@@ -556,4 +589,79 @@ public class Builder {
                     .get());
     }
   }
+
+  private static class QuotaImpl implements QuotaCreator {
+    private final org.apache.kafka.clients.admin.Admin admin;
+
+    QuotaImpl(org.apache.kafka.clients.admin.Admin admin) {
+      this.admin = admin;
+    }
+
+    @Override
+    public Ip ip(String ip) {
+      return new Ip() {
+        private int connectionRate = Integer.MAX_VALUE;
+
+        @Override
+        public Ip connectionRate(int value) {
+          this.connectionRate = value;
+          return this;
+        }
+
+        @Override
+        public void create() {
+          Utils.handleException(
+              () ->
+                  admin
+                      .alterClientQuotas(
+                          List.of(
+                              new ClientQuotaAlteration(
+                                  new ClientQuotaEntity(Map.of(ClientQuotaEntity.IP, ip)),
+                                  List.of(
+                                      new ClientQuotaAlteration.Op(
+                                          "connection_creation_rate", (double) connectionRate)))))
+                      .all()
+                      .get());
+        }
+      };
+    }
+
+    @Override
+    public Client clientId(String id) {
+      return new Client() {
+        private int produceRate = Integer.MAX_VALUE;
+        private int consumeRate = Integer.MAX_VALUE;
+
+        @Override
+        public Client produceRate(int value) {
+          this.produceRate = value;
+          return this;
+        }
+
+        @Override
+        public Client consumeRate(int value) {
+          this.consumeRate = value;
+          return this;
+        }
+
+        @Override
+        public void create() {
+          Utils.handleException(
+              () ->
+                  admin
+                      .alterClientQuotas(
+                          List.of(
+                              new ClientQuotaAlteration(
+                                  new ClientQuotaEntity(Map.of(ClientQuotaEntity.CLIENT_ID, id)),
+                                  List.of(
+                                      new ClientQuotaAlteration.Op(
+                                          "producer_byte_rate", (double) produceRate),
+                                      new ClientQuotaAlteration.Op(
+                                          "consumer_byte_rate", (double) consumeRate)))))
+                      .all()
+                      .get());
+        }
+      };
+    }
+  }
 }
diff --git a/app/src/main/java/org/astraea/admin/Quota.java b/app/src/main/java/org/astraea/admin/Quota.java
new file mode 100644
index 00000000..e955ab02
--- /dev/null
+++ b/app/src/main/java/org/astraea/admin/Quota.java
@@ -0,0 +1,135 @@
+package org.astraea.admin;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.apache.kafka.common.quota.ClientQuotaEntity;
+
+public class Quota {
+
+  static Collection<Quota> of(Map<ClientQuotaEntity, Map<String, Double>> data) {
+    return data.entrySet().stream()
+        .flatMap(
+            clientQuotaEntityMapEntry ->
+                clientQuotaEntityMapEntry.getKey().entries().entrySet().stream()
+                    .flatMap(
+                        stringStringEntry ->
+                            clientQuotaEntityMapEntry.getValue().entrySet().stream()
+                                .map(
+                                    v ->
+                                        new Quota(
+                                            target(stringStringEntry.getKey()),
+                                            stringStringEntry.getValue(),
+                                            action(v.getKey()),
+                                            v.getValue()))))
+        .collect(Collectors.toUnmodifiableList());
+  }
+
+  static Target target(String value) {
+    return Arrays.stream(Target.values())
+        .filter(p -> p.nameOfKafka.equalsIgnoreCase(value))
+        .findFirst()
+        .orElseThrow(() -> new RuntimeException("unknown: " + value));
+  }
+
+  static Action action(String value) {
+    return Arrays.stream(Action.values())
+        .filter(p -> p.nameOfKafka.equalsIgnoreCase(value))
+        .findFirst()
+        .orElseThrow(() -> new RuntimeException("unknown: " + value));
+  }
+
+  enum Target {
+    USER(ClientQuotaEntity.USER),
+    CLIENT_ID(ClientQuotaEntity.CLIENT_ID),
+    IP(ClientQuotaEntity.IP);
+    private final String nameOfKafka;
+
+    Target(String nameOfKafka) {
+      this.nameOfKafka = nameOfKafka;
+    }
+
+    String nameOfKafka() {
+      return nameOfKafka;
+    }
+  }
+
+  enum Action {
+    PRODUCER_BYTE_RATE("producer_byte_rate"),
+    CONSUMER_BYTE_RATE("consumer_byte_rate"),
+    REQUEST_PERCENTAGE("request_percentage"),
+    CONTROLLER_MUTATION_RATE("controller_mutation_rate"),
+    IP_CONNECTION_RATE("connection_creation_rate");
+    private final String nameOfKafka;
+
+    Action(String nameOfKafka) {
+      this.nameOfKafka = nameOfKafka;
+    }
+
+    String nameOfKafka() {
+      return nameOfKafka;
+    }
+  }
+
+  private final Target target;
+
+  private final String targetValue;
+  private final Action action;
+  private final double actionValue;
+
+  public Quota(Target target, String targetValue, Action action, double actionValue) {
+    this.target = target;
+    this.targetValue = targetValue;
+    this.action = action;
+    this.actionValue = actionValue;
+  }
+
+  public Target target() {
+    return target;
+  }
+
+  public String targetValue() {
+    return targetValue;
+  }
+
+  public Action action() {
+    return action;
+  }
+
+  public double actionValue() {
+    return actionValue;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    Quota quota = (Quota) o;
+    return Double.compare(quota.actionValue, actionValue) == 0
+        && target == quota.target
+        && Objects.equals(targetValue, quota.targetValue)
+        && action == quota.action;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(target, targetValue, action, actionValue);
+  }
+
+  @Override
+  public String toString() {
+    return "Quota{"
+        + "target="
+        + target
+        + ", targetValue='"
+        + targetValue
+        + '\''
+        + ", action="
+        + action
+        + ", actionValue="
+        + actionValue
+        + '}';
+  }
+}
diff --git a/app/src/main/java/org/astraea/admin/QuotaCreator.java b/app/src/main/java/org/astraea/admin/QuotaCreator.java
new file mode 100644
index 00000000..f2a82968
--- /dev/null
+++ b/app/src/main/java/org/astraea/admin/QuotaCreator.java
@@ -0,0 +1,54 @@
+package org.astraea.admin;
+
+/**
+ * Kafka quota APIs are too incomprehensible to use, so we re-design quota APIs to builder pattern.
+ */
+public interface QuotaCreator {
+
+  /** the valid quota for IP is connection rate only. */
+  interface Ip {
+    /**
+     * @param value An int representing the upper bound of connections
+     * @return this object
+     */
+    Ip connectionRate(int value);
+
+    void create();
+  }
+
+  /**
+   * the valid quota for IP are produce, consume and controller. However, we show only produce and
+   * consume here
+   */
+  interface Client {
+    /**
+     * @param value A rate representing the upper bound (bytes/sec) for producer traffic
+     * @return this object
+     */
+    Client produceRate(int value);
+
+    /**
+     * @param value A rate representing the upper bound (bytes/sec) for consumer traffic
+     * @return this object
+     */
+    Client consumeRate(int value);
+
+    void create();
+  }
+
+  /**
+   * start to set quota for specify ip address
+   *
+   * @param ip to add quota
+   * @return ip quota creator
+   */
+  Ip ip(String ip);
+
+  /**
+   * start to set quota for specify client id
+   *
+   * @param id to add quota
+   * @return client id quota creator
+   */
+  Client clientId(String id);
+}
diff --git a/app/src/main/java/org/astraea/admin/Migrator.java b/app/src/main/java/org/astraea/admin/ReplicaMigrator.java
similarity index 89%
rename from app/src/main/java/org/astraea/admin/Migrator.java
rename to app/src/main/java/org/astraea/admin/ReplicaMigrator.java
index a18f73a6..31e7fc5d 100644
--- a/app/src/main/java/org/astraea/admin/Migrator.java
+++ b/app/src/main/java/org/astraea/admin/ReplicaMigrator.java
@@ -5,14 +5,14 @@ import java.util.Map;
 import java.util.Set;
 
 /** used to migrate partitions to another broker or broker folder. */
-public interface Migrator {
+public interface ReplicaMigrator {
   /**
    * move all partitions (leader replica and follower replicas) of topic
    *
    * @param topic topic name
    * @return this migrator
    */
-  Migrator topic(String topic);
+  ReplicaMigrator topic(String topic);
 
   /**
    * move one partition (leader replica and follower replicas) of topic
@@ -21,7 +21,7 @@ public interface Migrator {
    * @param partition partition id
    * @return this migrator
    */
-  Migrator partition(String topic, int partition);
+  ReplicaMigrator partition(String topic, int partition);
 
   /**
    * move partitions to specify brokers. Noted that this method won't invoke leader election
diff --git a/app/src/main/java/org/astraea/admin/Creator.java b/app/src/main/java/org/astraea/admin/TopicCreator.java
similarity index 65%
rename from app/src/main/java/org/astraea/admin/Creator.java
rename to app/src/main/java/org/astraea/admin/TopicCreator.java
index 0d2aad97..46499b58 100644
--- a/app/src/main/java/org/astraea/admin/Creator.java
+++ b/app/src/main/java/org/astraea/admin/TopicCreator.java
@@ -4,19 +4,19 @@ import java.time.Duration;
 import java.util.Map;
 import org.apache.kafka.common.config.TopicConfig;
 
-public interface Creator {
-  Creator topic(String topic);
+public interface TopicCreator {
+  TopicCreator topic(String topic);
 
-  Creator numberOfPartitions(int numberOfPartitions);
+  TopicCreator numberOfPartitions(int numberOfPartitions);
 
-  Creator numberOfReplicas(short numberOfReplicas);
+  TopicCreator numberOfReplicas(short numberOfReplicas);
 
   /**
    * set the cleanup policy to compact. By default, the policy is `delete`/
    *
    * @return this creator
    */
-  default Creator compacted() {
+  default TopicCreator compacted() {
     return config(TopicConfig.CLEANUP_POLICY_CONFIG, TopicConfig.CLEANUP_POLICY_COMPACT);
   }
 
@@ -26,22 +26,22 @@ public interface Creator {
    * @param value the max lag
    * @return this creator
    */
-  default Creator compactionMaxLag(Duration value) {
+  default TopicCreator compactionMaxLag(Duration value) {
     return compacted()
         .config(TopicConfig.MAX_COMPACTION_LAG_MS_CONFIG, String.valueOf(value.toMillis()));
   }
 
-  default Creator compression(Compression compression) {
+  default TopicCreator compression(Compression compression) {
     return config(TopicConfig.COMPRESSION_TYPE_CONFIG, compression.nameOfKafka());
   }
 
-  Creator config(String key, String value);
+  TopicCreator config(String key, String value);
 
   /**
    * @param configs used to create new topic
    * @return this creator
    */
-  Creator configs(Map<String, String> configs);
+  TopicCreator configs(Map<String, String> configs);
 
   /** start to create topic. */
   void create();


Test patch:
diff --git a/app/src/test/java/org/astraea/admin/AdminTest.java b/app/src/test/java/org/astraea/admin/AdminTest.java
index 2f5b57d7..572ba7cd 100644
--- a/app/src/test/java/org/astraea/admin/AdminTest.java
+++ b/app/src/test/java/org/astraea/admin/AdminTest.java
@@ -414,4 +414,90 @@ public class AdminTest extends RequireBrokerCluster {
       Assertions.assertEquals(1, producerState.size());
     }
   }
+
+  @Test
+  void testIpQuota() throws InterruptedException {
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.quotaCreator().ip("192.168.11.11").connectionRate(10).create();
+      TimeUnit.SECONDS.sleep(2);
+
+      java.util.function.Consumer<List<Quota>> checker =
+          (quotas) -> {
+            Assertions.assertEquals(1, quotas.size());
+            Assertions.assertEquals(
+                Set.of(Quota.Target.IP),
+                quotas.stream().map(Quota::target).collect(Collectors.toSet()));
+            Assertions.assertEquals(
+                Set.of(Quota.Action.IP_CONNECTION_RATE),
+                quotas.stream().map(Quota::action).collect(Collectors.toSet()));
+            Assertions.assertEquals("192.168.11.11", quotas.iterator().next().targetValue());
+            Assertions.assertEquals(10, quotas.iterator().next().actionValue());
+          };
+
+      // only target
+      checker.accept(
+          admin.quotas(Quota.Target.IP).stream().collect(Collectors.toUnmodifiableList()));
+
+      // only target and name
+      checker.accept(
+          admin.quotas(Quota.Target.IP, "192.168.11.11").stream()
+              .collect(Collectors.toUnmodifiableList()));
+    }
+  }
+
+  @Test
+  void testClientQuota() throws InterruptedException {
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.quotaCreator().clientId("my-id").produceRate(10).consumeRate(100).create();
+      TimeUnit.SECONDS.sleep(2);
+
+      java.util.function.Consumer<List<Quota>> checker =
+          (quotas) -> {
+            Assertions.assertEquals(2, quotas.size());
+            Assertions.assertEquals(
+                Set.of(Quota.Target.CLIENT_ID),
+                quotas.stream().map(Quota::target).collect(Collectors.toSet()));
+            Assertions.assertEquals(
+                Set.of(Quota.Action.PRODUCER_BYTE_RATE, Quota.Action.CONSUMER_BYTE_RATE),
+                quotas.stream().map(Quota::action).collect(Collectors.toSet()));
+            Assertions.assertEquals(
+                10,
+                quotas.stream()
+                    .filter(q -> q.action() == Quota.Action.PRODUCER_BYTE_RATE)
+                    .findFirst()
+                    .get()
+                    .actionValue());
+            Assertions.assertEquals(
+                "my-id",
+                quotas.stream()
+                    .filter(q -> q.action() == Quota.Action.PRODUCER_BYTE_RATE)
+                    .findFirst()
+                    .get()
+                    .targetValue());
+            Assertions.assertEquals(
+                100,
+                quotas.stream()
+                    .filter(q -> q.action() == Quota.Action.CONSUMER_BYTE_RATE)
+                    .findFirst()
+                    .get()
+                    .actionValue());
+            Assertions.assertEquals(
+                "my-id",
+                quotas.stream()
+                    .filter(q -> q.action() == Quota.Action.CONSUMER_BYTE_RATE)
+                    .findFirst()
+                    .get()
+                    .targetValue());
+          };
+
+      // only target
+      checker.accept(
+          admin.quotas(Quota.Target.CLIENT_ID).stream().collect(Collectors.toUnmodifiableList()));
+
+      // only target and name
+      checker.accept(
+          admin.quotas(Quota.Target.CLIENT_ID, "my-id").stream()
+              .collect(Collectors.toUnmodifiableList()));
+    }
+  }
 }
diff --git a/app/src/test/java/org/astraea/admin/QuotaTest.java b/app/src/test/java/org/astraea/admin/QuotaTest.java
new file mode 100644
index 00000000..6ad7ee38
--- /dev/null
+++ b/app/src/test/java/org/astraea/admin/QuotaTest.java
@@ -0,0 +1,20 @@
+package org.astraea.admin;
+
+import java.util.Arrays;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class QuotaTest {
+
+  @Test
+  void testTarget() {
+    Arrays.stream(Quota.Target.values())
+        .forEach(t -> Assertions.assertEquals(t, Quota.target(t.nameOfKafka())));
+  }
+
+  @Test
+  void testAction() {
+    Arrays.stream(Quota.Action.values())
+        .forEach(t -> Assertions.assertEquals(t, Quota.action(t.nameOfKafka())));
+  }
+}