Skip to content

[server] Isolate consumers by store encryption policy #2953

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
minhmo1620 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[server] Isolate consumers by store encryption policy #2953

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 120 additions & 34 deletions ...nci-client/src/main/java/com/linkedin/davinci/kafka/consumer/AggKafkaConsumerService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package com.linkedin.davinci.kafka.consumer;

import static com.linkedin.venice.ConfigKeys.CLUSTER_ENCRYPTION_ENABLED;
import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS;

import com.fasterxml.jackson.core.type.TypeReference;
Expand All @@ -25,6 +26,7 @@
import com.linkedin.venice.utils.SystemTime;
import com.linkedin.venice.utils.Time;
import com.linkedin.venice.utils.Utils;
import com.linkedin.venice.utils.VeniceProperties;
import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
import io.tehuti.metrics.MetricsRepository;
import it.unimi.dsi.fastutil.objects.Object2IntMap;
Expand All @@ -34,6 +36,7 @@
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.Executors;
Expand All @@ -50,7 +53,8 @@

/**
* {@link AggKafkaConsumerService} supports Kafka consumer pool for multiple Kafka clusters from different data centers;
* for each Kafka bootstrap server url, {@link AggKafkaConsumerService} will create one {@link KafkaConsumerService}.
* for each Kafka bootstrap server URL and consumer decryption mode, {@link AggKafkaConsumerService} will create one
* {@link KafkaConsumerService}.
*/
public class AggKafkaConsumerService extends AbstractVeniceService {
private static final Logger LOGGER = LogManager.getLogger(AggKafkaConsumerService.class);
Expand All @@ -68,7 +72,7 @@ public class AggKafkaConsumerService extends AbstractVeniceService {
private final boolean liveConfigBasedKafkaThrottlingEnabled;
private final boolean isKafkaConsumerOffsetCollectionEnabled;
private final KafkaConsumerService.ConsumerAssignmentStrategy sharedConsumerAssignmentStrategy;
private final Map<String, AbstractKafkaConsumerService> kafkaServerToConsumerServiceMap =
private final Map<ConsumerServiceKey, AbstractKafkaConsumerService> kafkaServerToConsumerServiceMap =
new VeniceConcurrentHashMap<>();
private final Map<String, String> kafkaClusterUrlToAliasMap;
private final Object2IntMap<String> kafkaClusterUrlToIdMap;
Expand All @@ -93,6 +97,42 @@ public class AggKafkaConsumerService extends AbstractVeniceService {
new VeniceJsonSerializer<>(new TypeReference<Map<String, Map<String, TopicPartitionIngestionInfo>>>() {
});

static final class ConsumerServiceKey {
private final String resolvedKafkaUrl;
private final boolean decryptionEnabled;

ConsumerServiceKey(String resolvedKafkaUrl, boolean decryptionEnabled) {
this.resolvedKafkaUrl = Objects.requireNonNull(resolvedKafkaUrl);
this.decryptionEnabled = decryptionEnabled;
}

String getResolvedKafkaUrl() {
return resolvedKafkaUrl;
}

@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (!(o instanceof ConsumerServiceKey)) {
return false;
}
ConsumerServiceKey that = (ConsumerServiceKey) o;
return decryptionEnabled == that.decryptionEnabled && resolvedKafkaUrl.equals(that.resolvedKafkaUrl);
}

@Override
public int hashCode() {
return Objects.hash(resolvedKafkaUrl, decryptionEnabled);
}

@Override
public String toString() {
return resolvedKafkaUrl + " (decryptionEnabled=" + decryptionEnabled + ")";
}
}

public AggKafkaConsumerService(
final PubSubPropertiesSupplier pubSubPropertiesSupplier,
final VeniceServerConfig serverConfig,
Expand Down Expand Up @@ -200,7 +240,7 @@ public void stopInner() throws Exception {
protected static Runnable getStuckConsumerDetectionAndRepairRunnable(
Logger logger,
Time time,
Map<String, AbstractKafkaConsumerService> kafkaServerToConsumerServiceMap,
Map<?, AbstractKafkaConsumerService> kafkaServerToConsumerServiceMap,
Map<String, StoreIngestionTask> versionTopicStoreIngestionTaskMapping,
long stuckConsumerRepairThresholdMs,
long nonExistingTopicIngestionTaskKillThresholdMs,
Expand Down Expand Up @@ -306,12 +346,12 @@ protected static Runnable getStuckConsumerDetectionAndRepairRunnable(
private static void reportStaleTopicPartitions(
Logger logger,
Time time,
Map<String, AbstractKafkaConsumerService> kafkaServerToConsumerServiceMap,
Map<?, AbstractKafkaConsumerService> kafkaServerToConsumerServiceMap,
long consumerPollTrackerStaleThresholdMs) {
StringBuilder stringBuilder = new StringBuilder();
long now = time.getMilliseconds();
// Detect and log any subscribed topic partitions that are not polling records
for (Map.Entry<String, AbstractKafkaConsumerService> consumerService: kafkaServerToConsumerServiceMap.entrySet()) {
for (Map.Entry<?, AbstractKafkaConsumerService> consumerService: kafkaServerToConsumerServiceMap.entrySet()) {
Map<PubSubTopicPartition, Long> staleTopicPartitions =
consumerService.getValue().getStaleTopicPartitions(now - consumerPollTrackerStaleThresholdMs);
if (!staleTopicPartitions.isEmpty()) {
Expand Down Expand Up @@ -354,14 +394,18 @@ private static void reportStaleTopicPartitions(
* or null if there isn't any.
*/
AbstractKafkaConsumerService getKafkaConsumerService(final String kafkaURL) {
AbstractKafkaConsumerService consumerService = kafkaServerToConsumerServiceMap.get(kafkaURL);
if (consumerService == null && kafkaClusterUrlResolver != null) {
// The resolver is needed to resolve a special format of kafka URL to the original kafka URL
consumerService = kafkaServerToConsumerServiceMap.get(kafkaClusterUrlResolver.apply(kafkaURL));
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaURL, false);
if (consumerService == null) {
consumerService = getKafkaConsumerService(kafkaURL, true);
}
return consumerService;
}

AbstractKafkaConsumerService getKafkaConsumerService(final String kafkaURL, boolean decryptionEnabled) {
String resolvedKafkaUrl = kafkaClusterUrlResolver == null ? kafkaURL : kafkaClusterUrlResolver.apply(kafkaURL);
return kafkaServerToConsumerServiceMap.get(new ConsumerServiceKey(resolvedKafkaUrl, decryptionEnabled));
}
Comment on lines +404 to +407

/**
* Create a new {@link KafkaConsumerService} given consumerProperties which must contain a value for "bootstrap.servers".
* If a {@link KafkaConsumerService} for the given "bootstrap.servers" (Kafka URL) has already been created, this method
Expand All @@ -377,9 +421,12 @@ public synchronized AbstractKafkaConsumerService createKafkaConsumerService(fina
throw new IllegalArgumentException("Kafka URL must be set in the consumer properties config. Got: " + kafkaUrl);
}
String resolvedKafkaUrl = kafkaClusterUrlResolver == null ? kafkaUrl : kafkaClusterUrlResolver.apply(kafkaUrl);
final AbstractKafkaConsumerService alreadyCreatedConsumerService = getKafkaConsumerService(resolvedKafkaUrl);
boolean decryptionEnabled = new VeniceProperties(consumerProperties).getBoolean(CLUSTER_ENCRYPTION_ENABLED, false);
ConsumerServiceKey consumerServiceKey = new ConsumerServiceKey(resolvedKafkaUrl, decryptionEnabled);
final AbstractKafkaConsumerService alreadyCreatedConsumerService =
kafkaServerToConsumerServiceMap.get(consumerServiceKey);
if (alreadyCreatedConsumerService != null) {
LOGGER.info("KafkaConsumerService has already been created for Kafka cluster with URL: {}", resolvedKafkaUrl);
LOGGER.info("KafkaConsumerService has already been created for {}", consumerServiceKey);
return alreadyCreatedConsumerService;
}

Expand Down Expand Up @@ -411,7 +458,7 @@ public synchronized AbstractKafkaConsumerService createKafkaConsumerService(fina
getCrossTpProcessingPoolForPoolType(poolType));

AbstractKafkaConsumerService consumerService =
kafkaServerToConsumerServiceMap.computeIfAbsent(resolvedKafkaUrl, url -> {
kafkaServerToConsumerServiceMap.computeIfAbsent(consumerServiceKey, key -> {
if (serverConfig
.getConsumerPoolStrategyType() == KafkaConsumerServiceDelegator.ConsumerPoolStrategyType.CURRENT_VERSION_PRIORITIZATION) {
return new KafkaConsumerServiceDelegator(serverConfig, consumerServiceBuilder);
Expand Down Expand Up @@ -452,18 +499,24 @@ public boolean hasConsumerAssignedFor(
final String kafkaURL,
PubSubTopic versionTopic,
PubSubTopicPartition pubSubTopicPartition) {
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaURL);
if (consumerService == null) {
return false;
for (boolean decryptionEnabled: new boolean[] { false, true }) {
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaURL, decryptionEnabled);
if (consumerService != null) {
SharedKafkaConsumer consumer =
consumerService.getConsumerAssignedToVersionTopicPartition(versionTopic, pubSubTopicPartition);
if (consumer != null && consumer.hasSubscription(pubSubTopicPartition)) {
return true;
}
}
}
SharedKafkaConsumer consumer =
consumerService.getConsumerAssignedToVersionTopicPartition(versionTopic, pubSubTopicPartition);
return consumer != null && consumer.hasSubscription(pubSubTopicPartition);
return false;
}

boolean hasConsumerAssignedFor(PubSubTopic versionTopic, PubSubTopicPartition pubSubTopicPartition) {
for (String kafkaUrl: kafkaServerToConsumerServiceMap.keySet()) {
if (hasConsumerAssignedFor(kafkaUrl, versionTopic, pubSubTopicPartition)) {
for (AbstractKafkaConsumerService consumerService: kafkaServerToConsumerServiceMap.values()) {
SharedKafkaConsumer consumer =
consumerService.getConsumerAssignedToVersionTopicPartition(versionTopic, pubSubTopicPartition);
if (consumer != null && consumer.hasSubscription(pubSubTopicPartition)) {
return true;
}
}
Expand Down Expand Up @@ -515,9 +568,28 @@ public ConsumedDataReceiver<List<DefaultPubSubMessage>> subscribeConsumerFor(
PubSubPosition lastOffset,
boolean inclusive) {
PubSubTopic versionTopic = storeIngestionTask.getVersionTopic();
boolean decryptionEnabled = StoreIngestionTask.resolveConsumerEncryptionEnabled(
serverConfig.getClusterProperties(),
metadataRepository.getStore(versionTopic.getStoreName()));
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaURL, decryptionEnabled);
return subscribeConsumerFor(
consumerService,
kafkaURL,
storeIngestionTask,
partitionReplicaIngestionContext,
lastOffset,
inclusive);
}

ConsumedDataReceiver<List<DefaultPubSubMessage>> subscribeConsumerFor(
AbstractKafkaConsumerService consumerService,
final String kafkaURL,
StoreIngestionTask storeIngestionTask,
PartitionReplicaIngestionContext partitionReplicaIngestionContext,
PubSubPosition lastOffset,
boolean inclusive) {
PubSubTopic versionTopic = storeIngestionTask.getVersionTopic();
PubSubTopicPartition pubSubTopicPartition = partitionReplicaIngestionContext.getPubSubTopicPartition();
AbstractKafkaConsumerService consumerService =
getKafkaConsumerService(kafkaClusterUrlResolver == null ? kafkaURL : kafkaClusterUrlResolver.apply(kafkaURL));
if (consumerService == null) {
throw new VeniceException(
"Kafka consumer service must exist for version topic: " + versionTopic + " in Kafka cluster: " + kafkaURL);
Expand Down Expand Up @@ -549,10 +621,14 @@ public long getLatestOffsetBasedOnMetrics(
final String kafkaURL,
PubSubTopic versionTopic,
PubSubTopicPartition pubSubTopicPartition) {
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaURL);
return consumerService == null
? -1
: consumerService.getLatestOffsetBasedOnMetrics(versionTopic, pubSubTopicPartition);
for (boolean decryptionEnabled: new boolean[] { false, true }) {
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaURL, decryptionEnabled);
if (consumerService != null
&& consumerService.getConsumerAssignedToVersionTopicPartition(versionTopic, pubSubTopicPartition) != null) {
return consumerService.getLatestOffsetBasedOnMetrics(versionTopic, pubSubTopicPartition);
}
}
return -1;
}

/**
Expand Down Expand Up @@ -596,18 +672,21 @@ void resumeConsumerFor(PubSubTopic versionTopic, PubSubTopicPartition pubSubTopi
*/
Set<String> getKafkaUrlsFor(PubSubTopic versionTopic) {
Set<String> kafkaUrls = new HashSet<>(kafkaServerToConsumerServiceMap.size());
for (Map.Entry<String, AbstractKafkaConsumerService> entry: kafkaServerToConsumerServiceMap.entrySet()) {
for (Map.Entry<ConsumerServiceKey, AbstractKafkaConsumerService> entry: kafkaServerToConsumerServiceMap
.entrySet()) {
if (entry.getValue().hasAnySubscriptionFor(versionTopic)) {
kafkaUrls.add(entry.getKey());
kafkaUrls.add(entry.getKey().getResolvedKafkaUrl());
}
}
return kafkaUrls;
}

byte[] getIngestionInfoFor(PubSubTopic versionTopic, PubSubTopicPartition pubSubTopicPartition) throws IOException {
Map<String, Map<String, TopicPartitionIngestionInfo>> topicPartitionIngestionContext = new HashMap<>();
for (String kafkaUrl: kafkaServerToConsumerServiceMap.keySet()) {
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaUrl);
for (Map.Entry<ConsumerServiceKey, AbstractKafkaConsumerService> consumerServiceEntry: kafkaServerToConsumerServiceMap
.entrySet()) {
String kafkaUrl = consumerServiceEntry.getKey().getResolvedKafkaUrl();
AbstractKafkaConsumerService consumerService = consumerServiceEntry.getValue();
Map<PubSubTopicPartition, TopicPartitionIngestionInfo> topicPartitionIngestionInfoMap =
consumerService.getIngestionInfoFor(versionTopic, pubSubTopicPartition, false);
for (Map.Entry<PubSubTopicPartition, TopicPartitionIngestionInfo> entry: topicPartitionIngestionInfoMap
Comment on lines +686 to 692
Expand All @@ -629,12 +708,19 @@ public String getIngestionInfoFor(
if (kafkaUrl == null) {
return "kafkaUrl is not found for region: " + regionName;
}
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaUrl);
if (consumerService == null) {
Map<PubSubTopicPartition, TopicPartitionIngestionInfo> topicPartitionIngestionInfoMap = new HashMap<>();
boolean consumerServiceFound = false;
for (boolean decryptionEnabled: new boolean[] { false, true }) {
AbstractKafkaConsumerService consumerService = getKafkaConsumerService(kafkaUrl, decryptionEnabled);
if (consumerService != null) {
consumerServiceFound = true;
topicPartitionIngestionInfoMap
.putAll(consumerService.getIngestionInfoFor(versionTopic, pubSubTopicPartition, true));
}
}
if (!consumerServiceFound) {
return "Kafka consumer service is not found for kafkaUrl: " + kafkaUrl + ", region: " + regionName;
}
Map<PubSubTopicPartition, TopicPartitionIngestionInfo> topicPartitionIngestionInfoMap =
consumerService.getIngestionInfoFor(versionTopic, pubSubTopicPartition, true);
return KafkaConsumerService.convertTopicPartitionIngestionInfoMapToStr(topicPartitionIngestionInfoMap);
}

Expand Down
16 changes: 14 additions & 2 deletions ...da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import static com.linkedin.davinci.kafka.consumer.LeaderFollowerStateType.LEADER;
import static com.linkedin.davinci.kafka.consumer.LeaderFollowerStateType.STANDBY;
import static com.linkedin.davinci.validation.DataIntegrityValidator.DISABLED;
import static com.linkedin.venice.ConfigKeys.CLUSTER_ENCRYPTION_ENABLED;
import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS;
import static com.linkedin.venice.LogMessages.KILLED_JOB_MESSAGE;
import static com.linkedin.venice.kafka.protocol.enums.ControlMessageType.START_OF_SEGMENT;
Expand Down Expand Up @@ -4932,7 +4933,7 @@ void consumerSubscribe(PubSubTopicPartition pubSubTopicPartition, PubSubPosition
}
final boolean consumeRemotely = !Objects.equals(resolvedKafkaURL, localKafkaServer);
// TODO: Move remote KafkaConsumerService creating operations into the aggKafkaConsumerService.
aggKafkaConsumerService
AbstractKafkaConsumerService kafkaConsumerService = aggKafkaConsumerService
.createKafkaConsumerService(createKafkaConsumerProperties(kafkaProps, resolvedKafkaURL, consumeRemotely));
PartitionConsumptionState pcs = pubSubTopicPartition == null
? null
Expand All @@ -4947,6 +4948,7 @@ void consumerSubscribe(PubSubTopicPartition pubSubTopicPartition, PubSubPosition
// localKafkaServer doesn't have suffix but kafkaURL may have suffix,
// and we don't want to pass the resolvedKafkaURL as it will be passed to data receiver for parsing cluster id
aggKafkaConsumerService.subscribeConsumerFor(
kafkaConsumerService,
kafkaURL,
this,
partitionReplicaIngestionContext,
Expand Down Expand Up @@ -5797,7 +5799,8 @@ protected Properties createKafkaConsumerProperties(
Properties localConsumerProps,
String remoteKafkaSourceAddress,
boolean consumeRemotely) {
Properties newConsumerProps = serverConfig.getClusterProperties().getPropertiesCopy();
VeniceProperties clusterProperties = serverConfig.getClusterProperties();
Properties newConsumerProps = clusterProperties.getPropertiesCopy();
newConsumerProps.putAll(localConsumerProps);
newConsumerProps.setProperty(KAFKA_BOOTSTRAP_SERVERS, remoteKafkaSourceAddress);
VeniceProperties customizedConsumerConfigs = consumeRemotely
Expand All @@ -5806,9 +5809,18 @@ protected Properties createKafkaConsumerProperties(
if (!customizedConsumerConfigs.isEmpty()) {
newConsumerProps.putAll(customizedConsumerConfigs.toProperties());
}
newConsumerProps.setProperty(
CLUSTER_ENCRYPTION_ENABLED,
Boolean.toString(resolveConsumerEncryptionEnabled(clusterProperties, storeRepository.getStore(storeName))));
return newConsumerProps;
}

static boolean resolveConsumerEncryptionEnabled(VeniceProperties clusterProperties, Store store) {
return clusterProperties.containsKey(CLUSTER_ENCRYPTION_ENABLED)
? clusterProperties.getBoolean(CLUSTER_ENCRYPTION_ENABLED)
: store != null && store.isEncryptionEnabled();
}

/**
* A function that would apply on a specific partition to check whether the partition is ready to serve.
*/
Expand Down
Loading
Loading