Repository: skiptests/astraea
Pull request #211 (2022-02-07 14:05:35)

Enable JMX on kafka tools
這樣方便我們觀察performance tool (by JMX)


Patch:
diff --git a/app/src/main/java/org/astraea/performance/Distribution.java b/app/src/main/java/org/astraea/performance/Distribution.java
index 8148aad6..bbfc5955 100644
--- a/app/src/main/java/org/astraea/performance/Distribution.java
+++ b/app/src/main/java/org/astraea/performance/Distribution.java
@@ -19,6 +19,8 @@ public interface Distribution {
       var name = rawArgument.split(":")[0];
 
       switch (name) {
+        case "fixed":
+          return fixed();
         case "uniform":
           if (args.length > 1) return uniform(Integer.parseInt(args[1]));
           return uniform();
@@ -34,6 +36,15 @@ public interface Distribution {
     }
   }
 
+  static Distribution fixed() {
+    var rand = new Random();
+    return fixed(rand.nextLong());
+  }
+
+  static Distribution fixed(long value) {
+    return () -> value;
+  }
+
   /** A distribution for providing a random long number from range [0, 2147483647) */
   static Distribution uniform() {
     return uniform(Integer.MAX_VALUE);
diff --git a/app/src/main/java/org/astraea/performance/Manager.java b/app/src/main/java/org/astraea/performance/Manager.java
index 2fb59b53..56a9396b 100644
--- a/app/src/main/java/org/astraea/performance/Manager.java
+++ b/app/src/main/java/org/astraea/performance/Manager.java
@@ -104,9 +104,8 @@ public class Manager {
   }
 
   /** Randomly choose a key according to the distribution. */
-  public Optional<byte[]> getKey() {
-    if (distribution == null) return Optional.empty();
-    return Optional.of((String.valueOf(distribution.get())).getBytes());
+  public byte[] getKey() {
+    return (String.valueOf(distribution.get())).getBytes();
   }
 
   /** Randomly generate content before {@link #getContent()} is called. */
diff --git a/app/src/main/java/org/astraea/performance/Performance.java b/app/src/main/java/org/astraea/performance/Performance.java
index 5a4cd928..8846a60b 100644
--- a/app/src/main/java/org/astraea/performance/Performance.java
+++ b/app/src/main/java/org/astraea/performance/Performance.java
@@ -185,7 +185,7 @@ public class Performance {
             .sender()
             .topic(param.topic)
             .partition(partitions.get(rand.nextInt(partitions.size())))
-            .key(manager.getKey().orElse(null))
+            .key(manager.getKey())
             .value(payload.get())
             .timestamp(start)
             .run()
@@ -310,9 +310,9 @@ public class Performance {
     @Parameter(
         names = {"--key.distribution"},
         description =
-            "String: Distribution name. Available distribution names: \"uniform\", \"zipfian\", \"latest\". Default: (No key)",
+            "String: Distribution name. Available distribution names: \"fixed\" \"uniform\", \"zipfian\", \"latest\". Default: uniform",
         converter = Distribution.DistributionConverter.class)
-    Distribution distribution = null;
+    Distribution distribution = Distribution.uniform();
 
     @Parameter(
         names = {"--specify.broker"},
diff --git a/docker/start_kafka_tool.sh b/docker/start_kafka_tool.sh
index b51a116d..65997236 100755
--- a/docker/start_kafka_tool.sh
+++ b/docker/start_kafka_tool.sh
@@ -9,6 +9,10 @@ declare -r DOCKER_FOLDER=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null
 declare -r DOCKERFILE=$DOCKER_FOLDER/kafka_tool.dockerfile
 declare -r BUILD=${BUILD:-false}
 declare -r RUN=${RUN:-true}
+declare -r ADDRESS=$([[ "$(which ipconfig)" != "" ]] && ipconfig getifaddr en0 || hostname -i)
+declare -r JMX_PORT=${JMX_PORT:-$(($(($RANDOM % 10000)) + 10000))}
+declare -r JMX_OPTS="-Dcom.sun.management.jmxremote -Dcom.sun.management.jmxremote.authenticate=false -Dcom.sun.management.jmxremote.ssl=false \
+                     -Dcom.sun.management.jmxremote.port=$JMX_PORT -Dcom.sun.management.jmxremote.rmi.port=$JMX_PORT -Djava.rmi.server.hostname=$ADDRESS"
 
 # ===================================[functions]===================================
 
@@ -72,9 +76,11 @@ function buildImageIfNeed() {
 
 function runContainer() {
   local args=$1
+  echo "JMX address: $ADDRESS:$JMX_PORT"
   docker run --rm --init \
+    -p $JMX_PORT:$JMX_PORT \
     $IMAGE_NAME \
-    /bin/bash -c "java -jar /tmp/app.jar $args"
+    /bin/bash -c "java $JMX_OPTS -jar /tmp/app.jar $args"
 }
 
 # ===================================[main]===================================


Test patch:
diff --git a/app/src/test/java/org/astraea/performance/DistributionTest.java b/app/src/test/java/org/astraea/performance/DistributionTest.java
index 2cb38401..6ff1e06f 100644
--- a/app/src/test/java/org/astraea/performance/DistributionTest.java
+++ b/app/src/test/java/org/astraea/performance/DistributionTest.java
@@ -1,9 +1,23 @@
 package org.astraea.performance;
 
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 public class DistributionTest {
+
+  @Test
+  void testFixed() {
+    var distribution = Distribution.fixed();
+    Assertions.assertEquals(
+        1,
+        IntStream.range(0, 10)
+            .mapToObj(ignored -> distribution.get())
+            .collect(Collectors.toSet())
+            .size());
+  }
+
   @Test
   void testLatest() throws InterruptedException {
     var distribution = Distribution.latest();
diff --git a/app/src/test/java/org/astraea/performance/ManagerTest.java b/app/src/test/java/org/astraea/performance/ManagerTest.java
index 57f72668..31c51912 100644
--- a/app/src/test/java/org/astraea/performance/ManagerTest.java
+++ b/app/src/test/java/org/astraea/performance/ManagerTest.java
@@ -60,18 +60,18 @@ public class ManagerTest {
 
     argument.distribution = Distribution.uniform();
     var manager = new Manager(argument, List.of(), List.of());
-    Assertions.assertTrue(manager.getKey().isPresent());
+    Assertions.assertTrue(manager.getKey().length > 0);
 
     argument.distribution = Distribution.zipfian(10);
     manager = new Manager(argument, List.of(), List.of());
-    Assertions.assertTrue(manager.getKey().isPresent());
+    Assertions.assertTrue(manager.getKey().length > 0);
 
     argument.distribution = Distribution.latest();
     manager = new Manager(argument, List.of(), List.of());
-    Assertions.assertTrue(manager.getKey().isPresent());
+    Assertions.assertTrue(manager.getKey().length > 0);
 
-    argument.distribution = null;
+    argument.distribution = Distribution.fixed();
     manager = new Manager(argument, List.of(), List.of());
-    Assertions.assertTrue(manager.getKey().isEmpty());
+    Assertions.assertTrue(manager.getKey().length > 0);
   }
 }