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..702b75c976d 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,8 @@ 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)); + .readDictionaryFromKafka(pushJobSetting.kafkaInputTopic, getSourceDictionaryConsumerProperties()); if (sourceDict != null) { pushJobSetting.sourceDictionary = ByteUtils.extractByteArray(sourceDict); } @@ -1666,12 +1662,35 @@ private Optional getCompressionDictionary() throws VeniceException { return Optional.of(emptyPushZstdDictionary.get()); } + private VeniceProperties getSourceDictionaryConsumerProperties() { + return getSourceDictionaryConsumerProperties(pushJobSetting.repushSourcePubsubBroker); + } + + @VisibleForTesting + VeniceProperties getSourceDictionaryConsumerProperties(String sourcePubsubBroker) { + return buildSourceDictionaryConsumerProperties( + props, + pushJobSetting.enableSSL ? sslProperties.get() : new Properties(), + sourcePubsubBroker); + } + + @VisibleForTesting + static VeniceProperties buildSourceDictionaryConsumerProperties( + VeniceProperties jobProperties, + Properties sslProperties, + String sourcePubsubBroker) { + Properties consumerProperties = jobProperties.toProperties(); + consumerProperties.putAll(sslProperties); + consumerProperties.setProperty(PUBSUB_BROKER_ADDRESS, sourcePubsubBroker); + consumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, sourcePubsubBroker); + return new VeniceProperties(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()) .setCompressionDictSize( props.getInt( COMPRESSION_DICTIONARY_SIZE_LIMIT, @@ -1686,19 +1705,14 @@ private ByteBuffer fetchOrBuildCompressionDictionary() throws VeniceException { LOGGER.info("Rebuild a new Zstd dictionary from the input topic: {}", pushJobSetting.kafkaInputTopic); paramBuilder.setKafkaInputBroker(pushJobSetting.repushSourcePubsubBroker) .setTopicName(pushJobSetting.kafkaInputTopic) + .setConsumerProperties(getSourceDictionaryConsumerProperties().toProperties()) .setSourceVersionCompressionStrategy(pushJobSetting.sourceKafkaInputVersionInfo.getCompressionStrategy()); KafkaInputDictTrainer dictTrainer = new KafkaInputDictTrainer(paramBuilder.build()); 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)); + .readDictionaryFromKafka(pushJobSetting.kafkaInputTopic, getSourceDictionaryConsumerProperties()); } } LOGGER.info( @@ -1740,8 +1754,9 @@ private ByteBuffer fetchOrBuildCompressionDictionary() throws VeniceException { "Rebuild a new Zstd dictionary from the source topic: {} in Kafka: {}", sourceTopicName, sourceKafkaUrl); - paramBuilder.setKafkaInputBroker(repushInfoResponse.getRepushInfo().getKafkaBrokerUrl()) + paramBuilder.setKafkaInputBroker(sourceKafkaUrl) .setTopicName(sourceTopicName) + .setConsumerProperties(getSourceDictionaryConsumerProperties(sourceKafkaUrl).toProperties()) .setSourceVersionCompressionStrategy( repushInfoResponse.getRepushInfo().getVersion().getCompressionStrategy()); KafkaInputDictTrainer dictTrainer = new KafkaInputDictTrainer(paramBuilder.build()); 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..c8ae6fb05bf 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,8 @@ public ParamBuilder setKeySchema(String keySchema) { return this; } - public ParamBuilder setSslProperties(Properties sslProperties) { - this.sslProperties = sslProperties; + public ParamBuilder setConsumerProperties(Properties consumerProperties) { + this.consumerProperties = consumerProperties; return this; } @@ -166,11 +166,11 @@ protected KafkaInputDictTrainer( this.trainerSupplier = trainerSupplier; this.sourceVersionCompressionStrategy = param.sourceVersionCompressionStrategy; Properties properties = new Properties(); + properties.putAll(param.consumerProperties); properties.setProperty(VENICE_REPUSH_SOURCE_PUBSUB_BROKER, param.kafkaInputBroker); 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.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..d0194d750b6 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,10 @@ 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.ConfigKeys.PUBSUB_TYPE_ID_TO_POSITION_CLASS_NAME_MAP; 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 +31,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 +48,57 @@ */ public class VenicePushJobRepushTest extends VenicePushJobTestBase { + @Test + public void testSourceDictionaryConsumerPropertiesRetainPubSubConfigAndOverrideBrokers() { + Properties jobProperties = new Properties(); + jobProperties.setProperty( + PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS, + "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory"); + jobProperties.setProperty( + PUBSUB_TYPE_ID_TO_POSITION_CLASS_NAME_MAP, + "1:com.linkedin.venice.pubsub.adapter.xinfra.XinfraPosition"); + 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"); + + try (VenicePushJob pushJob = getSpyVenicePushJob(jobProperties, null)) { + VeniceProperties consumerProperties = pushJob.getSourceDictionaryConsumerProperties("source-broker"); + + assertEquals( + consumerProperties.getString(PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS), + "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory"); + assertEquals( + consumerProperties.getString(PUBSUB_TYPE_ID_TO_POSITION_CLASS_NAME_MAP), + "1:com.linkedin.venice.pubsub.adapter.xinfra.XinfraPosition"); + assertEquals(consumerProperties.getString("xc.pubsub.broker.url.to.region.name.map"), "northguard:ei4"); + assertEquals(consumerProperties.getString(PUBSUB_BROKER_ADDRESS), "source-broker"); + assertEquals(consumerProperties.getString(KAFKA_BOOTSTRAP_SERVERS), "source-broker"); + assertEquals( + jobProperties.getProperty(PUBSUB_BROKER_ADDRESS), + "destination-broker", + "Building consumer properties must not mutate the job properties"); + } + } + + @Test + public void testSourceDictionaryConsumerPropertiesApplySslOverrides() { + Properties jobProperties = new Properties(); + jobProperties.setProperty(PUBSUB_SECURITY_PROTOCOL, "PLAINTEXT"); + jobProperties.setProperty("ssl.keystore.location", "stale-keystore"); + jobProperties.setProperty("xc.tls.key.store.type", "PKCS12"); + + Properties sslProperties = new Properties(); + sslProperties.setProperty(PUBSUB_SECURITY_PROTOCOL, "SSL"); + sslProperties.setProperty("ssl.keystore.location", "credential-keystore"); + + VeniceProperties consumerProperties = VenicePushJob + .buildSourceDictionaryConsumerProperties(new VeniceProperties(jobProperties), sslProperties, "source-broker"); + + assertEquals(consumerProperties.getString(PUBSUB_SECURITY_PROTOCOL), "SSL"); + assertEquals(consumerProperties.getString("ssl.keystore.location"), "credential-keystore"); + assertEquals(consumerProperties.getString("xc.tls.key.store.type"), "PKCS12"); + } + @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..17e10b58d11 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 @@ -1,5 +1,11 @@ package com.linkedin.venice.hadoop.input.kafka; +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.ConfigKeys.PUBSUB_TYPE_ID_TO_POSITION_CLASS_NAME_MAP; +import static com.linkedin.venice.vpj.VenicePushJobConstants.KAFKA_INPUT_TOPIC; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; @@ -7,6 +13,10 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.fail; import com.github.luben.zstd.ZstdDictTrainer; import com.linkedin.venice.compression.CompressionStrategy; @@ -37,6 +47,7 @@ import org.apache.avro.Schema; import org.apache.hadoop.mapred.InputSplit; import org.apache.hadoop.mapred.RecordReader; +import org.mockito.ArgumentCaptor; import org.testng.annotations.Test; @@ -52,6 +63,13 @@ private KafkaInputDictTrainer.Param getParam(int sampleSize) { } private KafkaInputDictTrainer.Param getParam(int sampleSize, CompressionStrategy sourceVersionCompressionStrategy) { + return getParam(sampleSize, sourceVersionCompressionStrategy, new Properties()); + } + + private KafkaInputDictTrainer.Param getParam( + int sampleSize, + CompressionStrategy sourceVersionCompressionStrategy, + Properties consumerProperties) { Map allSchemas = Utils.getAllSchemasFromResources(AvroProtocolDefinition.KAFKA_MESSAGE_ENVELOPE); Map allSchemaStr = allSchemas.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().toString())); @@ -60,14 +78,14 @@ private KafkaInputDictTrainer.Param getParam(int sampleSize, CompressionStrategy .setKeySchema("\"string\"") .setCompressionDictSize(900 * 1024) .setDictSampleSize(sampleSize) - .setSslProperties(new Properties()) + .setConsumerProperties(consumerProperties) .setSourceVersionCompressionStrategy(sourceVersionCompressionStrategy) .setNewKMESchemasFromController(allSchemaStr) .build(); } - @Test(expectedExceptions = VeniceException.class, expectedExceptionsMessageRegExp = "No record.*") - public void testEmptyTopic() throws IOException { + @Test + public void testConsumerPropertiesPropagateToEmptyTopicTraining() throws IOException { KafkaInputFormat mockFormat = mock(KafkaInputFormat.class); PubSubTopicPartition topicPartition = new PubSubTopicPartitionImpl(PUB_SUB_TOPIC_REPOSITORY.getTopic("test_topic"), 0); @@ -79,12 +97,48 @@ public void testEmptyTopic() throws IOException { doReturn(false).when(mockRecordReader).next(any(), any()); doReturn(mockRecordReader).when(mockFormat).getRecordReader(any(), any(), any(), any()); + Properties consumerProperties = new Properties(); + consumerProperties.setProperty( + PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS, + "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory"); + consumerProperties.setProperty( + PUBSUB_TYPE_ID_TO_POSITION_CLASS_NAME_MAP, + "1:com.linkedin.venice.pubsub.adapter.xinfra.XinfraPosition"); + consumerProperties.setProperty("xc.pubsub.broker.url.to.region.name.map", "northguard:ei4"); + consumerProperties.setProperty(PUBSUB_BROKER_ADDRESS, "test_url"); + consumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, "test_url"); + consumerProperties.setProperty(PUBSUB_SECURITY_PROTOCOL, "SSL"); + consumerProperties.setProperty("ssl.keystore.location", "credential-keystore"); + KafkaInputDictTrainer trainer = new KafkaInputDictTrainer( mockFormat, Optional.empty(), - getParam(100), + getParam(100, CompressionStrategy.NO_OP, consumerProperties), getCompressorBuilder(new NoopCompressor())); - trainer.trainDict(Optional.of(mock(PubSubConsumerAdapter.class))); + try { + trainer.trainDict(Optional.of(mock(PubSubConsumerAdapter.class))); + fail("Expected training on an empty topic to fail"); + } catch (VeniceException e) { + assertTrue(e.getMessage().startsWith("No record")); + } + + ArgumentCaptor consumerPropertiesCaptor = ArgumentCaptor.forClass(VeniceProperties.class); + verify(mockFormat).getSplits(consumerPropertiesCaptor.capture()); + VeniceProperties actualConsumerProperties = consumerPropertiesCaptor.getValue(); + assertEquals( + actualConsumerProperties.getString(PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS), + "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory"); + assertEquals( + actualConsumerProperties.getString(PUBSUB_TYPE_ID_TO_POSITION_CLASS_NAME_MAP), + "1:com.linkedin.venice.pubsub.adapter.xinfra.XinfraPosition"); + assertEquals(actualConsumerProperties.getString("xc.pubsub.broker.url.to.region.name.map"), "northguard:ei4"); + assertEquals(actualConsumerProperties.getString(PUBSUB_BROKER_ADDRESS), "test_url"); + assertEquals(actualConsumerProperties.getString(KAFKA_BOOTSTRAP_SERVERS), "test_url"); + assertEquals(actualConsumerProperties.getString(PUBSUB_SECURITY_PROTOCOL), "SSL"); + assertEquals(actualConsumerProperties.getString("ssl.keystore.location"), "credential-keystore"); + assertFalse( + consumerProperties.containsKey(KAFKA_INPUT_TOPIC), + "Building trainer properties must not mutate the supplied consumer properties"); } interface ResettableRecordReader extends RecordReader {