use of org.apache.kafka.streams.processor.internals.assignment.StickyTaskAssignor in project kafka by apache.
the class StreamPartitionAssignor method assign.
/*
* This assigns tasks to consumer clients in the following steps.
*
* 0. check all repartition source topics and use internal topic manager to make sure
* they have been created with the right number of partitions.
*
* 1. using user customized partition grouper to generate tasks along with their
* assigned partitions; also make sure that the task's corresponding changelog topics
* have been created with the right number of partitions.
*
* 2. using TaskAssignor to assign tasks to consumer clients.
* - Assign a task to a client which was running it previously.
* If there is no such client, assign a task to a client which has its valid local state.
* - A client may have more than one stream threads.
* The assignor tries to assign tasks to a client proportionally to the number of threads.
* - We try not to assign the same set of tasks to two different clients
* We do the assignment in one-pass. The result may not satisfy above all.
*
* 3. within each client, tasks are assigned to consumer clients in round-robin manner.
*/
@Override
public Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions) {
// construct the client metadata from the decoded subscription info
Map<UUID, ClientMetadata> clientsMetadata = new HashMap<>();
for (Map.Entry<String, Subscription> entry : subscriptions.entrySet()) {
String consumerId = entry.getKey();
Subscription subscription = entry.getValue();
SubscriptionInfo info = SubscriptionInfo.decode(subscription.userData());
// create the new client metadata if necessary
ClientMetadata clientMetadata = clientsMetadata.get(info.processId);
if (clientMetadata == null) {
clientMetadata = new ClientMetadata(info.userEndPoint);
clientsMetadata.put(info.processId, clientMetadata);
}
// add the consumer to the client
clientMetadata.addConsumer(consumerId, info);
}
log.info("stream-thread [{}] Constructed client metadata {} from the member subscriptions.", streamThread.getName(), clientsMetadata);
// ---------------- Step Zero ---------------- //
// parse the topology to determine the repartition source topics,
// making sure they are created with the number of partitions as
// the maximum of the depending sub-topologies source topics' number of partitions
Map<Integer, TopologyBuilder.TopicsInfo> topicGroups = streamThread.builder.topicGroups();
Map<String, InternalTopicMetadata> repartitionTopicMetadata = new HashMap<>();
for (TopologyBuilder.TopicsInfo topicsInfo : topicGroups.values()) {
for (InternalTopicConfig topic : topicsInfo.repartitionSourceTopics.values()) {
repartitionTopicMetadata.put(topic.name(), new InternalTopicMetadata(topic));
}
}
boolean numPartitionsNeeded;
do {
numPartitionsNeeded = false;
for (TopologyBuilder.TopicsInfo topicsInfo : topicGroups.values()) {
for (String topicName : topicsInfo.repartitionSourceTopics.keySet()) {
int numPartitions = repartitionTopicMetadata.get(topicName).numPartitions;
// try set the number of partitions for this repartition topic if it is not set yet
if (numPartitions == UNKNOWN) {
for (TopologyBuilder.TopicsInfo otherTopicsInfo : topicGroups.values()) {
Set<String> otherSinkTopics = otherTopicsInfo.sinkTopics;
if (otherSinkTopics.contains(topicName)) {
// use the maximum of all its source topic partitions as the number of partitions
for (String sourceTopicName : otherTopicsInfo.sourceTopics) {
Integer numPartitionsCandidate;
// map().join().join(map())
if (repartitionTopicMetadata.containsKey(sourceTopicName)) {
numPartitionsCandidate = repartitionTopicMetadata.get(sourceTopicName).numPartitions;
} else {
numPartitionsCandidate = metadata.partitionCountForTopic(sourceTopicName);
if (numPartitionsCandidate == null) {
repartitionTopicMetadata.get(topicName).numPartitions = NOT_AVAILABLE;
}
}
if (numPartitionsCandidate != null && numPartitionsCandidate > numPartitions) {
numPartitions = numPartitionsCandidate;
}
}
}
}
// another iteration is needed
if (numPartitions == UNKNOWN)
numPartitionsNeeded = true;
else
repartitionTopicMetadata.get(topicName).numPartitions = numPartitions;
}
}
}
} while (numPartitionsNeeded);
// augment the metadata with the newly computed number of partitions for all the
// repartition source topics
Map<TopicPartition, PartitionInfo> allRepartitionTopicPartitions = new HashMap<>();
for (Map.Entry<String, InternalTopicMetadata> entry : repartitionTopicMetadata.entrySet()) {
String topic = entry.getKey();
Integer numPartitions = entry.getValue().numPartitions;
for (int partition = 0; partition < numPartitions; partition++) {
allRepartitionTopicPartitions.put(new TopicPartition(topic, partition), new PartitionInfo(topic, partition, null, new Node[0], new Node[0]));
}
}
// ensure the co-partitioning topics within the group have the same number of partitions,
// and enforce the number of partitions for those repartition topics to be the same if they
// are co-partitioned as well.
ensureCopartitioning(streamThread.builder.copartitionGroups(), repartitionTopicMetadata, metadata);
// make sure the repartition source topics exist with the right number of partitions,
// create these topics if necessary
prepareTopic(repartitionTopicMetadata);
metadataWithInternalTopics = metadata.withPartitions(allRepartitionTopicPartitions);
log.debug("stream-thread [{}] Created repartition topics {} from the parsed topology.", streamThread.getName(), allRepartitionTopicPartitions.values());
// ---------------- Step One ---------------- //
// get the tasks as partition groups from the partition grouper
Set<String> allSourceTopics = new HashSet<>();
Map<Integer, Set<String>> sourceTopicsByGroup = new HashMap<>();
for (Map.Entry<Integer, TopologyBuilder.TopicsInfo> entry : topicGroups.entrySet()) {
allSourceTopics.addAll(entry.getValue().sourceTopics);
sourceTopicsByGroup.put(entry.getKey(), entry.getValue().sourceTopics);
}
Map<TaskId, Set<TopicPartition>> partitionsForTask = streamThread.partitionGrouper.partitionGroups(sourceTopicsByGroup, metadataWithInternalTopics);
// check if all partitions are assigned, and there are no duplicates of partitions in multiple tasks
Set<TopicPartition> allAssignedPartitions = new HashSet<>();
Map<Integer, Set<TaskId>> tasksByTopicGroup = new HashMap<>();
for (Map.Entry<TaskId, Set<TopicPartition>> entry : partitionsForTask.entrySet()) {
Set<TopicPartition> partitions = entry.getValue();
for (TopicPartition partition : partitions) {
if (allAssignedPartitions.contains(partition)) {
log.warn("stream-thread [{}] Partition {} is assigned to more than one tasks: {}", streamThread.getName(), partition, partitionsForTask);
}
}
allAssignedPartitions.addAll(partitions);
TaskId id = entry.getKey();
Set<TaskId> ids = tasksByTopicGroup.get(id.topicGroupId);
if (ids == null) {
ids = new HashSet<>();
tasksByTopicGroup.put(id.topicGroupId, ids);
}
ids.add(id);
}
for (String topic : allSourceTopics) {
List<PartitionInfo> partitionInfoList = metadataWithInternalTopics.partitionsForTopic(topic);
if (!partitionInfoList.isEmpty()) {
for (PartitionInfo partitionInfo : partitionInfoList) {
TopicPartition partition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
if (!allAssignedPartitions.contains(partition)) {
log.warn("stream-thread [{}] Partition {} is not assigned to any tasks: {}", streamThread.getName(), partition, partitionsForTask);
}
}
} else {
log.warn("stream-thread [{}] No partitions found for topic {}", streamThread.getName(), topic);
}
}
// add tasks to state change log topic subscribers
Map<String, InternalTopicMetadata> changelogTopicMetadata = new HashMap<>();
for (Map.Entry<Integer, TopologyBuilder.TopicsInfo> entry : topicGroups.entrySet()) {
final int topicGroupId = entry.getKey();
final Map<String, InternalTopicConfig> stateChangelogTopics = entry.getValue().stateChangelogTopics;
for (InternalTopicConfig topicConfig : stateChangelogTopics.values()) {
// the expected number of partitions is the max value of TaskId.partition + 1
int numPartitions = UNKNOWN;
if (tasksByTopicGroup.get(topicGroupId) != null) {
for (TaskId task : tasksByTopicGroup.get(topicGroupId)) {
if (numPartitions < task.partition + 1)
numPartitions = task.partition + 1;
}
InternalTopicMetadata topicMetadata = new InternalTopicMetadata(topicConfig);
topicMetadata.numPartitions = numPartitions;
changelogTopicMetadata.put(topicConfig.name(), topicMetadata);
} else {
log.debug("stream-thread [{}] No tasks found for topic group {}", streamThread.getName(), topicGroupId);
}
}
}
prepareTopic(changelogTopicMetadata);
log.debug("stream-thread [{}] Created state changelog topics {} from the parsed topology.", streamThread.getName(), changelogTopicMetadata);
// ---------------- Step Two ---------------- //
// assign tasks to clients
Map<UUID, ClientState> states = new HashMap<>();
for (Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
states.put(entry.getKey(), entry.getValue().state);
}
log.debug("stream-thread [{}] Assigning tasks {} to clients {} with number of replicas {}", streamThread.getName(), partitionsForTask.keySet(), states, numStandbyReplicas);
final StickyTaskAssignor<UUID> taskAssignor = new StickyTaskAssignor<>(states, partitionsForTask.keySet());
taskAssignor.assign(numStandbyReplicas);
log.info("stream-thread [{}] Assigned tasks to clients as {}.", streamThread.getName(), states);
// ---------------- Step Three ---------------- //
// construct the global partition assignment per host map
partitionsByHostState = new HashMap<>();
for (Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
HostInfo hostInfo = entry.getValue().hostInfo;
if (hostInfo != null) {
final Set<TopicPartition> topicPartitions = new HashSet<>();
final ClientState state = entry.getValue().state;
for (final TaskId id : state.activeTasks()) {
topicPartitions.addAll(partitionsForTask.get(id));
}
partitionsByHostState.put(hostInfo, topicPartitions);
}
}
// within the client, distribute tasks to its owned consumers
Map<String, Assignment> assignment = new HashMap<>();
for (Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
final Set<String> consumers = entry.getValue().consumers;
final ClientState state = entry.getValue().state;
final ArrayList<TaskId> taskIds = new ArrayList<>(state.assignedTaskCount());
final int numActiveTasks = state.activeTaskCount();
taskIds.addAll(state.activeTasks());
taskIds.addAll(state.standbyTasks());
final int numConsumers = consumers.size();
int i = 0;
for (String consumer : consumers) {
Map<TaskId, Set<TopicPartition>> standby = new HashMap<>();
ArrayList<AssignedPartition> assignedPartitions = new ArrayList<>();
final int numTaskIds = taskIds.size();
for (int j = i; j < numTaskIds; j += numConsumers) {
TaskId taskId = taskIds.get(j);
if (j < numActiveTasks) {
for (TopicPartition partition : partitionsForTask.get(taskId)) {
assignedPartitions.add(new AssignedPartition(taskId, partition));
}
} else {
Set<TopicPartition> standbyPartitions = standby.get(taskId);
if (standbyPartitions == null) {
standbyPartitions = new HashSet<>();
standby.put(taskId, standbyPartitions);
}
standbyPartitions.addAll(partitionsForTask.get(taskId));
}
}
Collections.sort(assignedPartitions);
List<TaskId> active = new ArrayList<>();
List<TopicPartition> activePartitions = new ArrayList<>();
for (AssignedPartition partition : assignedPartitions) {
active.add(partition.taskId);
activePartitions.add(partition.partition);
}
// finally, encode the assignment before sending back to coordinator
assignment.put(consumer, new Assignment(activePartitions, new AssignmentInfo(active, standby, partitionsByHostState).encode()));
i++;
}
}
return assignment;
}
use of org.apache.kafka.streams.processor.internals.assignment.StickyTaskAssignor in project apache-kafka-on-k8s by banzaicloud.
the class StreamsPartitionAssignor method assign.
/*
* This assigns tasks to consumer clients in the following steps.
*
* 0. check all repartition source topics and use internal topic manager to make sure
* they have been created with the right number of partitions.
*
* 1. using user customized partition grouper to generate tasks along with their
* assigned partitions; also make sure that the task's corresponding changelog topics
* have been created with the right number of partitions.
*
* 2. using TaskAssignor to assign tasks to consumer clients.
* - Assign a task to a client which was running it previously.
* If there is no such client, assign a task to a client which has its valid local state.
* - A client may have more than one stream threads.
* The assignor tries to assign tasks to a client proportionally to the number of threads.
* - We try not to assign the same set of tasks to two different clients
* We do the assignment in one-pass. The result may not satisfy above all.
*
* 3. within each client, tasks are assigned to consumer clients in round-robin manner.
*/
@Override
public Map<String, Assignment> assign(final Cluster metadata, final Map<String, Subscription> subscriptions) {
// construct the client metadata from the decoded subscription info
final Map<UUID, ClientMetadata> clientsMetadata = new HashMap<>();
int minUserMetadataVersion = SubscriptionInfo.LATEST_SUPPORTED_VERSION;
for (final Map.Entry<String, Subscription> entry : subscriptions.entrySet()) {
final String consumerId = entry.getKey();
final Subscription subscription = entry.getValue();
final SubscriptionInfo info = SubscriptionInfo.decode(subscription.userData());
final int usedVersion = info.version();
if (usedVersion > SubscriptionInfo.LATEST_SUPPORTED_VERSION) {
throw new IllegalStateException("Unknown metadata version: " + usedVersion + "; latest supported version: " + SubscriptionInfo.LATEST_SUPPORTED_VERSION);
}
if (usedVersion < minUserMetadataVersion) {
minUserMetadataVersion = usedVersion;
}
// create the new client metadata if necessary
ClientMetadata clientMetadata = clientsMetadata.get(info.processId());
if (clientMetadata == null) {
clientMetadata = new ClientMetadata(info.userEndPoint());
clientsMetadata.put(info.processId(), clientMetadata);
}
// add the consumer to the client
clientMetadata.addConsumer(consumerId, info);
}
log.debug("Constructed client metadata {} from the member subscriptions.", clientsMetadata);
// ---------------- Step Zero ---------------- //
// parse the topology to determine the repartition source topics,
// making sure they are created with the number of partitions as
// the maximum of the depending sub-topologies source topics' number of partitions
final Map<Integer, InternalTopologyBuilder.TopicsInfo> topicGroups = taskManager.builder().topicGroups();
final Map<String, InternalTopicMetadata> repartitionTopicMetadata = new HashMap<>();
for (final InternalTopologyBuilder.TopicsInfo topicsInfo : topicGroups.values()) {
for (final InternalTopicConfig topic : topicsInfo.repartitionSourceTopics.values()) {
repartitionTopicMetadata.put(topic.name(), new InternalTopicMetadata(topic));
}
}
boolean numPartitionsNeeded;
do {
numPartitionsNeeded = false;
for (final InternalTopologyBuilder.TopicsInfo topicsInfo : topicGroups.values()) {
for (final String topicName : topicsInfo.repartitionSourceTopics.keySet()) {
int numPartitions = repartitionTopicMetadata.get(topicName).numPartitions;
// try set the number of partitions for this repartition topic if it is not set yet
if (numPartitions == UNKNOWN) {
for (final InternalTopologyBuilder.TopicsInfo otherTopicsInfo : topicGroups.values()) {
final Set<String> otherSinkTopics = otherTopicsInfo.sinkTopics;
if (otherSinkTopics.contains(topicName)) {
// use the maximum of all its source topic partitions as the number of partitions
for (final String sourceTopicName : otherTopicsInfo.sourceTopics) {
final Integer numPartitionsCandidate;
// map().join().join(map())
if (repartitionTopicMetadata.containsKey(sourceTopicName)) {
numPartitionsCandidate = repartitionTopicMetadata.get(sourceTopicName).numPartitions;
} else {
numPartitionsCandidate = metadata.partitionCountForTopic(sourceTopicName);
if (numPartitionsCandidate == null) {
repartitionTopicMetadata.get(topicName).numPartitions = NOT_AVAILABLE;
}
}
if (numPartitionsCandidate != null && numPartitionsCandidate > numPartitions) {
numPartitions = numPartitionsCandidate;
}
}
}
}
// another iteration is needed
if (numPartitions == UNKNOWN) {
numPartitionsNeeded = true;
} else {
repartitionTopicMetadata.get(topicName).numPartitions = numPartitions;
}
}
}
}
} while (numPartitionsNeeded);
// ensure the co-partitioning topics within the group have the same number of partitions,
// and enforce the number of partitions for those repartition topics to be the same if they
// are co-partitioned as well.
ensureCopartitioning(taskManager.builder().copartitionGroups(), repartitionTopicMetadata, metadata);
// make sure the repartition source topics exist with the right number of partitions,
// create these topics if necessary
prepareTopic(repartitionTopicMetadata);
// augment the metadata with the newly computed number of partitions for all the
// repartition source topics
final Map<TopicPartition, PartitionInfo> allRepartitionTopicPartitions = new HashMap<>();
for (final Map.Entry<String, InternalTopicMetadata> entry : repartitionTopicMetadata.entrySet()) {
final String topic = entry.getKey();
final int numPartitions = entry.getValue().numPartitions;
for (int partition = 0; partition < numPartitions; partition++) {
allRepartitionTopicPartitions.put(new TopicPartition(topic, partition), new PartitionInfo(topic, partition, null, new Node[0], new Node[0]));
}
}
final Cluster fullMetadata = metadata.withPartitions(allRepartitionTopicPartitions);
taskManager.setClusterMetadata(fullMetadata);
log.debug("Created repartition topics {} from the parsed topology.", allRepartitionTopicPartitions.values());
// ---------------- Step One ---------------- //
// get the tasks as partition groups from the partition grouper
final Set<String> allSourceTopics = new HashSet<>();
final Map<Integer, Set<String>> sourceTopicsByGroup = new HashMap<>();
for (final Map.Entry<Integer, InternalTopologyBuilder.TopicsInfo> entry : topicGroups.entrySet()) {
allSourceTopics.addAll(entry.getValue().sourceTopics);
sourceTopicsByGroup.put(entry.getKey(), entry.getValue().sourceTopics);
}
final Map<TaskId, Set<TopicPartition>> partitionsForTask = partitionGrouper.partitionGroups(sourceTopicsByGroup, fullMetadata);
// check if all partitions are assigned, and there are no duplicates of partitions in multiple tasks
final Set<TopicPartition> allAssignedPartitions = new HashSet<>();
final Map<Integer, Set<TaskId>> tasksByTopicGroup = new HashMap<>();
for (final Map.Entry<TaskId, Set<TopicPartition>> entry : partitionsForTask.entrySet()) {
final Set<TopicPartition> partitions = entry.getValue();
for (final TopicPartition partition : partitions) {
if (allAssignedPartitions.contains(partition)) {
log.warn("Partition {} is assigned to more than one tasks: {}", partition, partitionsForTask);
}
}
allAssignedPartitions.addAll(partitions);
final TaskId id = entry.getKey();
Set<TaskId> ids = tasksByTopicGroup.get(id.topicGroupId);
if (ids == null) {
ids = new HashSet<>();
tasksByTopicGroup.put(id.topicGroupId, ids);
}
ids.add(id);
}
for (final String topic : allSourceTopics) {
final List<PartitionInfo> partitionInfoList = fullMetadata.partitionsForTopic(topic);
if (!partitionInfoList.isEmpty()) {
for (final PartitionInfo partitionInfo : partitionInfoList) {
final TopicPartition partition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
if (!allAssignedPartitions.contains(partition)) {
log.warn("Partition {} is not assigned to any tasks: {}", partition, partitionsForTask);
}
}
} else {
log.warn("No partitions found for topic {}", topic);
}
}
// add tasks to state change log topic subscribers
final Map<String, InternalTopicMetadata> changelogTopicMetadata = new HashMap<>();
for (final Map.Entry<Integer, InternalTopologyBuilder.TopicsInfo> entry : topicGroups.entrySet()) {
final int topicGroupId = entry.getKey();
final Map<String, InternalTopicConfig> stateChangelogTopics = entry.getValue().stateChangelogTopics;
for (final InternalTopicConfig topicConfig : stateChangelogTopics.values()) {
// the expected number of partitions is the max value of TaskId.partition + 1
int numPartitions = UNKNOWN;
if (tasksByTopicGroup.get(topicGroupId) != null) {
for (final TaskId task : tasksByTopicGroup.get(topicGroupId)) {
if (numPartitions < task.partition + 1)
numPartitions = task.partition + 1;
}
final InternalTopicMetadata topicMetadata = new InternalTopicMetadata(topicConfig);
topicMetadata.numPartitions = numPartitions;
changelogTopicMetadata.put(topicConfig.name(), topicMetadata);
} else {
log.debug("No tasks found for topic group {}", topicGroupId);
}
}
}
prepareTopic(changelogTopicMetadata);
log.debug("Created state changelog topics {} from the parsed topology.", changelogTopicMetadata.values());
// ---------------- Step Two ---------------- //
// assign tasks to clients
final Map<UUID, ClientState> states = new HashMap<>();
for (final Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
states.put(entry.getKey(), entry.getValue().state);
}
log.debug("Assigning tasks {} to clients {} with number of replicas {}", partitionsForTask.keySet(), states, numStandbyReplicas);
final StickyTaskAssignor<UUID> taskAssignor = new StickyTaskAssignor<>(states, partitionsForTask.keySet());
taskAssignor.assign(numStandbyReplicas);
log.info("Assigned tasks to clients as {}.", states);
// ---------------- Step Three ---------------- //
// construct the global partition assignment per host map
final Map<HostInfo, Set<TopicPartition>> partitionsByHostState = new HashMap<>();
if (minUserMetadataVersion == 2) {
for (final Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
final HostInfo hostInfo = entry.getValue().hostInfo;
if (hostInfo != null) {
final Set<TopicPartition> topicPartitions = new HashSet<>();
final ClientState state = entry.getValue().state;
for (final TaskId id : state.activeTasks()) {
topicPartitions.addAll(partitionsForTask.get(id));
}
partitionsByHostState.put(hostInfo, topicPartitions);
}
}
}
taskManager.setPartitionsByHostState(partitionsByHostState);
// within the client, distribute tasks to its owned consumers
final Map<String, Assignment> assignment = new HashMap<>();
for (final Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
final Set<String> consumers = entry.getValue().consumers;
final ClientState state = entry.getValue().state;
final List<List<TaskId>> interleavedActive = interleaveTasksByGroupId(state.activeTasks(), consumers.size());
final List<List<TaskId>> interleavedStandby = interleaveTasksByGroupId(state.standbyTasks(), consumers.size());
int consumerTaskIndex = 0;
for (final String consumer : consumers) {
final Map<TaskId, Set<TopicPartition>> standby = new HashMap<>();
final ArrayList<AssignedPartition> assignedPartitions = new ArrayList<>();
final List<TaskId> assignedActiveList = interleavedActive.get(consumerTaskIndex);
for (final TaskId taskId : assignedActiveList) {
for (final TopicPartition partition : partitionsForTask.get(taskId)) {
assignedPartitions.add(new AssignedPartition(taskId, partition));
}
}
if (!state.standbyTasks().isEmpty()) {
final List<TaskId> assignedStandbyList = interleavedStandby.get(consumerTaskIndex);
for (final TaskId taskId : assignedStandbyList) {
Set<TopicPartition> standbyPartitions = standby.get(taskId);
if (standbyPartitions == null) {
standbyPartitions = new HashSet<>();
standby.put(taskId, standbyPartitions);
}
standbyPartitions.addAll(partitionsForTask.get(taskId));
}
}
consumerTaskIndex++;
Collections.sort(assignedPartitions);
final List<TaskId> active = new ArrayList<>();
final List<TopicPartition> activePartitions = new ArrayList<>();
for (final AssignedPartition partition : assignedPartitions) {
active.add(partition.taskId);
activePartitions.add(partition.partition);
}
// finally, encode the assignment before sending back to coordinator
assignment.put(consumer, new Assignment(activePartitions, new AssignmentInfo(minUserMetadataVersion, active, standby, partitionsByHostState).encode()));
}
}
return assignment;
}
Aggregations