Skip to content
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.admin.TopicDescription;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.config.SaslConfigs;
import org.apache.kafka.common.config.SslConfigs;
import org.apache.kafka.common.errors.NotEnoughReplicasException;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.errors.TopicExistsException;
Expand All @@ -33,6 +35,7 @@
import org.slf4j.Logger;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
Expand All @@ -44,6 +47,21 @@
* Utilities for Trogdor TaskWorkers.
*/
public final class WorkerUtils {

private static final List<String> CLIENT_SECURITY_KEYS = Arrays.asList(
AdminClientConfig.SECURITY_PROTOCOL_CONFIG,
SaslConfigs.SASL_MECHANISM,
SaslConfigs.SASL_JAAS_CONFIG,
SaslConfigs.SASL_KERBEROS_SERVICE_NAME,
SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG,
SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG,
SslConfigs.SSL_KEY_PASSWORD_CONFIG,
SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG,
SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG,
SslConfigs.SSL_PROTOCOL_CONFIG,
SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you not use .withClientSslSupport().withClientSaslSupport()?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did not know about those methods. Let me look. Thanks!

);

/**
* Handle an exception in a TaskWorker.
*
Expand Down Expand Up @@ -93,12 +111,12 @@ 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<String, String> clientConf,
Map<String, NewTopic> 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, clientConf)) {
createTopics(log, adminClient, topics, failOnExisting);
} catch (Exception e) {
log.warn("Failed to create or verify topics {}", topics, e);
Expand Down Expand Up @@ -227,10 +245,18 @@ private static void verifyTopics(
}
}

private static AdminClient createAdminClient(String bootstrapServers) {
private static AdminClient createAdminClient(
String bootstrapServers, Map<String, String> clientConf) {
Properties props = new Properties();
props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
props.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, CREATE_TOPICS_REQUEST_TIMEOUT);

// copy security related props
for (Map.Entry<String, String> clientPropsEntry: clientConf.entrySet()) {
if (CLIENT_SECURITY_KEYS.contains(clientPropsEntry.getKey())) {
props.put(clientPropsEntry.getKey(), clientPropsEntry.getValue());
}
}
return AdminClient.create(props);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,8 @@ public void run() {
String name = topicIndexToName(i);
newTopics.put(name, new NewTopic(name, spec.numPartitions(), spec.replicationFactor()));
}
WorkerUtils.createTopics(log, spec.bootstrapServers(), newTopics, false);
WorkerUtils.createTopics(log, spec.bootstrapServers(), spec.producerConf(),
newTopics, false);

executor.submit(new SendRecords());
} catch (Throwable e) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ public void run() {
throw new ConfigException("Invalid null or empty partitionAssignments.");
}
WorkerUtils.createTopics(
log, spec.bootstrapServers(),
log, spec.bootstrapServers(), Collections.<String, String>emptyMap(),
Collections.singletonMap(TOPIC_NAME,
new NewTopic(TOPIC_NAME, spec.partitionAssignments())),
true);
Expand Down