diff --git a/tools/src/main/java/org/apache/kafka/trogdor/common/WorkerUtils.java b/tools/src/main/java/org/apache/kafka/trogdor/common/WorkerUtils.java index 99c13c09febaf..98dbf3836b692 100644 --- a/tools/src/main/java/org/apache/kafka/trogdor/common/WorkerUtils.java +++ b/tools/src/main/java/org/apache/kafka/trogdor/common/WorkerUtils.java @@ -74,6 +74,23 @@ public static int perSecToPerPeriod(float perSec, long periodMs) { return (int) perPeriod; } + /** + * Adds all properties from commonConf and then from clientConf to given 'props' (in + * that order, over-writing properties with the same keys). + * @param props Properties object that may contain zero or more properties + * @param commonConf Map with common client properties + * @param clientConf Map with client properties + */ + public static void addConfigsToProperties( + Properties props, Map commonConf, Map clientConf) { + for (Map.Entry commonEntry : commonConf.entrySet()) { + props.setProperty(commonEntry.getKey(), commonEntry.getValue()); + } + for (Map.Entry entry : clientConf.entrySet()) { + props.setProperty(entry.getKey(), entry.getValue()); + } + } + private static final int CREATE_TOPICS_REQUEST_TIMEOUT = 25000; private static final int CREATE_TOPICS_CALL_TIMEOUT = 180000; private static final int MAX_CREATE_TOPICS_BATCH_SIZE = 10; @@ -85,6 +102,9 @@ public static int perSecToPerPeriod(float perSec, long periodMs) { * * @param log The logger to use. * @param bootstrapServers The bootstrap server list. + * @param commonClientConf Common client config + * @param adminClientConf AdminClient config. This config has precedence over fields in + * common client config. * @param topics Maps topic names to partition assignments. * @param failOnExisting If true, the method will throw TopicExistsException if one or * more topics already exist. Otherwise, the existing topics are @@ -93,12 +113,14 @@ public static int perSecToPerPeriod(float perSec, long periodMs) { * number of partitions, the method throws RuntimeException. */ public static void createTopics( - Logger log, String bootstrapServers, + Logger log, String bootstrapServers, Map commonClientConf, + Map adminClientConf, Map topics, boolean failOnExisting) throws Throwable { // this method wraps the call to createTopics() that takes admin client, so that we can // unit test the functionality with MockAdminClient. The exception is caught and // re-thrown so that admin client is closed when the method returns. - try (AdminClient adminClient = createAdminClient(bootstrapServers)) { + try (AdminClient adminClient + = createAdminClient(bootstrapServers, commonClientConf, adminClientConf)) { createTopics(log, adminClient, topics, failOnExisting); } catch (Exception e) { log.warn("Failed to create or verify topics {}", topics, e); @@ -227,10 +249,15 @@ private static void verifyTopics( } } - private static AdminClient createAdminClient(String bootstrapServers) { + private static AdminClient createAdminClient( + String bootstrapServers, Map commonClientConf, + Map adminClientConf) { Properties props = new Properties(); props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); props.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, CREATE_TOPICS_REQUEST_TIMEOUT); + // first add common client config, and then admin client config to properties, possibly + // over-writing default or common properties. + addConfigsToProperties(props, commonClientConf, adminClientConf); return AdminClient.create(props); } } diff --git a/tools/src/main/java/org/apache/kafka/trogdor/task/TaskSpec.java b/tools/src/main/java/org/apache/kafka/trogdor/task/TaskSpec.java index 84ed75a8c8ac9..af7a76f854119 100644 --- a/tools/src/main/java/org/apache/kafka/trogdor/task/TaskSpec.java +++ b/tools/src/main/java/org/apache/kafka/trogdor/task/TaskSpec.java @@ -21,6 +21,8 @@ import com.fasterxml.jackson.annotation.JsonTypeInfo; import org.apache.kafka.trogdor.common.JsonUtil; +import java.util.Collections; +import java.util.Map; import java.util.Objects; @@ -102,4 +104,8 @@ public final int hashCode() { public String toString() { return JsonUtil.toJsonString(this); } + + protected Map configOrEmptyMap(Map config) { + return (config == null) ? Collections.emptyMap() : config; + } } diff --git a/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchSpec.java b/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchSpec.java index 7b1bedd183036..ec6e3096469b6 100644 --- a/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchSpec.java +++ b/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchSpec.java @@ -26,7 +26,6 @@ import java.util.Collections; import java.util.Map; -import java.util.TreeMap; import java.util.Set; /** @@ -45,6 +44,8 @@ public class ProduceBenchSpec extends TaskSpec { private final PayloadGenerator keyGenerator; private final PayloadGenerator valueGenerator; private final Map producerConf; + private final Map adminClientConf; + private final Map commonClientConf; private final int totalTopics; private final int activeTopics; private final String topicPrefix; @@ -61,6 +62,8 @@ public ProduceBenchSpec(@JsonProperty("startMs") long startMs, @JsonProperty("keyGenerator") PayloadGenerator keyGenerator, @JsonProperty("valueGenerator") PayloadGenerator valueGenerator, @JsonProperty("producerConf") Map producerConf, + @JsonProperty("commonClientConf") Map commonClientConf, + @JsonProperty("adminClientConf") Map adminClientConf, @JsonProperty("totalTopics") int totalTopics, @JsonProperty("activeTopics") int activeTopics, @JsonProperty("topicPrefix") String topicPrefix, @@ -75,7 +78,9 @@ public ProduceBenchSpec(@JsonProperty("startMs") long startMs, new SequentialPayloadGenerator(4, 0) : keyGenerator; this.valueGenerator = valueGenerator == null ? new ConstantPayloadGenerator(512, new byte[0]) : valueGenerator; - this.producerConf = (producerConf == null) ? new TreeMap() : producerConf; + this.producerConf = configOrEmptyMap(producerConf); + this.commonClientConf = configOrEmptyMap(commonClientConf); + this.adminClientConf = configOrEmptyMap(adminClientConf); this.totalTopics = totalTopics; this.activeTopics = activeTopics; this.topicPrefix = (topicPrefix == null) ? DEFAULT_TOPIC_PREFIX : topicPrefix; @@ -120,6 +125,16 @@ public Map producerConf() { return producerConf; } + @JsonProperty + public Map commonClientConf() { + return commonClientConf; + } + + @JsonProperty + public Map adminClientConf() { + return adminClientConf; + } + @JsonProperty public int totalTopics() { return totalTopics; diff --git a/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchWorker.java b/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchWorker.java index e291bae408842..a891b83bcc163 100644 --- a/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchWorker.java +++ b/tools/src/main/java/org/apache/kafka/trogdor/workload/ProduceBenchWorker.java @@ -109,9 +109,11 @@ public void run() { Map newTopics = new HashMap<>(); for (int i = 0; i < spec.totalTopics(); i++) { String name = topicIndexToName(i); - newTopics.put(name, new NewTopic(name, spec.numPartitions(), spec.replicationFactor())); + newTopics.put(name, new NewTopic(name, spec.numPartitions(), + spec.replicationFactor())); } - WorkerUtils.createTopics(log, spec.bootstrapServers(), newTopics, false); + WorkerUtils.createTopics(log, spec.bootstrapServers(), spec.commonClientConf(), + spec.adminClientConf(), newTopics, false); executor.submit(new SendRecords()); } catch (Throwable e) { @@ -182,9 +184,9 @@ public class SendRecords implements Callable { new StatusUpdater(histogram), 1, 1, TimeUnit.MINUTES); Properties props = new Properties(); props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, spec.bootstrapServers()); - for (Map.Entry entry : spec.producerConf().entrySet()) { - props.setProperty(entry.getKey(), entry.getValue()); - } + // add common client configs to producer properties, and then user-specified producer + // configs + WorkerUtils.addConfigsToProperties(props, spec.commonClientConf(), spec.producerConf()); this.producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer()); this.keys = new PayloadIterator(spec.keyGenerator()); this.values = new PayloadIterator(spec.valueGenerator()); diff --git a/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorker.java b/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorker.java index a05785c04e0d2..08b11ac560368 100644 --- a/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorker.java +++ b/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorker.java @@ -124,7 +124,7 @@ public void run() { throw new ConfigException("Invalid null or empty partitionAssignments."); } WorkerUtils.createTopics( - log, spec.bootstrapServers(), + log, spec.bootstrapServers(), spec.commonClientConf(), spec.adminClientConf(), Collections.singletonMap(TOPIC_NAME, new NewTopic(TOPIC_NAME, spec.partitionAssignments())), true); @@ -184,6 +184,8 @@ class ProducerRunnable implements Runnable { props.put(ProducerConfig.CLIENT_ID_CONFIG, "producer." + id); props.put(ProducerConfig.ACKS_CONFIG, "all"); props.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, 105000); + // user may over-write the defaults with common client config and producer config + WorkerUtils.addConfigsToProperties(props, spec.commonClientConf(), spec.producerConf()); producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer()); int perPeriod = WorkerUtils. @@ -275,6 +277,8 @@ class ConsumerRunnable implements Runnable { props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); props.put(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, 105000); props.put(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 100000); + // user may over-write the defaults with common client config and consumer config + WorkerUtils.addConfigsToProperties(props, spec.commonClientConf(), spec.consumerConf()); consumer = new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer()); consumer.subscribe(Collections.singleton(TOPIC_NAME)); diff --git a/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorkloadSpec.java b/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorkloadSpec.java index 00bd833c05265..3d0e3ef2e6b4a 100644 --- a/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorkloadSpec.java +++ b/tools/src/main/java/org/apache/kafka/trogdor/workload/RoundTripWorkloadSpec.java @@ -26,6 +26,7 @@ import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.NavigableMap; import java.util.Set; import java.util.TreeMap; @@ -41,12 +42,20 @@ public class RoundTripWorkloadSpec extends TaskSpec { private final NavigableMap> partitionAssignments; private final PayloadGenerator valueGenerator; private final int maxMessages; + private final Map commonClientConf; + private final Map producerConf; + private final Map consumerConf; + private final Map adminClientConf; @JsonCreator public RoundTripWorkloadSpec(@JsonProperty("startMs") long startMs, @JsonProperty("durationMs") long durationMs, @JsonProperty("clientNode") String clientNode, @JsonProperty("bootstrapServers") String bootstrapServers, + @JsonProperty("commonClientConf") Map commonClientConf, + @JsonProperty("adminClientConf") Map adminClientConf, + @JsonProperty("consumerConf") Map consumerConf, + @JsonProperty("producerConf") Map producerConf, @JsonProperty("targetMessagesPerSec") int targetMessagesPerSec, @JsonProperty("partitionAssignments") NavigableMap> partitionAssignments, @JsonProperty("valueGenerator") PayloadGenerator valueGenerator, @@ -60,6 +69,10 @@ public RoundTripWorkloadSpec(@JsonProperty("startMs") long startMs, this.valueGenerator = valueGenerator == null ? new UniformRandomPayloadGenerator(32, 123, 10) : valueGenerator; this.maxMessages = maxMessages; + this.commonClientConf = configOrEmptyMap(commonClientConf); + this.adminClientConf = configOrEmptyMap(adminClientConf); + this.producerConf = configOrEmptyMap(producerConf); + this.consumerConf = configOrEmptyMap(consumerConf); } @JsonProperty @@ -92,6 +105,26 @@ public int maxMessages() { return maxMessages; } + @JsonProperty + public Map commonClientConf() { + return commonClientConf; + } + + @JsonProperty + public Map adminClientConf() { + return adminClientConf; + } + + @JsonProperty + public Map producerConf() { + return producerConf; + } + + @JsonProperty + public Map consumerConf() { + return consumerConf; + } + @Override public TaskController newController(String id) { return new TaskController() { diff --git a/tools/src/test/java/org/apache/kafka/trogdor/common/JsonSerializationTest.java b/tools/src/test/java/org/apache/kafka/trogdor/common/JsonSerializationTest.java index dee7614292a65..76b206bf135c6 100644 --- a/tools/src/test/java/org/apache/kafka/trogdor/common/JsonSerializationTest.java +++ b/tools/src/test/java/org/apache/kafka/trogdor/common/JsonSerializationTest.java @@ -49,8 +49,8 @@ public void testDeserializationDoesNotProduceNulls() throws Exception { verify(new WorkerRunning(null, 0, null)); verify(new WorkerStopping(null, 0, null)); verify(new ProduceBenchSpec(0, 0, null, null, - 0, 0, null, null, null, 0, 0, "test-topic", 1, (short) 3)); - verify(new RoundTripWorkloadSpec(0, 0, null, null, + 0, 0, null, null, null, null, null, 0, 0, "test-topic", 1, (short) 3)); + verify(new RoundTripWorkloadSpec(0, 0, null, null, null, null, null, null, 0, null, null, 0)); verify(new SampleTaskSpec(0, 0, 0, null)); } diff --git a/tools/src/test/java/org/apache/kafka/trogdor/common/WorkerUtilsTest.java b/tools/src/test/java/org/apache/kafka/trogdor/common/WorkerUtilsTest.java index 22b784600f650..fbe23899bc9dc 100644 --- a/tools/src/test/java/org/apache/kafka/trogdor/common/WorkerUtilsTest.java +++ b/tools/src/test/java/org/apache/kafka/trogdor/common/WorkerUtilsTest.java @@ -20,6 +20,7 @@ import org.apache.kafka.clients.admin.TopicDescription; +import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.common.TopicPartitionInfo; import org.apache.kafka.common.Node; @@ -41,6 +42,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Properties; public class WorkerUtilsTest { @@ -208,4 +210,56 @@ public void testCreateNonExistingTopicsWithZeroTopicsDoesNothing() throws Throwa assertEquals(0, adminClient.listTopics().names().get().size()); } + @Test + public void testAddConfigsToPropertiesAddsAllConfigs() { + Properties props = new Properties(); + props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put(ProducerConfig.ACKS_CONFIG, "all"); + + Properties resultProps = new Properties(); + resultProps.putAll(props); + resultProps.put(ProducerConfig.CLIENT_ID_CONFIG, "test-client"); + resultProps.put(ProducerConfig.LINGER_MS_CONFIG, "1000"); + + WorkerUtils.addConfigsToProperties( + props, + Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, "test-client"), + Collections.singletonMap(ProducerConfig.LINGER_MS_CONFIG, "1000")); + assertEquals(resultProps, props); + } + + @Test + public void testCommonConfigOverwritesDefaultProps() { + Properties props = new Properties(); + props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put(ProducerConfig.ACKS_CONFIG, "all"); + + Properties resultProps = new Properties(); + resultProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + resultProps.put(ProducerConfig.ACKS_CONFIG, "1"); + resultProps.put(ProducerConfig.LINGER_MS_CONFIG, "1000"); + + WorkerUtils.addConfigsToProperties( + props, + Collections.singletonMap(ProducerConfig.ACKS_CONFIG, "1"), + Collections.singletonMap(ProducerConfig.LINGER_MS_CONFIG, "1000")); + assertEquals(resultProps, props); + } + + @Test + public void testClientConfigOverwritesBothDefaultAndCommonConfigs() { + Properties props = new Properties(); + props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put(ProducerConfig.ACKS_CONFIG, "all"); + + Properties resultProps = new Properties(); + resultProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + resultProps.put(ProducerConfig.ACKS_CONFIG, "0"); + + WorkerUtils.addConfigsToProperties( + props, + Collections.singletonMap(ProducerConfig.ACKS_CONFIG, "1"), + Collections.singletonMap(ProducerConfig.ACKS_CONFIG, "0")); + assertEquals(resultProps, props); + } }