Repository: skiptests/astraea
Pull request #349 (2022-05-16 18:26:13)

Enable to read/write quotas by restful
新增可以讀寫quota的restful APIs


Patch:
diff --git a/app/src/main/java/org/astraea/admin/Builder.java b/app/src/main/java/org/astraea/admin/Builder.java
index 7eeec61d..1950b67c 100644
--- a/app/src/main/java/org/astraea/admin/Builder.java
+++ b/app/src/main/java/org/astraea/admin/Builder.java
@@ -654,20 +654,22 @@ public class Builder {
 
         @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());
+          var q = new ArrayList<ClientQuotaAlteration.Op>();
+          if (produceRate != Integer.MAX_VALUE)
+            q.add(new ClientQuotaAlteration.Op("producer_byte_rate", (double) produceRate));
+          if (consumeRate != Integer.MAX_VALUE)
+            q.add(new ClientQuotaAlteration.Op("consumer_byte_rate", (double) consumeRate));
+          if (!q.isEmpty())
+            Utils.handleException(
+                () ->
+                    admin
+                        .alterClientQuotas(
+                            List.of(
+                                new ClientQuotaAlteration(
+                                    new ClientQuotaEntity(Map.of(ClientQuotaEntity.CLIENT_ID, id)),
+                                    q)))
+                        .all()
+                        .get());
         }
       };
     }
diff --git a/app/src/main/java/org/astraea/admin/Quota.java b/app/src/main/java/org/astraea/admin/Quota.java
index e955ab02..5c4d0584 100644
--- a/app/src/main/java/org/astraea/admin/Quota.java
+++ b/app/src/main/java/org/astraea/admin/Quota.java
@@ -41,7 +41,7 @@ public class Quota {
         .orElseThrow(() -> new RuntimeException("unknown: " + value));
   }
 
