Repository: skiptests/astraea
Pull request #117 (2021-11-21 09:23:07)

Enable performance tool to apply compression
壓縮演算法也是測效能很重要的一環


Patch:
diff --git a/app/src/main/java/org/astraea/performance/Performance.java b/app/src/main/java/org/astraea/performance/Performance.java
index 8142173d..563d2206 100644
--- a/app/src/main/java/org/astraea/performance/Performance.java
+++ b/app/src/main/java/org/astraea/performance/Performance.java
@@ -1,6 +1,8 @@
 package org.astraea.performance;
 
+import com.beust.jcommander.IStringConverter;
 import com.beust.jcommander.Parameter;
+import com.beust.jcommander.ParameterException;
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Map;
@@ -10,8 +12,11 @@ import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
+import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
 import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.record.CompressionType;
 import org.astraea.argument.ArgumentUtil;
 import org.astraea.argument.BasicArgumentWithPropFile;
 import org.astraea.concurrent.ThreadPool;
@@ -100,7 +105,7 @@ public class Performance {
                         i ->
                             producerExecutor(
                                 Producer.builder()
-                                    .configs(param.props())
+                                    .configs(param.producerProps())
                                     .partitionClassName(param.partitioner)
                                     .build(),
                                 param,
@@ -240,11 +245,36 @@ public class Performance {
         validateWith = ArgumentUtil.NotEmptyString.class)
     String partitioner = DefaultPartitioner.class.getName();
 
+    public Map<String, Object> producerProps() {
+      var props = props();
+      props.put(ProducerConfig.COMPRESSION_TYPE_CONFIG, compression.name);
+      if (!this.jmxServers.isEmpty()) props.put("jmx_servers", this.jmxServers);
+      return props;
+    }
+
+    @Parameter(
+        names = {"--compression"},
+        description =
+            "String: the compression algorithm used by producer. Available algorithm are none, gzip, snappy, lz4, and zstd",
+        converter = CompressionArgument.class)
+    CompressionType compression = CompressionType.NONE;
+  }
+
+  static class CompressionArgument implements IStringConverter<CompressionType> {
+
     @Override
-    public Map<String, Object> props() {
-      Map<String, Object> prop = properties(propFile);
-      if (!this.jmxServers.isEmpty()) prop.put("jmx_servers", this.jmxServers);
-      return prop;
+    public CompressionType convert(String value) {
+      try {
+        return CompressionType.forName(value.toLowerCase());
+      } catch (IllegalArgumentException e) {
+        throw new ParameterException(
+            "the "
+                + value
+                + " is unsupported. The supported algorithms are "
+                + Stream.of(CompressionType.values())
+                    .map(CompressionType::name)
+                    .collect(Collectors.joining(",")));
+      }
     }
   }
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/performance/CompressionArgumentTest.java b/app/src/test/java/org/astraea/performance/CompressionArgumentTest.java
new file mode 100644
index 00000000..30d96a54
--- /dev/null
+++ b/app/src/test/java/org/astraea/performance/CompressionArgumentTest.java
@@ -0,0 +1,41 @@
+package org.astraea.performance;
+
+import com.beust.jcommander.ParameterException;
+import java.util.stream.Stream;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.record.CompressionType;
+import org.astraea.argument.ArgumentUtil;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class CompressionArgumentTest {
+
+  @Test
+  void testConversion() {
+    var arg = new Performance.CompressionArgument();
+    Stream.of(CompressionType.values())
+        .forEach(type -> Assertions.assertEquals(type, arg.convert(type.name)));
+    Assertions.assertThrows(ParameterException.class, () -> arg.convert("aaa"));
+  }
+
+  @Test
+  void testParse() {
+    var arg =
+        ArgumentUtil.parseArgument(
+            new Performance.Argument(),
+            new String[] {"--bootstrap.servers", "aa", "--compression", "gzip"});
+    Assertions.assertEquals(CompressionType.GZIP, arg.compression);
+    Assertions.assertEquals(
+        "gzip", arg.producerProps().get(ProducerConfig.COMPRESSION_TYPE_CONFIG));
+    Assertions.assertNull(arg.props().get(ProducerConfig.COMPRESSION_TYPE_CONFIG));
+  }
+
+  @Test
+  void testJmxServer() {
+    var arg =
+        ArgumentUtil.parseArgument(
+            new Performance.Argument(),
+            new String[] {"--bootstrap.servers", "aa", "--jmx.servers", "aaa"});
+    Assertions.assertEquals("aaa", arg.producerProps().get("jmx_servers"));
+  }
+}