diff --git a/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java b/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java index 3f5968e6f94..974acd3e47f 100755 --- a/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java +++ b/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java @@ -6,6 +6,7 @@ import static com.linkedin.venice.ConfigKeys.KAFKA_PRODUCER_REQUEST_TIMEOUT_MS; import static com.linkedin.venice.ConfigKeys.KAFKA_PRODUCER_RETRIES_CONFIG; import static com.linkedin.venice.ConfigKeys.MULTI_REGION; +import static com.linkedin.venice.ConfigKeys.PUBSUB_BROKER_ADDRESS; import static com.linkedin.venice.ConfigKeys.VENICE_PARTITIONERS; import static com.linkedin.venice.VeniceConstants.DEFAULT_SSL_FACTORY_CLASS_NAME; import static com.linkedin.venice.status.BatchJobHeartbeatConfigs.HEARTBEAT_ENABLED_CONFIG; @@ -866,13 +867,9 @@ public void run() { if (pushJobSetting.isSourceKafka) { if (pushJobSetting.sourceVersionCompressionStrategy == CompressionStrategy.ZSTD_WITH_DICT) { LOGGER.info("Source version uses ZSTD_WITH_DICT. Fetching source dictionary."); - Properties kafkaConsumerProperties = new Properties(); - if (pushJobSetting.enableSSL) { - kafkaConsumerProperties.putAll(this.sslProperties.get()); - } - kafkaConsumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, pushJobSetting.repushSourcePubsubBroker); - ByteBuffer sourceDict = DictionaryUtils - .readDictionaryFromKafka(pushJobSetting.kafkaInputTopic, new VeniceProperties(kafkaConsumerProperties)); + ByteBuffer sourceDict = DictionaryUtils.readDictionaryFromKafka( + pushJobSetting.kafkaInputTopic, + new VeniceProperties(getRepushDictionaryConsumerProperties())); if (sourceDict != null) { pushJobSetting.sourceDictionary = ByteUtils.extractByteArray(sourceDict); } @@ -1666,12 +1663,46 @@ private Optional getCompressionDictionary() throws VeniceException { return Optional.of(emptyPushZstdDictionary.get()); } + /** + * Build the pub-sub consumer properties used to read the repush source topic's compression dictionary. + * Seeds from the full job config so the configured pub-sub consumer adapter factory (e.g. xinfra) and its + * client configs are honored, then overlays SSL and the repush source broker. Seeding only SSL + broker would + * drop {@code pubsub.consumer.adapter.factory.class} and silently fall back to the default Apache Kafka + * consumer, which cannot read the source topic on a non-Kafka (e.g. xinfra) deployment. + */ + private Properties getRepushDictionaryConsumerProperties() { + return buildRepushDictionaryConsumerProperties( + props, + pushJobSetting.enableSSL ? sslProperties.get() : new Properties(), + pushJobSetting.repushSourcePubsubBroker); + } + + /** + * Assemble the pub-sub consumer properties used to read the repush source topic's compression dictionary. The + * source dictionary must be read with the same pub-sub client as the rest of the repush, so seed from the full + * job configuration (which carries {@code pubsub.consumer.adapter.factory.class} and any client-specific + * settings, such as xinfra's routing maps), overlay SSL, and point the broker properties at the repush source + * broker. Seeding only SSL and the broker address drops the adapter-factory class and silently falls back to + * the default Apache Kafka consumer, which cannot read the source topic on a non-Kafka (e.g. xinfra) deployment. + */ + @VisibleForTesting + static Properties buildRepushDictionaryConsumerProperties( + VeniceProperties jobProperties, + Properties sslProperties, + String repushSourcePubsubBroker) { + Properties consumerProperties = jobProperties.toProperties(); + consumerProperties.putAll(sslProperties); + consumerProperties.setProperty(PUBSUB_BROKER_ADDRESS, repushSourcePubsubBroker); + consumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, repushSourcePubsubBroker); + return consumerProperties; + } + private ByteBuffer fetchOrBuildCompressionDictionary() throws VeniceException { // Prepare the param builder, which can be used by different scenarios. KafkaInputDictTrainer.ParamBuilder paramBuilder = new KafkaInputDictTrainer.ParamBuilder() .setKeySchema(AvroCompatibilityHelper.toParsingForm(pushJobSetting.storeKeySchema)) .setNewKMESchemasFromController(pushJobSetting.newKmeSchemasFromController) - .setSslProperties(pushJobSetting.enableSSL ? sslProperties.get() : new Properties()) + .setConsumerProperties(getRepushDictionaryConsumerProperties()) .setCompressionDictSize( props.getInt( COMPRESSION_DICTIONARY_SIZE_LIMIT, @@ -1691,14 +1722,9 @@ private ByteBuffer fetchOrBuildCompressionDictionary() throws VeniceException { return ByteBuffer.wrap(dictTrainer.trainDict()); } else { LOGGER.info("Reading Zstd dictionary from input topic: {}", pushJobSetting.kafkaInputTopic); - // set up ssl properties and kafka consumer properties - Properties kafkaConsumerProperties = new Properties(); - if (pushJobSetting.enableSSL) { - kafkaConsumerProperties.putAll(this.sslProperties.get()); - } - kafkaConsumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, pushJobSetting.repushSourcePubsubBroker); - return DictionaryUtils - .readDictionaryFromKafka(pushJobSetting.kafkaInputTopic, new VeniceProperties(kafkaConsumerProperties)); + return DictionaryUtils.readDictionaryFromKafka( + pushJobSetting.kafkaInputTopic, + new VeniceProperties(getRepushDictionaryConsumerProperties())); } } LOGGER.info( diff --git a/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/input/kafka/KafkaInputDictTrainer.java b/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/input/kafka/KafkaInputDictTrainer.java index b1695dcf2e2..04cfdb6a6cb 100644 --- a/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/input/kafka/KafkaInputDictTrainer.java +++ b/clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/input/kafka/KafkaInputDictTrainer.java @@ -50,7 +50,7 @@ public static class Param { private final String kafkaInputBroker; private final String topicName; private final String keySchema; - private final Properties sslProperties; + private final Properties consumerProperties; private final int compressionDictSize; private final int dictSampleSize; private final CompressionStrategy sourceVersionCompressionStrategy; @@ -62,7 +62,7 @@ public static class Param { this.kafkaInputBroker = builder.kafkaInputBroker; this.topicName = builder.topicName; this.keySchema = builder.keySchema; - this.sslProperties = builder.sslProperties; + this.consumerProperties = builder.consumerProperties; this.compressionDictSize = builder.compressionDictSize; this.dictSampleSize = builder.dictSampleSize; this.sourceVersionCompressionStrategy = builder.sourceVersionCompressionStrategy; @@ -75,7 +75,7 @@ public static class ParamBuilder { private String kafkaInputBroker; private String topicName; private String keySchema; - private Properties sslProperties; + private Properties consumerProperties; private int compressionDictSize; private int dictSampleSize; private CompressionStrategy sourceVersionCompressionStrategy; @@ -97,8 +97,14 @@ public ParamBuilder setKeySchema(String keySchema) { return this; } - public ParamBuilder setSslProperties(Properties sslProperties) { - this.sslProperties = sslProperties; + /** + * Properties the source-topic consumer is built from. These must include the job's pub-sub client + * configuration (e.g. {@code pubsub.consumer.adapter.factory.class} and any client-specific configs such as + * xinfra's) so the dictionary is read with the same pub-sub client as the rest of the repush; otherwise + * {@code PubSubClientsFactory.createConsumerFactory(...)} silently defaults to the Apache Kafka consumer. + */ + public ParamBuilder setConsumerProperties(Properties consumerProperties) { + this.consumerProperties = consumerProperties; return this; } @@ -170,7 +176,7 @@ protected KafkaInputDictTrainer( properties.setProperty(KAFKA_INPUT_TOPIC, param.topicName); properties.setProperty(KAFKA_SOURCE_KEY_SCHEMA_STRING_PROP, param.keySchema); this.sourceTopicName = param.topicName; - properties.putAll(param.sslProperties); + properties.putAll(param.consumerProperties); properties.setProperty(COMPRESSION_DICTIONARY_SIZE_LIMIT, Integer.toString(param.compressionDictSize)); properties.setProperty(COMPRESSION_DICTIONARY_SAMPLE_SIZE, Integer.toString(param.dictSampleSize)); properties diff --git a/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobRepushTest.java b/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobRepushTest.java index f731eafa834..20a033c9694 100644 --- a/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobRepushTest.java +++ b/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobRepushTest.java @@ -1,5 +1,9 @@ package com.linkedin.venice.hadoop; +import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS; +import static com.linkedin.venice.ConfigKeys.PUBSUB_BROKER_ADDRESS; +import static com.linkedin.venice.ConfigKeys.PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS; +import static com.linkedin.venice.ConfigKeys.PUBSUB_SECURITY_PROTOCOL; import static com.linkedin.venice.vpj.VenicePushJobConstants.ALLOW_REGULAR_PUSH_WITH_TTL_REPUSH; import static com.linkedin.venice.vpj.VenicePushJobConstants.COMPLIANCE_PUSH; import static com.linkedin.venice.vpj.VenicePushJobConstants.KAFKA_INPUT_MAX_RECORDS_PER_MAPPER; @@ -26,6 +30,7 @@ import com.linkedin.venice.meta.StoreInfo; import com.linkedin.venice.meta.Version; import com.linkedin.venice.utils.Time; +import com.linkedin.venice.utils.VeniceProperties; import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -42,6 +47,65 @@ */ public class VenicePushJobRepushTest extends VenicePushJobTestBase { + @Test + public void testBuildRepushDictionaryConsumerPropertiesRetainsAdapterConfigAndOverridesBrokers() { + Properties jobProperties = new Properties(); + jobProperties.setProperty( + PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS, + "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory"); + jobProperties.setProperty("xc.pubsub.broker.url.to.region.name.map", "northguard:ei4"); + jobProperties.setProperty(PUBSUB_BROKER_ADDRESS, "destination-broker"); + jobProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, "legacy-broker"); + + Properties consumerProperties = VenicePushJob.buildRepushDictionaryConsumerProperties( + new VeniceProperties(jobProperties), + new Properties(), + "source-broker"); + + // The configured pub-sub adapter factory and its client-specific configs must survive so the source + // dictionary is read with the same adapter as the rest of the repush instead of the implicit Kafka default. + assertEquals( + consumerProperties.getProperty(PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS), + "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory"); + assertEquals(consumerProperties.getProperty("xc.pubsub.broker.url.to.region.name.map"), "northguard:ei4"); + // Both broker properties are pointed at the repush source broker. + assertEquals(consumerProperties.getProperty(PUBSUB_BROKER_ADDRESS), "source-broker"); + assertEquals(consumerProperties.getProperty(KAFKA_BOOTSTRAP_SERVERS), "source-broker"); + } + + @Test + public void testBuildRepushDictionaryConsumerPropertiesAppliesSslOverrides() { + Properties jobProperties = new Properties(); + jobProperties.setProperty(PUBSUB_SECURITY_PROTOCOL, "PLAINTEXT"); + jobProperties.setProperty("ssl.keystore.location", "stale-keystore"); + + Properties sslProperties = new Properties(); + sslProperties.setProperty(PUBSUB_SECURITY_PROTOCOL, "SSL"); + sslProperties.setProperty("ssl.keystore.location", "credential-keystore"); + + Properties consumerProperties = VenicePushJob + .buildRepushDictionaryConsumerProperties(new VeniceProperties(jobProperties), sslProperties, "source-broker"); + + assertEquals(consumerProperties.getProperty(PUBSUB_SECURITY_PROTOCOL), "SSL"); + assertEquals(consumerProperties.getProperty("ssl.keystore.location"), "credential-keystore"); + } + + @Test + public void testBuildRepushDictionaryConsumerPropertiesDoesNotMutateJobProperties() { + Properties jobProperties = new Properties(); + jobProperties.setProperty(PUBSUB_BROKER_ADDRESS, "destination-broker"); + + VenicePushJob.buildRepushDictionaryConsumerProperties( + new VeniceProperties(jobProperties), + new Properties(), + "source-broker"); + + assertEquals( + jobProperties.getProperty(PUBSUB_BROKER_ADDRESS), + "destination-broker", + "Building consumer properties must not mutate the job properties"); + } + @Test(expectedExceptions = VeniceException.class, expectedExceptionsMessageRegExp = ".*Repush with TTL is only supported while using Kafka Input Format.*") public void testRepushTTLJobWithNonKafkaInput() { Properties repushProps = new Properties(); diff --git a/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/input/kafka/TestKafkaInputDictTrainer.java b/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/input/kafka/TestKafkaInputDictTrainer.java index 30d8242af8d..6c6fcc12135 100644 --- a/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/input/kafka/TestKafkaInputDictTrainer.java +++ b/clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/input/kafka/TestKafkaInputDictTrainer.java @@ -60,7 +60,7 @@ private KafkaInputDictTrainer.Param getParam(int sampleSize, CompressionStrategy .setKeySchema("\"string\"") .setCompressionDictSize(900 * 1024) .setDictSampleSize(sampleSize) - .setSslProperties(new Properties()) + .setConsumerProperties(new Properties()) .setSourceVersionCompressionStrategy(sourceVersionCompressionStrategy) .setNewKMESchemasFromController(allSchemaStr) .build();