-  enum Target {
+  public enum Target {
     USER(ClientQuotaEntity.USER),
     CLIENT_ID(ClientQuotaEntity.CLIENT_ID),
     IP(ClientQuotaEntity.IP);
@@ -51,12 +51,12 @@ public class Quota {
       this.nameOfKafka = nameOfKafka;
     }
 
-    String nameOfKafka() {
+    public String nameOfKafka() {
       return nameOfKafka;
     }
   }
 
-  enum Action {
+  public enum Action {
     PRODUCER_BYTE_RATE("producer_byte_rate"),
     CONSUMER_BYTE_RATE("consumer_byte_rate"),
     REQUEST_PERCENTAGE("request_percentage"),
@@ -68,7 +68,7 @@ public class Quota {
       this.nameOfKafka = nameOfKafka;
     }
 
-    String nameOfKafka() {
+    public String nameOfKafka() {
       return nameOfKafka;
     }
   }
diff --git a/app/src/main/java/org/astraea/web/QuotaHandler.java b/app/src/main/java/org/astraea/web/QuotaHandler.java
new file mode 100644
index 00000000..bcfb6be5
--- /dev/null
+++ b/app/src/main/java/org/astraea/web/QuotaHandler.java
@@ -0,0 +1,87 @@
+package org.astraea.web;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import org.astraea.admin.Admin;
+
+public class QuotaHandler implements Handler {
+
+  static final String IP_KEY = "ip";
+  static final String CLIENT_ID_KEY = "client-id";
+  static final String CONNECTION_RATE_KEY = "connection-rate";
+  static final String PRODUCE_RATE_KEY = "produce-rate";
+  static final String CONSUME_RATE_KEY = "consume-rate";
+
+  private final Admin admin;
+
+  QuotaHandler(Admin admin) {
+    this.admin = admin;
+  }
+
+  @Override
+  public Quotas get(Optional<String> target, Map<String, String> queries) {
+    if (queries.containsKey(IP_KEY))
+      return new Quotas(admin.quotas(org.astraea.admin.Quota.Target.IP, queries.get(IP_KEY)));
+    if (queries.containsKey(CLIENT_ID_KEY))
+      return new Quotas(
+          admin.quotas(org.astraea.admin.Quota.Target.CLIENT_ID, queries.get(CLIENT_ID_KEY)));
+    return new Quotas(admin.quotas());
+  }
+
+  @Override
+  public JsonObject post(PostRequest request) {
+    if (request.get(IP_KEY).isPresent()) {
+      admin
+          .quotaCreator()
+          .ip(request.value(IP_KEY))
+          .connectionRate(request.intValue(CONNECTION_RATE_KEY))
+          .create();
+      return new Quotas(admin.quotas(org.astraea.admin.Quota.Target.IP, request.value(IP_KEY)));
+    }
+    if (request.get(CLIENT_ID_KEY).isPresent()) {
+      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))
+          .create();
+      return new Quotas(
+          admin.quotas(org.astraea.admin.Quota.Target.CLIENT_ID, request.value(CLIENT_ID_KEY)));
+    }
+    return ErrorObject.for404("You must define either " + CLIENT_ID_KEY + " or " + IP_KEY);
+  }
+
+  static class Quota implements JsonObject {
+    final String target;
+
+    final String targetValue;
+    final String action;
+    final double actionValue;
+
+    public Quota(org.astraea.admin.Quota quota) {
+      this(
+          quota.target().nameOfKafka(),
+          quota.targetValue(),
+          quota.action().nameOfKafka(),
+          quota.actionValue());
+    }
+
+    public Quota(String target, String targetValue, String action, double actionValue) {
+      this.target = target;
+      this.targetValue = targetValue;
+      this.action = action;
+      this.actionValue = actionValue;
+    }
+  }
+
+  static class Quotas implements JsonObject {
+    final List<Quota> quotas;
+
+    Quotas(Collection<org.astraea.admin.Quota> quotas) {
+      this.quotas = quotas.stream().map(Quota::new).collect(Collectors.toUnmodifiableList());
+    }
+  }
+}
diff --git a/app/src/main/java/org/astraea/web/WebService.java b/app/src/main/java/org/astraea/web/WebService.java
index a3dd08d6..d0821feb 100644
--- a/app/src/main/java/org/astraea/web/WebService.java
+++ b/app/src/main/java/org/astraea/web/WebService.java
@@ -19,6 +19,7 @@ public class WebService {
     server.createContext("/groups", new GroupHandler(Admin.of(arg.configs())));
     server.createContext("/brokers", new BrokerHandler(Admin.of(arg.configs())));
     server.createContext("/producers", new ProducerHandler(Admin.of(arg.configs())));
+    server.createContext("/quotas", new QuotaHandler(Admin.of(arg.configs())));
     server.start();
   }
 


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 572ba7cd..cf76715e 100644
--- a/app/src/test/java/org/astraea/admin/AdminTest.java
+++ b/app/src/test/java/org/astraea/admin/AdminTest.java
@@ -445,6 +445,17 @@ public class AdminTest extends RequireBrokerCluster {
     }
   }
 
+  @Test
+  void testMultipleIpQuota() throws InterruptedException {
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.quotaCreator().ip("192.168.11.11").connectionRate(10).create();
+      admin.quotaCreator().ip("192.168.11.11").connectionRate(12).create();
+      admin.quotaCreator().ip("192.168.11.11").connectionRate(9).create();
+      TimeUnit.SECONDS.sleep(2);
+      Assertions.assertEquals(1, admin.quotas(Quota.Target.IP, "192.168.11.11").size());
+    }
+  }
+
   @Test
   void testClientQuota() throws InterruptedException {
     try (var admin = Admin.of(bootstrapServers())) {
@@ -500,4 +511,14 @@ public class AdminTest extends RequireBrokerCluster {
               .collect(Collectors.toUnmodifiableList()));
     }
   }
