Skip to content

KAFKA-15045: (KIP-924 pt. 15) Implement #defaultStandbyTaskAssignment and finish rack-aware standby optimization #16129

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’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ableegoldman merged 6 commits into from
May 30, 2024
Merged

KAFKA-15045: (KIP-924 pt. 15) Implement #defaultStandbyTaskAssignment and finish rack-aware standby optimization #16129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f51153c
KAFKA-15045: (KIP-924 pt. 15) More TaskAssigmentUtils implementation
apourchet May 29, 2024
56c990c
finished rough draft
apourchet May 29, 2024
87dad64
formatting
apourchet May 29, 2024
597ae87
addressed bug comments
apourchet May 30, 2024
8eacaca
first pass to address most comments
apourchet May 30, 2024
01d5f4d
remove odd rack hardcode case
apourchet May 30, 2024
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
51 changes: 34 additions & 17 deletions ...s/src/main/java/org/apache/kafka/streams/processor/assignment/KafkaStreamsAssignment.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,13 @@
*/
package org.apache.kafka.streams.processor.assignment;

import static java.util.Collections.unmodifiableMap;

import java.time.Instant;
import java.util.HashSet;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.apache.kafka.streams.processor.TaskId;

Expand All @@ -31,7 +33,7 @@
public class KafkaStreamsAssignment {

private final ProcessId processId;
private final Map<TaskId, AssignedTask> assignment;
private final Map<TaskId, AssignedTask> tasks;
private final Optional<Instant> followupRebalanceDeadline;

/**
Expand All @@ -45,7 +47,8 @@ public class KafkaStreamsAssignment {
* @return a new KafkaStreamsAssignment object with the given processId and assignment
*/
public static KafkaStreamsAssignment of(final ProcessId processId, final Set<AssignedTask> assignment) {
return new KafkaStreamsAssignment(processId, assignment, Optional.empty());
final Map<TaskId, AssignedTask> tasks = assignment.stream().collect(Collectors.toMap(AssignedTask::id, Function.identity()));
return new KafkaStreamsAssignment(processId, tasks, Optional.empty());
}

/**
Expand All @@ -62,14 +65,14 @@ public static KafkaStreamsAssignment of(final ProcessId processId, final Set<Ass
* @return a new KafkaStreamsAssignment object with the same processId and assignment but with the given rebalanceDeadline
*/
public KafkaStreamsAssignment withFollowupRebalance(final Instant rebalanceDeadline) {
return new KafkaStreamsAssignment(this.processId(), this.assignment(), Optional.of(rebalanceDeadline));
return new KafkaStreamsAssignment(this.processId(), this.tasks(), Optional.of(rebalanceDeadline));
}

private KafkaStreamsAssignment(final ProcessId processId,
final Set<AssignedTask> assignment,
final Map<TaskId, AssignedTask> tasks,
final Optional<Instant> followupRebalanceDeadline) {
this.processId = processId;
this.assignment = assignment.stream().collect(Collectors.toMap(AssignedTask::id, t -> t));
this.tasks = tasks;
this.followupRebalanceDeadline = followupRebalanceDeadline;
}

Expand All @@ -83,24 +86,18 @@ public ProcessId processId() {

/**
*
* @return a set of assigned tasks that are part of this {@code KafkaStreamsAssignment}
* @return a read-only set of assigned tasks that are part of this {@code KafkaStreamsAssignment}
*/
public Set<AssignedTask> assignment() {
// TODO change assignment to return a map so we aren't forced to copy this into a Set
return new HashSet<>(assignment.values());
}

// TODO: merge this with #assignment by having it return a Map<TaskId, AssignedTask>
public Set<TaskId> assignedTaskIds() {
return assignment.keySet();
public Map<TaskId, AssignedTask> tasks() {
return unmodifiableMap(tasks);
}

public void assignTask(final AssignedTask newTask) {
assignment.put(newTask.id(), newTask);
tasks.put(newTask.id(), newTask);
}

public void removeTask(final AssignedTask removedTask) {
assignment.remove(removedTask.id());
tasks.remove(removedTask.id());
}

/**
Expand Down Expand Up @@ -140,5 +137,25 @@ public TaskId id() {
public Type type() {
return taskType;
}

@Override
public int hashCode() {
final int prime = 31;
int result = prime + this.id.hashCode();
result = prime * result + this.type().hashCode();
return result;
}

@Override
public boolean equals(final Object obj) {
if (this == obj)
return true;
if (obj == null)
return false;
if (getClass() != obj.getClass())
return false;
final AssignedTask other = (AssignedTask) obj;
return this.id.equals(other.id()) && this.taskType == other.taskType;
}
}
}
481 changes: 461 additions & 20 deletions streams/src/main/java/org/apache/kafka/streams/processor/assignment/TaskAssignmentUtils.java
View file
Open in desktop

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion ...main/java/org/apache/kafka/streams/processor/assignment/assignors/StickyTaskAssignor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,7 @@ public void processOptimizedAssignments(final Map<ProcessId, KafkaStreamsAssignm

for (final Map.Entry<ProcessId, KafkaStreamsAssignment> entry : optimizedAssignments.entrySet()) {
final ProcessId processId = entry.getKey();
final Set<AssignedTask> assignedTasks = optimizedAssignments.get(processId).assignment();
final Set<AssignedTask> assignedTasks = new HashSet<>(optimizedAssignments.get(processId).tasks().values());
Comment thread
ableegoldman marked this conversation as resolved.
newAssignments.put(processId, assignedTasks);

for (final AssignedTask task : assignedTasks) {
Expand Down
53 changes: 28 additions & 25 deletions .../src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import java.time.Instant;
import java.util.Optional;
import java.util.function.Function;
import org.apache.kafka.clients.admin.Admin;
import org.apache.kafka.clients.admin.ListOffsetsResult;
import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
Expand Down Expand Up @@ -555,18 +556,21 @@ private ApplicationState buildApplicationState(final TopologyMetadata topologyMe

RackUtils.annotateTopicPartitionsWithRackInfo(cluster, internalTopicManager, allTopicPartitions);

final Set<TaskInfo> logicalTasks = logicalTaskIds.stream().map(taskId -> {
final Set<String> stateStoreNames = topologyMetadata
.stateStoreNameToSourceTopicsForTopology(taskId.topologyName())
.keySet();
final Set<TaskTopicPartition> topicPartitions = topicPartitionsForTask.get(taskId);
return new DefaultTaskInfo(
taskId,
!stateStoreNames.isEmpty(),
stateStoreNames,
topicPartitions
);
}).collect(Collectors.toSet());
final Map<TaskId, TaskInfo> logicalTasks = logicalTaskIds.stream().collect(Collectors.toMap(
Function.identity(),
taskId -> {
final Set<String> stateStoreNames = topologyMetadata
.stateStoreNameToSourceTopicsForTopology(taskId.topologyName())

@ableegoldman ableegoldman May 30, 2024

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.

Ah somehow I missed this before -- this is actually returning all the state stores for this topology, it's not specific to the taskId. This was an existing issue so we don't need to fix it in this PR, it can be addressed in a followup. It might be a bit complicated so I'll take a look at how we can get this info

Would've caught this during testing since we definitely want tests with mixed stateless-and-stateful tasks, but still good to fix ASAP

.keySet();
final Set<TaskTopicPartition> topicPartitions = topicPartitionsForTask.get(taskId);
return new DefaultTaskInfo(
taskId,
!stateStoreNames.isEmpty(),
stateStoreNames,
topicPartitions
);
}
));

return new DefaultApplicationState(
assignmentConfigs.toPublicAssignmentConfigs(),
Expand Down Expand Up @@ -728,12 +732,12 @@ private void checkAllPartitions(final Set<String> allSourceTopics,
* populate the stateful tasks that have been assigned to the clients
*/
private UserTaskAssignmentListener assignTasksToClients(final Cluster fullMetadata,
final Set<String> allSourceTopics,
final Map<Subtopology, TopicsInfo> topicGroups,
final Map<UUID, ClientMetadata> clientMetadataMap,
final Map<TaskId, Set<TopicPartition>> partitionsForTask,
final Map<UUID, Map<String, Optional<String>>> racksForProcessConsumer,
final Set<TaskId> statefulTasks) {
final Set<String> allSourceTopics,
final Map<Subtopology, TopicsInfo> topicGroups,
final Map<UUID, ClientMetadata> clientMetadataMap,
final Map<TaskId, Set<TopicPartition>> partitionsForTask,
final Map<UUID, Map<String, Optional<String>>> racksForProcessConsumer,
final Set<TaskId> statefulTasks) {
if (!statefulTasks.isEmpty()) {
throw new TaskAssignmentException("The stateful tasks should not be populated before assigning tasks to clients");
}
Expand Down Expand Up @@ -775,7 +779,7 @@ private UserTaskAssignmentListener assignTasksToClients(final Cluster fullMetada

final Optional<org.apache.kafka.streams.processor.assignment.TaskAssignor> userTaskAssignor =
userTaskAssignorSupplier.get();
UserTaskAssignmentListener userTaskAssignmentListener = (GroupAssignment assignment, GroupSubscription subscription) -> { };
final UserTaskAssignmentListener userTaskAssignmentListener;
if (userTaskAssignor.isPresent()) {
final ApplicationState applicationState = buildApplicationState(
taskManager.topologyMetadata(),
Expand All @@ -785,12 +789,11 @@ private UserTaskAssignmentListener assignTasksToClients(final Cluster fullMetada
);
final org.apache.kafka.streams.processor.assignment.TaskAssignor assignor = userTaskAssignor.get();
final TaskAssignment taskAssignment = assignor.assign(applicationState);
processStreamsPartitionAssignment(clientMetadataMap, taskAssignment);
final AssignmentError assignmentError = validateTaskAssignment(applicationState, taskAssignment);
userTaskAssignmentListener = (GroupAssignment assignment, GroupSubscription subscription) -> {
assignor.onAssignmentComputed(assignment, subscription, assignmentError);
};
processStreamsPartitionAssignment(clientMetadataMap, taskAssignment);
userTaskAssignmentListener = (assignment, subscription) -> assignor.onAssignmentComputed(assignment, subscription, assignmentError);
} else {
userTaskAssignmentListener = (assignment, subscription) -> { };
final TaskAssignor taskAssignor = createTaskAssignor(lagComputationSuccessful);
final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor(
fullMetadata,
Expand Down Expand Up @@ -1564,7 +1567,7 @@ private AssignmentError validateTaskAssignment(final ApplicationState applicatio
final Map<TaskId, ProcessId> standbyTasksInOutput = new HashMap<>();
for (final KafkaStreamsAssignment assignment : assignments) {
final Set<TaskId> tasksForAssignment = new HashSet<>();
for (final KafkaStreamsAssignment.AssignedTask task : assignment.assignment()) {
for (final KafkaStreamsAssignment.AssignedTask task : assignment.tasks().values()) {
if (activeTasksInOutput.containsKey(task.id()) && task.type() == KafkaStreamsAssignment.AssignedTask.Type.ACTIVE) {
log.error("Assignment is invalid: active task {} was assigned to multiple KafkaStreams clients: {} and {}",
task.id(), assignment.processId().id(), activeTasksInOutput.get(task.id()).id());
Expand Down Expand Up @@ -1614,7 +1617,7 @@ private AssignmentError validateTaskAssignment(final ApplicationState applicatio

final Set<TaskId> taskIdsInInput = applicationState.allTasks().keySet();
for (final KafkaStreamsAssignment assignment : assignments) {
for (final KafkaStreamsAssignment.AssignedTask task : assignment.assignment()) {
for (final KafkaStreamsAssignment.AssignedTask task : assignment.tasks().values()) {
if (!taskIdsInInput.contains(task.id())) {
log.error("Assignment is invalid: task {} assigned to KafkaStreams client {} was unknown",
task.id(), assignment.processId().id());
Expand Down
16 changes: 8 additions & 8 deletions ...ms/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,8 @@ public ClientState(final UUID processId, final Map<String, String> clientTags) {
}

ClientState(final UUID processId, final int capacity, final Map<String, String> clientTags) {
previousStandbyTasks.taskIds(new TreeSet<>());
previousActiveTasks.taskIds(new TreeSet<>());
previousStandbyTasks.setTaskIds(new TreeSet<>());
previousActiveTasks.setTaskIds(new TreeSet<>());
taskOffsetSums = new TreeMap<>();
taskLagTotals = new TreeMap<>();
this.capacity = capacity;
Expand All @@ -110,8 +110,8 @@ public ClientState(final Set<TaskId> previousActiveTasks,
final Map<String, String> clientTags,
final int capacity,
final UUID processId) {
this.previousStandbyTasks.taskIds(unmodifiableSet(new TreeSet<>(previousStandbyTasks)));
this.previousActiveTasks.taskIds(unmodifiableSet(new TreeSet<>(previousActiveTasks)));
this.previousStandbyTasks.setTaskIds(unmodifiableSet(new TreeSet<>(previousStandbyTasks)));
this.previousActiveTasks.setTaskIds(unmodifiableSet(new TreeSet<>(previousActiveTasks)));
taskOffsetSums = emptyMap();
this.taskLagTotals = unmodifiableMap(taskLagTotals);
this.capacity = capacity;
Expand Down Expand Up @@ -489,14 +489,14 @@ public SortedMap<String, Set<TaskId>> taskIdsByPreviousConsumer() {
}

public void setAssignedTasks(final KafkaStreamsAssignment assignment) {
final Set<TaskId> activeTasks = assignment.assignment().stream()
final Set<TaskId> activeTasks = assignment.tasks().values().stream()
.filter(task -> task.type() == ACTIVE).map(KafkaStreamsAssignment.AssignedTask::id)
.collect(Collectors.toSet());
final Set<TaskId> standbyTasks = assignment.assignment().stream()
final Set<TaskId> standbyTasks = assignment.tasks().values().stream()
.filter(task -> task.type() == STANDBY).map(KafkaStreamsAssignment.AssignedTask::id)
.collect(Collectors.toSet());
assignedActiveTasks.taskIds(activeTasks);
assignedStandbyTasks.taskIds(standbyTasks);
assignedActiveTasks.setTaskIds(activeTasks);
assignedStandbyTasks.setTaskIds(standbyTasks);
}

public String currentAssignment() {
Expand Down
2 changes: 1 addition & 1 deletion ...rc/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ class ClientStateTask {
this.consumerToTaskIds = consumerToTaskIds;
}

void taskIds(final Set<TaskId> clientToTaskIds) {
void setTaskIds(final Set<TaskId> clientToTaskIds) {
taskIds = clientToTaskIds;
}

Expand Down
14 changes: 7 additions & 7 deletions .../java/org/apache/kafka/streams/processor/internals/assignment/ConstrainedPrioritySet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,22 +30,22 @@
/**
* Wraps a priority queue of clients and returns the next valid candidate(s) based on the current task assignment
*/
class ConstrainedPrioritySet {
public class ConstrainedPrioritySet {

private final PriorityQueue<UUID> clientsByTaskLoad;
private final BiFunction<UUID, TaskId, Boolean> constraint;
private final Set<UUID> uniqueClients = new HashSet<>();

ConstrainedPrioritySet(final BiFunction<UUID, TaskId, Boolean> constraint,
final Function<UUID, Double> weight) {
public ConstrainedPrioritySet(final BiFunction<UUID, TaskId, Boolean> constraint,
final Function<UUID, Double> weight) {
this.constraint = constraint;
clientsByTaskLoad = new PriorityQueue<>(Comparator.comparing(weight).thenComparing(clientId -> clientId));
}

/**
* @return the next least loaded client that satisfies the given criteria, or null if none do
*/
UUID poll(final TaskId task, final Function<UUID, Boolean> extraConstraint) {
public UUID poll(final TaskId task, final Function<UUID, Boolean> extraConstraint) {
final Set<UUID> invalidPolledClients = new HashSet<>();
while (!clientsByTaskLoad.isEmpty()) {
final UUID candidateClient = pollNextClient();
Expand All @@ -66,17 +66,17 @@ UUID poll(final TaskId task, final Function<UUID, Boolean> extraConstraint) {
/**
* @return the next least loaded client that satisfies the given criteria, or null if none do
*/
UUID poll(final TaskId task) {
public UUID poll(final TaskId task) {
return poll(task, client -> true);
}

void offerAll(final Collection<UUID> clients) {
public void offerAll(final Collection<UUID> clients) {
for (final UUID client : clients) {
offer(client);
}
}

void offer(final UUID client) {
public void offer(final UUID client) {
if (uniqueClients.contains(client)) {
clientsByTaskLoad.remove(client);
} else {
Expand Down
6 changes: 2 additions & 4 deletions ...java/org/apache/kafka/streams/processor/internals/assignment/DefaultApplicationState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,7 @@
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;
import org.apache.kafka.streams.processor.assignment.TaskInfo;
import org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor.ClientMetadata;
import org.apache.kafka.streams.processor.TaskId;
Expand All @@ -42,10 +40,10 @@ public class DefaultApplicationState implements ApplicationState {
private final Map<Boolean, Map<ProcessId, KafkaStreamsState>> cachedKafkaStreamStates;

public DefaultApplicationState(final AssignmentConfigs assignmentConfigs,
final Set<TaskInfo> tasks,
final Map<TaskId, TaskInfo> tasks,
final Map<UUID, ClientMetadata> clientStates) {
this.assignmentConfigs = assignmentConfigs;
this.tasks = unmodifiableMap(tasks.stream().collect(Collectors.toMap(TaskInfo::id, task -> task)));
this.tasks = unmodifiableMap(tasks);
this.clientStates = clientStates;
this.cachedKafkaStreamStates = new HashMap<>();
}
Expand Down
2 changes: 1 addition & 1 deletion ...n/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ int getCost(final TaskId taskId,

// This is number is picked based on testing. Usually the optimization for standby assignment
// stops after 3 rounds
private static final int STANDBY_OPTIMIZER_MAX_ITERATION = 4;
public static final int STANDBY_OPTIMIZER_MAX_ITERATION = 4;

private final Cluster fullMetadata;
private final Map<TaskId, Set<TopicPartition>> partitionsForTask;
Expand Down