mirror of https://github.com/apache/kafka.git
KAFKA-18319: Add task assignor interfaces (#18270)
Introduces interfaces for defining task assignors. Task assignors are pure functions, mapping the state of the group and a topology to a target assignment. We include a mock assignor, which we will be able to use when testing / benchmarking without the complexities of the sticky task assignor and the high-availability task assignor. We may remove the mock assignor in before the streams rebalance protocol goes GA. The consumer groups introduce these interfaces to establish a clear separation between the group coordinator code and the pluggable assignors, which may live outside the group coordinator code. We have removed pluggable assignors in KIP-1071, but I think it still makes sense to keep these interfaces for having a clean interface for people to code against. This will pay off, if we plan on making the task assignors pluggable later. Reviewers: Bill Bejeck <bbejeck@gmail.com>, David Jacot <djacot@confluent.io>
This commit is contained in:
Lucas Brutschy
GitHub
parent
0344f8f5ae
commit
dd0fd55316
|
|
@ -0,0 +1,60 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.Collections;
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
import java.util.Optional;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* The assignment specification for a Streams group member.
|
||||
*
|
||||
* @param instanceId The instance ID if provided.
|
||||
* @param rackId The rack ID if provided.
|
||||
* @param activeTasks Reconciled active tasks
|
||||
* @param standbyTasks Reconciled standby tasks
|
||||
* @param warmupTasks Reconciled warm-up tasks
|
||||
* @param processId The process ID.
|
||||
* @param clientTags The client tags for a rack-aware assignment.
|
||||
* @param taskOffsets The last received cumulative task offsets of assigned tasks or dormant tasks.
|
||||
*/
|
||||
public record AssignmentMemberSpec(Optional<String> instanceId,
|
||||
Optional<String> rackId,
|
||||
Map<String, Set<Integer>> activeTasks,
|
||||
Map<String, Set<Integer>> standbyTasks,
|
||||
Map<String, Set<Integer>> warmupTasks,
|
||||
String processId,
|
||||
Map<String, String> clientTags,
|
||||
Map<TaskId, Long> taskOffsets,
|
||||
Map<TaskId, Long> taskEndOffsets
|
||||
) {
|
||||
|
||||
public AssignmentMemberSpec {
|
||||
Objects.requireNonNull(instanceId);
|
||||
Objects.requireNonNull(rackId);
|
||||
activeTasks = Collections.unmodifiableMap(Objects.requireNonNull(activeTasks));
|
||||
standbyTasks = Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks));
|
||||
warmupTasks = Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks));
|
||||
Objects.requireNonNull(processId);
|
||||
clientTags = Collections.unmodifiableMap(Objects.requireNonNull(clientTags));
|
||||
taskOffsets = Collections.unmodifiableMap(Objects.requireNonNull(taskOffsets));
|
||||
taskEndOffsets = Collections.unmodifiableMap(Objects.requireNonNull(taskEndOffsets));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* The task assignment for a streams group.
|
||||
*
|
||||
* @param members The member assignments keyed by member id.
|
||||
*/
|
||||
public record GroupAssignment(Map<String, MemberAssignment> members) {
|
||||
|
||||
public GroupAssignment {
|
||||
Objects.requireNonNull(members);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,36 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* The group metadata specifications required to compute the target assignment.
|
||||
*/
|
||||
public interface GroupSpec {
|
||||
|
||||
/**
|
||||
* @return Member metadata keyed by member Id.
|
||||
*/
|
||||
Map<String, AssignmentMemberSpec> members();
|
||||
|
||||
/**
|
||||
* @return Any configurations passed to the assignor.
|
||||
*/
|
||||
Map<String, String> assignmentConfigs();
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,36 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* The assignment specification for a Streams group.
|
||||
*
|
||||
* @param members The member metadata keyed by member ID.
|
||||
* @param assignmentConfigs Any configurations passed to the assignor.
|
||||
*/
|
||||
public record GroupSpecImpl(Map<String, AssignmentMemberSpec> members,
|
||||
Map<String, String> assignmentConfigs) implements GroupSpec {
|
||||
|
||||
public GroupSpecImpl {
|
||||
Objects.requireNonNull(members);
|
||||
Objects.requireNonNull(assignmentConfigs);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,42 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.Collections;
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* The task assignment for a Streams group member.
|
||||
*
|
||||
* @param activeTasks The target tasks assigned to this member keyed by subtopologyId.
|
||||
*/
|
||||
public record MemberAssignment(Map<String, Set<Integer>> activeTasks,
|
||||
Map<String, Set<Integer>> standbyTasks,
|
||||
Map<String, Set<Integer>> warmupTasks) {
|
||||
|
||||
public MemberAssignment {
|
||||
Objects.requireNonNull(activeTasks);
|
||||
Objects.requireNonNull(standbyTasks);
|
||||
Objects.requireNonNull(warmupTasks);
|
||||
}
|
||||
|
||||
public static MemberAssignment empty() {
|
||||
return new MemberAssignment(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,106 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.Comparator;
|
||||
import java.util.HashMap;
|
||||
import java.util.HashSet;
|
||||
import java.util.Map;
|
||||
import java.util.PriorityQueue;
|
||||
import java.util.Set;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
/**
|
||||
* Mock implementation of {@link TaskAssignor} that assigns tasks to members in a round-robin fashion, with a bit of stickiness.
|
||||
*/
|
||||
public class MockAssignor implements TaskAssignor {
|
||||
|
||||
public static final String MOCK_ASSIGNOR_NAME = "mock";
|
||||
|
||||
@Override
|
||||
public String name() {
|
||||
return MOCK_ASSIGNOR_NAME;
|
||||
}
|
||||
|
||||
@Override
|
||||
public GroupAssignment assign(
|
||||
final GroupSpec groupSpec,
|
||||
final TopologyDescriber topologyDescriber
|
||||
) throws TaskAssignorException {
|
||||
|
||||
Map<String, MemberAssignment> newTargetAssignment = new HashMap<>();
|
||||
Map<String, String[]> subtopologyToActiveMember = new HashMap<>();
|
||||
|
||||
for (String subtopology : topologyDescriber.subtopologies()) {
|
||||
int numberOfPartitions = topologyDescriber.numTasks(subtopology);
|
||||
subtopologyToActiveMember.put(subtopology, new String[numberOfPartitions]);
|
||||
}
|
||||
|
||||
// Copy existing assignment and fill temporary data structures
|
||||
for (Map.Entry<String, AssignmentMemberSpec> memberEntry : groupSpec.members().entrySet()) {
|
||||
final String memberId = memberEntry.getKey();
|
||||
final AssignmentMemberSpec memberSpec = memberEntry.getValue();
|
||||
|
||||
Map<String, Set<Integer>> activeTasks = new HashMap<>(memberSpec.activeTasks());
|
||||
|
||||
newTargetAssignment.put(memberId, new MemberAssignment(activeTasks, new HashMap<>(), new HashMap<>()));
|
||||
for (Map.Entry<String, Set<Integer>> entry : activeTasks.entrySet()) {
|
||||
final String subtopologyId = entry.getKey();
|
||||
final Set<Integer> taskIds = entry.getValue();
|
||||
final String[] activeMembers = subtopologyToActiveMember.get(subtopologyId);
|
||||
for (int taskId : taskIds) {
|
||||
if (activeMembers[taskId] != null) {
|
||||
throw new TaskAssignorException(
|
||||
"Task " + taskId + " of subtopology " + subtopologyId + " is assigned to multiple members");
|
||||
}
|
||||
activeMembers[taskId] = memberId;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Define priority queue to sort members by task count
|
||||
PriorityQueue<MemberAndTaskCount> memberAndTaskCount = new PriorityQueue<>(Comparator.comparingInt(m -> m.taskCount));
|
||||
memberAndTaskCount.addAll(
|
||||
newTargetAssignment.keySet().stream()
|
||||
.map(memberId -> new MemberAndTaskCount(memberId,
|
||||
newTargetAssignment.get(memberId).activeTasks().values().stream().mapToInt(Set::size).sum()))
|
||||
.collect(Collectors.toSet())
|
||||
);
|
||||
|
||||
// Assign unassigned tasks to members with the fewest tasks
|
||||
for (Map.Entry<String, String[]> entry : subtopologyToActiveMember.entrySet()) {
|
||||
final String subtopologyId = entry.getKey();
|
||||
final String[] activeMembers = entry.getValue();
|
||||
for (int i = 0; i < activeMembers.length; i++) {
|
||||
if (activeMembers[i] == null) {
|
||||
final MemberAndTaskCount m = memberAndTaskCount.poll();
|
||||
if (m == null) {
|
||||
throw new TaskAssignorException("No member available to assign task " + i + " of subtopology " + subtopologyId);
|
||||
}
|
||||
newTargetAssignment.get(m.memberId).activeTasks().computeIfAbsent(subtopologyId, k -> new HashSet<>()).add(i);
|
||||
activeMembers[i] = m.memberId;
|
||||
memberAndTaskCount.add(new MemberAndTaskCount(m.memberId, m.taskCount + 1));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return new GroupAssignment(newTargetAssignment);
|
||||
}
|
||||
|
||||
private record MemberAndTaskCount(String memberId, int taskCount) {
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
/**
|
||||
* Server side task assignor used by streams groups.
|
||||
*/
|
||||
public interface TaskAssignor {
|
||||
|
||||
/**
|
||||
* Unique name for this assignor.
|
||||
*/
|
||||
String name();
|
||||
|
||||
/**
|
||||
* Assigns tasks to group members based on the given assignment specification and topic metadata.
|
||||
*
|
||||
* @param groupSpec The assignment spec which includes member metadata.
|
||||
* @param topologyDescriber The task metadata describer.
|
||||
* @return The new assignment for the group.
|
||||
*/
|
||||
GroupAssignment assign(
|
||||
GroupSpec groupSpec,
|
||||
TopologyDescriber topologyDescriber
|
||||
) throws TaskAssignorException;
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import org.apache.kafka.common.errors.ApiException;
|
||||
|
||||
/**
|
||||
* Exception thrown by {@link TaskAssignor#assign(GroupSpec, TopologyDescriber)}}. The exception is only used internally.
|
||||
*/
|
||||
public class TaskAssignorException extends ApiException {
|
||||
|
||||
public TaskAssignorException(String message) {
|
||||
super(message);
|
||||
}
|
||||
|
||||
public TaskAssignorException(String message, Throwable cause) {
|
||||
super(message, cause);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.Comparator;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* The identifier for a task
|
||||
*
|
||||
* @param subtopologyId The unique identifier of the subtopology.
|
||||
* @param partition The partition of the input topics this task is processing.
|
||||
*/
|
||||
public record TaskId(String subtopologyId, int partition) implements Comparable<TaskId> {
|
||||
|
||||
public TaskId {
|
||||
Objects.requireNonNull(subtopologyId);
|
||||
}
|
||||
|
||||
@Override
|
||||
public int compareTo(final TaskId other) {
|
||||
return Comparator.comparing(TaskId::subtopologyId)
|
||||
.thenComparingInt(TaskId::partition)
|
||||
.compare(this, other);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.NoSuchElementException;
|
||||
|
||||
/**
|
||||
* The subscribed topic describer is used by the {@link TaskAssignor} to obtain topic and task metadata of the groups topology.
|
||||
*/
|
||||
public interface TopologyDescriber {
|
||||
|
||||
/**
|
||||
* @return The list of subtopologies IDs.
|
||||
*/
|
||||
List<String> subtopologies();
|
||||
|
||||
/**
|
||||
* The number of tasks for the given subtopology.
|
||||
*
|
||||
* @param subtopologyId String identifying the subtopology.
|
||||
*
|
||||
* @return The number of tasks corresponding to the given subtopology ID.
|
||||
* @throws NoSuchElementException if subtopology does not exist in the topology.
|
||||
*/
|
||||
int numTasks(String subtopologyId) throws NoSuchElementException;
|
||||
|
||||
/**
|
||||
* Whether the given subtopology is stateful.
|
||||
*
|
||||
* @param subtopologyId String identifying the subtopology.
|
||||
* @return true if the subtopology is stateful, false otherwise.
|
||||
*/
|
||||
boolean isStateful(String subtopologyId);
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,62 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.Collections;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
|
||||
|
||||
public class GroupSpecImplTest {
|
||||
|
||||
private Map<String, AssignmentMemberSpec> members;
|
||||
private GroupSpecImpl groupSpec;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
members = new HashMap<>();
|
||||
|
||||
members.put("test-member", new AssignmentMemberSpec(
|
||||
Optional.of("test-instance"),
|
||||
Optional.of("test-rack"),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
));
|
||||
|
||||
groupSpec = new GroupSpecImpl(
|
||||
members,
|
||||
new HashMap<>()
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMembers() {
|
||||
assertEquals(members, groupSpec.members());
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,267 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.kafka.coordinator.group.streams.assignor;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.Collections;
|
||||
import java.util.HashMap;
|
||||
import java.util.HashSet;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.Set;
|
||||
|
||||
import static org.apache.kafka.common.utils.Utils.mkEntry;
|
||||
import static org.apache.kafka.common.utils.Utils.mkMap;
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertNotNull;
|
||||
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
|
||||
public class MockAssignorTest {
|
||||
|
||||
private final MockAssignor assignor = new MockAssignor();
|
||||
|
||||
@Test
|
||||
public void testZeroMembers() {
|
||||
|
||||
TaskAssignorException ex = assertThrows(TaskAssignorException.class, () -> assignor.assign(
|
||||
new GroupSpecImpl(
|
||||
Collections.emptyMap(),
|
||||
new HashMap<>()
|
||||
),
|
||||
new TopologyDescriberImpl(5, Collections.singletonList("test-subtopology"))
|
||||
));
|
||||
|
||||
assertEquals("No member available to assign task 0 of subtopology test-subtopology", ex.getMessage());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDoubleAssignment() {
|
||||
|
||||
final AssignmentMemberSpec memberSpec1 = new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
Collections.singletonMap("test-subtopology", new HashSet<>(List.of(0))),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
);
|
||||
|
||||
final AssignmentMemberSpec memberSpec2 = new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
Collections.singletonMap("test-subtopology", new HashSet<>(List.of(0))),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
);
|
||||
|
||||
TaskAssignorException ex = assertThrows(TaskAssignorException.class, () -> assignor.assign(
|
||||
new GroupSpecImpl(
|
||||
Map.of("member1", memberSpec1, "member2", memberSpec2),
|
||||
new HashMap<>()
|
||||
),
|
||||
new TopologyDescriberImpl(5, Collections.singletonList("test-subtopology"))
|
||||
));
|
||||
|
||||
assertEquals("Task 0 of subtopology test-subtopology is assigned to multiple members", ex.getMessage());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testBasicScenario() {
|
||||
|
||||
final GroupAssignment result = assignor.assign(
|
||||
new GroupSpecImpl(
|
||||
Collections.emptyMap(),
|
||||
new HashMap<>()
|
||||
),
|
||||
new TopologyDescriberImpl(5, Collections.emptyList())
|
||||
);
|
||||
|
||||
assertEquals(0, result.members().size());
|
||||
}
|
||||
|
||||
|
||||
@Test
|
||||
public void testSingleMember() {
|
||||
|
||||
final AssignmentMemberSpec memberSpec = new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
);
|
||||
|
||||
final GroupAssignment result = assignor.assign(
|
||||
new GroupSpecImpl(
|
||||
Collections.singletonMap("test_member", memberSpec),
|
||||
new HashMap<>()
|
||||
),
|
||||
new TopologyDescriberImpl(4, List.of("test-subtopology"))
|
||||
);
|
||||
|
||||
assertEquals(1, result.members().size());
|
||||
final MemberAssignment testMember = result.members().get("test_member");
|
||||
assertNotNull(testMember);
|
||||
assertEquals(mkMap(
|
||||
mkEntry("test-subtopology", Set.of(0, 1, 2, 3))
|
||||
), testMember.activeTasks());
|
||||
}
|
||||
|
||||
|
||||
@Test
|
||||
public void testTwoMembersTwoSubtopologies() {
|
||||
|
||||
final AssignmentMemberSpec memberSpec1 = new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
);
|
||||
|
||||
final AssignmentMemberSpec memberSpec2 = new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
);
|
||||
|
||||
final GroupAssignment result = assignor.assign(
|
||||
new GroupSpecImpl(
|
||||
mkMap(mkEntry("test_member1", memberSpec1), mkEntry("test_member2", memberSpec2)),
|
||||
new HashMap<>()
|
||||
),
|
||||
new TopologyDescriberImpl(4, List.of("test-subtopology1", "test-subtopology2"))
|
||||
);
|
||||
|
||||
final Map<String, Set<Integer>> expected1 = mkMap(
|
||||
mkEntry("test-subtopology1", Set.of(1, 3)),
|
||||
mkEntry("test-subtopology2", Set.of(1, 3))
|
||||
);
|
||||
final Map<String, Set<Integer>> expected2 = mkMap(
|
||||
mkEntry("test-subtopology1", Set.of(0, 2)),
|
||||
mkEntry("test-subtopology2", Set.of(0, 2))
|
||||
);
|
||||
|
||||
assertEquals(2, result.members().size());
|
||||
final MemberAssignment testMember1 = result.members().get("test_member1");
|
||||
final MemberAssignment testMember2 = result.members().get("test_member2");
|
||||
assertNotNull(testMember1);
|
||||
assertNotNull(testMember2);
|
||||
assertTrue(expected1.equals(testMember1.activeTasks()) || expected2.equals(testMember1.activeTasks()));
|
||||
assertTrue(expected1.equals(testMember2.activeTasks()) || expected2.equals(testMember2.activeTasks()));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTwoMembersTwoSubtopologiesStickiness() {
|
||||
|
||||
final AssignmentMemberSpec memberSpec1 = new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
mkMap(
|
||||
mkEntry("test-subtopology1", new HashSet<>(List.of(0, 2, 3))),
|
||||
mkEntry("test-subtopology2", new HashSet<>(List.of(0)))
|
||||
),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
);
|
||||
|
||||
final AssignmentMemberSpec memberSpec2 = new AssignmentMemberSpec(
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
mkMap(
|
||||
mkEntry("test-subtopology1", new HashSet<>(List.of(1))),
|
||||
mkEntry("test-subtopology2", new HashSet<>(List.of(3)))
|
||||
),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
"test-process",
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap(),
|
||||
Collections.emptyMap()
|
||||
);
|
||||
final GroupAssignment result = assignor.assign(
|
||||
new GroupSpecImpl(
|
||||
mkMap(mkEntry("test_member1", memberSpec1), mkEntry("test_member2", memberSpec2)),
|
||||
new HashMap<>()
|
||||
),
|
||||
new TopologyDescriberImpl(4, List.of("test-subtopology1", "test-subtopology2"))
|
||||
);
|
||||
|
||||
assertEquals(2, result.members().size());
|
||||
final MemberAssignment testMember1 = result.members().get("test_member1");
|
||||
final MemberAssignment testMember2 = result.members().get("test_member2");
|
||||
assertNotNull(testMember1);
|
||||
assertNotNull(testMember2);
|
||||
assertEquals(mkMap(
|
||||
mkEntry("test-subtopology1", Set.of(0, 2, 3)),
|
||||
mkEntry("test-subtopology2", Set.of(0))
|
||||
), testMember1.activeTasks());
|
||||
assertEquals(mkMap(
|
||||
mkEntry("test-subtopology1", Set.of(1)),
|
||||
mkEntry("test-subtopology2", Set.of(1, 2, 3))
|
||||
), testMember2.activeTasks());
|
||||
}
|
||||
|
||||
private record TopologyDescriberImpl(int numPartitions, List<String> subtopologies) implements TopologyDescriber {
|
||||
|
||||
@Override
|
||||
public List<String> subtopologies() {
|
||||
return subtopologies;
|
||||
}
|
||||
|
||||
@Override
|
||||
public int numTasks(String subtopologyId) {
|
||||
return numPartitions;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean isStateful(String subtopologyId) {
|
||||
return false;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Reference in New Issue