+
+  @Test
+  void testMultipleClientQuota() throws InterruptedException {
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.quotaCreator().clientId("my-id").consumeRate(100).create();
+      admin.quotaCreator().clientId("my-id").produceRate(999).create();
+      TimeUnit.SECONDS.sleep(2);
+      Assertions.assertEquals(2, admin.quotas(Quota.Target.CLIENT_ID, "my-id").size());
+    }
+  }
 }
diff --git a/app/src/test/java/org/astraea/web/QuotaHandlerTest.java b/app/src/test/java/org/astraea/web/QuotaHandlerTest.java
new file mode 100644
index 00000000..715ca4e2
--- /dev/null
+++ b/app/src/test/java/org/astraea/web/QuotaHandlerTest.java
@@ -0,0 +1,76 @@
+package org.astraea.web;
+
+import java.util.Map;
+import java.util.Optional;
+import org.astraea.admin.Admin;
+import org.astraea.admin.Quota;
+import org.astraea.service.RequireBrokerCluster;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class QuotaHandlerTest extends RequireBrokerCluster {
+
+  @Test
+  void testCreateQuota() {
+    var ip = "192.168.10.11";
+    try (Admin admin = Admin.of(bootstrapServers())) {
+      var handler = new QuotaHandler(admin);
+
+      var result =
+          Assertions.assertInstanceOf(
+              QuotaHandler.Quotas.class,
+              handler.post(
+                  PostRequest.of(
+                      Map.of(QuotaHandler.IP_KEY, ip, QuotaHandler.CONNECTION_RATE_KEY, "10"))));
+      Assertions.assertEquals(1, result.quotas.size());
+      Assertions.assertEquals(
+          Quota.Target.IP.nameOfKafka(), result.quotas.iterator().next().target);
+      Assertions.assertEquals(ip, result.quotas.iterator().next().targetValue);
+      Assertions.assertEquals(
+          Quota.Action.IP_CONNECTION_RATE.nameOfKafka(), result.quotas.iterator().next().action);
+      Assertions.assertEquals(10, result.quotas.iterator().next().actionValue);
+    }
+  }
+
+  @Test
+  void testQuery() {
+    var ip0 = "192.168.10.11";
+    var ip1 = "192.168.10.12";
+    try (Admin admin = Admin.of(bootstrapServers())) {
+      var handler = new QuotaHandler(admin);
+
+      handler.post(
+          PostRequest.of(Map.of(QuotaHandler.IP_KEY, ip0, QuotaHandler.CONNECTION_RATE_KEY, "10")));
+      handler.post(
+          PostRequest.of(Map.of(QuotaHandler.IP_KEY, ip1, QuotaHandler.CONNECTION_RATE_KEY, "20")));
+      Assertions.assertEquals(
+          1, handler.get(Optional.empty(), Map.of(QuotaHandler.IP_KEY, ip0)).quotas.size());
+      Assertions.assertEquals(
+          1, handler.get(Optional.empty(), Map.of(QuotaHandler.IP_KEY, ip1)).quotas.size());
+    }
+  }
+
+  @Test
+  void testQueryNonexistentQuota() {
+    var ip = "192.168.10.11";
+    try (Admin admin = Admin.of(bootstrapServers())) {
+      var handler = new QuotaHandler(admin);
+      Assertions.assertEquals(
+          0,
+          Assertions.assertInstanceOf(
+                  QuotaHandler.Quotas.class,
+                  handler.get(Optional.empty(), Map.of(Quota.Target.IP.nameOfKafka(), "unknown")))
+              .quotas
+              .size());
+
+      Assertions.assertEquals(
+          0,
+          Assertions.assertInstanceOf(
+                  QuotaHandler.Quotas.class,
+                  handler.get(
+                      Optional.empty(), Map.of(Quota.Target.CLIENT_ID.nameOfKafka(), "unknown")))
+              .quotas
+              .size());
+    }
+  }
+}