You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by "lucasbru (via GitHub)" <gi...@apache.org> on 2025/01/10 11:23:33 UTC
[PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
lucasbru opened a new pull request, #18476:
URL: https://github.com/apache/kafka/pull/18476
Implements the current assignment builder, analogous to the current assignment builder of consumer groups. The main difference is the underlying assigned resource, and slightly different logic around process IDs: We make sure to move a task only to a new client, once the task is not owned anymore by any client with the same process ID (sharing the same state directory) - in any role (active, standby or warm-up).
Compared to the feature branch, the main difference is that I refactored the separate treatment of active, standby and warm-up tasks into a compound datatype called `TaskTuple` (which is used in place of the more specific `Assignment` class). This also has effects on `StreamsGroupMember`.
### Committer Checklist (excluded from commit message)
- [ ] Verify design and implementation
- [ ] Verify test coverage and CI build status
- [ ] Verify documentation (including upgrade notes)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1922352806
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(
+ Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask
+ ) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ Set<Integer> currentPartitions = currentAssignment.getOrDefault(subtopologyId, Collections.emptySet());
Review Comment:
I did that on purpose. We'll otherwise call (String, Integer) a task and (Integer) a task. I found that very confusing. Also in the `TaskId` class, we call the integer the partition.
But if you think my solution is more confusing, we can also call (String, Integer) and Integer both a task. We'll get naming collisions - `resultAssignedTasks` for all subtopologies vs. `resultAssignedTasks` for one subtopology. I appended `ForThisSubtopology` to all the collections that contain only integers. I also extracted it into a separate method since it became super confusing to me - not sure it helped.
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1926531384
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
Review Comment:
Fine, done
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
Review Comment:
DOne
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,702 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
Review Comment:
I added coverage for both cases
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/StreamsGroupMember.java:
##########
@@ -432,32 +364,9 @@ private static List<StreamsGroupDescribeResponseData.TaskIds> taskIdsFromMap(
}
/**
- * @return True if the two provided members have different assigned active tasks.
- */
- public static boolean hasAssignedActiveTasksChanged(
- StreamsGroupMember member1,
- StreamsGroupMember member2
- ) {
- return !member1.assignedActiveTasks().equals(member2.assignedActiveTasks());
- }
-
- /**
- * @return True if the two provided members have different assigned active tasks.
- */
- public static boolean hasAssignedStandbyTasksChanged(
- StreamsGroupMember member1,
- StreamsGroupMember member2
- ) {
- return !member1.assignedStandbyTasks().equals(member2.assignedStandbyTasks());
- }
-
- /**
- * @return True if the two provided members have different assigned active tasks.
+ * @return True if the two provided members have different assigned tasks.
*/
- public static boolean hasAssignedWarmupTasksChanged(
- StreamsGroupMember member1,
- StreamsGroupMember member2
- ) {
- return !member1.assignedWarmupTasks().equals(member2.assignedWarmupTasks());
+ public static boolean hasAssignedTasksChanged(StreamsGroupMember member1, StreamsGroupMember member2) {
Review Comment:
Ah, because hasAssignedActiveTasksChanged, hasAssignedStandbyTasksChanged and hasAssignedWarmupTasksChanged also didn't have unit tests. I guess we forgot to add them. I added one for the refactored method.
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStableWithNewTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 4)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 3)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasksWithOwnedTasksNotHavingRevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, __) ->
+ subtopologyId2.equals(subtopologyId) ? "process" : null
+ )
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testRemainsInUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ CurrentAssignmentBuilder currentAssignmentBuilder = new CurrentAssignmentBuilder(
+ member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet());
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(null)
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5, 6)))
+ .build()
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 5)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
Review Comment:
Done
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,451 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(int targetAssignmentEpoch,
+ TaskTuple targetAssignment) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the process ID of a subtopology ID /
+ * partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(BiFunction<String, Integer, String> currentActiveTaskProcessId) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the process IDs of a subtopology
+ * ID / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the process IDs of a subtopology ID
+ * / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(TaskTuple ownedAssignment) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ hasUnreleasedTasks |= computeAssignmentDifferenceForOneSubtopology(
+ subtopologyId,
+ currentAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ targetAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ resultAssignedTasks,
+ resultTasksPendingRevocation,
+ resultTasksPendingAssignment,
+ isUnreleasedTask
+ );
+ }
+ return hasUnreleasedTasks;
+ }
+
+ private static boolean computeAssignmentDifferenceForOneSubtopology(final String subtopologyId,
+ final Set<Integer> currentTasksForThisSubtopology,
+ final Set<Integer> targetTasksForThisSubtopology,
+ final Map<String, Set<Integer>> resultAssignedTasks,
+ final Map<String, Set<Integer>> resultTasksPendingRevocation,
+ final Map<String, Set<Integer>> resultTasksPendingAssignment,
+ final BiPredicate<String, Integer> isUnreleasedTask) {
+ // Result Assigned Tasks = Current Tasks ∩ Target Tasks
+ // i.e. we remove all tasks from the current assignment that are not in the target
+ // assignment
+ Set<Integer> resultAssignedTasksForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultAssignedTasksForThisSubtopology.retainAll(targetTasksForThisSubtopology);
+
+ // Result Tasks Pending Revocation = Current Tasks - Result Assigned Tasks
+ // i.e. we will ask the member to revoke all tasks in its current assignment that
+ // are not in the target assignment
+ Set<Integer> resultTasksPendingRevocationForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultTasksPendingRevocationForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+
+ // Result Tasks Pending Assignment = Target Tasks - Result Assigned Tasks - Unreleased Tasks
+ // i.e. we will ask the member to assign all tasks in its target assignment,
+ // except those that are already assigned, and those that are unreleased
+ Set<Integer> resultTasksPendingAssignmentForThisSubtopology = new HashSet<>(targetTasksForThisSubtopology);
+ resultTasksPendingAssignmentForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+ boolean hasUnreleasedTasks = resultTasksPendingAssignmentForThisSubtopology.removeIf(taskId ->
+ isUnreleasedTask.test(subtopologyId, taskId)
+ );
+
+ if (!resultAssignedTasksForThisSubtopology.isEmpty()) {
+ resultAssignedTasks.put(subtopologyId, resultAssignedTasksForThisSubtopology);
+ }
+
+ if (!resultTasksPendingRevocationForThisSubtopology.isEmpty()) {
+ resultTasksPendingRevocation.put(subtopologyId, resultTasksPendingRevocationForThisSubtopology);
+ }
+
+ if (!resultTasksPendingAssignmentForThisSubtopology.isEmpty()) {
+ resultTasksPendingAssignment.put(subtopologyId, resultTasksPendingAssignmentForThisSubtopology);
+ }
+
+ return hasUnreleasedTasks;
+ }
+
+ /**
+ * Computes the next assignment.
+ *
+ * @param memberEpoch The epoch of the member to use. This may be different from
+ * the epoch in {@link CurrentAssignmentBuilder#member}.
+ * @param memberAssignedTasks The assigned tasks of the member to use.
+ * @return A new StreamsGroupMember.
+ */
+ private StreamsGroupMember computeNextAssignment(int memberEpoch,
+ TaskTuple memberAssignedTasks) {
+ Map<String, Set<Integer>> newActiveAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingAssignment = new HashMap<>();
+
+ boolean hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.activeTasks(),
+ targetAssignment.activeTasks(),
+ newActiveAssignedTasks,
+ newActiveTasksPendingRevocation,
+ newActiveTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ currentActiveTaskProcessId.apply(subtopologyId, partitionId) != null ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ );
+
+ hasUnreleasedTasks = computeAssignmentDifference(
Review Comment:
having separate booleans could make things clearer. Updated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1922191790
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
Review Comment:
I considered this, but in the javadoc I saw that this was very confusing. For example
```A new task tuple, containing all active tasks, standby tasks and warm-up tasks from both tuples.```
```A new [[Tasks]], containing all active tasks, standby tasks and warm-up tasks from both given [[Tasks]].```
I'll keep the tuple name for now.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
Review Comment:
Done
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
Review Comment:
Since you are asking - I don't think so. I find calculations in the test harder to read than just constants.
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
Review Comment:
Since you are asking - I don't think so. I find calculations in the test harder to read than just constants.
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStableWithNewTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 4)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 3)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasksWithOwnedTasksNotHavingRevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, __) ->
+ subtopologyId2.equals(subtopologyId) ? "process" : null
+ )
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testRemainsInUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ CurrentAssignmentBuilder currentAssignmentBuilder = new CurrentAssignmentBuilder(
+ member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet());
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(null)
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5, 6)))
+ .build()
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 5)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
Review Comment:
Done
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/TaskAssignmentTestUtil.java:
##########
@@ -17,37 +17,41 @@
package org.apache.kafka.coordinator.group.streams;
import java.util.AbstractMap;
-import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
-import java.util.Objects;
import java.util.Set;
public class TaskAssignmentTestUtil {
- public static Assignment mkAssignment(final Map<String, Set<Integer>> activeTasks,
- final Map<String, Set<Integer>> standbyTasks,
- final Map<String, Set<Integer>> warmupTasks) {
- return new Assignment(
- Collections.unmodifiableMap(Objects.requireNonNull(activeTasks)),
- Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks)),
- Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks))
- );
+ public enum TaskRole {
+ ACTIVE,
+ STANDBY,
+ WARMUP
+ }
+
+ @SafeVarargs
+ public static TaskTuple mkTaskTuple(TaskRole taskRole, Map.Entry<String, Set<Integer>>... entries) {
+ return switch (taskRole) {
+ case ACTIVE -> new TaskTuple(mkTasksPerSubtopology(entries), new HashMap<>(), new HashMap<>());
+ case STANDBY -> new TaskTuple(new HashMap<>(), mkTasksPerSubtopology(entries), new HashMap<>());
+ case WARMUP -> new TaskTuple(new HashMap<>(), new HashMap<>(), mkTasksPerSubtopology(entries));
+ };
}
- public static Map.Entry<String, Set<Integer>> mkTasks(String subtopologyId,
- Integer... tasks) {
+ public static Map.Entry<String, Set<Integer>> mkTasks(
+ String subtopologyId,
+ Integer... tasks
+ ) {
Review Comment:
Not intentional, I just do it intuitively because I like it better. Most classes in the group coordinator seem use this style, except for the ones you wrote. Anyway, reverted.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
Review Comment:
No, this can be checked centrally when the heartbeat comes in and not again in every class.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
+ Map<String, Set<Integer>> standbyTasks,
+ Map<String, Set<Integer>> warmupTasks) {
+
+ public TaskTuple {
+ activeTasks = Collections.unmodifiableMap(Objects.requireNonNull(activeTasks));
+ standbyTasks = Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks));
+ warmupTasks = Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks));
+ }
+
+ /**
+ * An empty task tuple.
+ */
+ public static final TaskTuple EMPTY = new TaskTuple(
+ Collections.emptyMap(),
+ Collections.emptyMap(),
+ Collections.emptyMap()
+ );
+
+ /**
+ * @return true if all collections in the tuple are empty.
+ */
+ public boolean isEmpty() {
+ return activeTasks.isEmpty() && standbyTasks.isEmpty() && warmupTasks.isEmpty();
+ }
+
+ /**
+ * Merges this task tuple with another task tuple.
+ *
+ * @param other The other task tuple.
+ * @return A new task tuple, containing all active tasks, standby tasks and warm-up tasks from both tuples.
+ */
+ public TaskTuple merge(TaskTuple other) {
+ Map<String, Set<Integer>> mergedActiveTasks = merge(activeTasks, other.activeTasks);
+ Map<String, Set<Integer>> mergedStandbyTasks = merge(standbyTasks, other.standbyTasks);
+ Map<String, Set<Integer>> mergedWarmupTasks = merge(warmupTasks, other.warmupTasks);
+ return new TaskTuple(mergedActiveTasks, mergedStandbyTasks, mergedWarmupTasks);
+ }
+
+ private Map<String, Set<Integer>> merge(final Map<String, Set<Integer>> tasks1, final Map<String, Set<Integer>> tasks2) {
Review Comment:
Done
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
+ Map<String, Set<Integer>> standbyTasks,
+ Map<String, Set<Integer>> warmupTasks) {
+
+ public TaskTuple {
+ activeTasks = Collections.unmodifiableMap(Objects.requireNonNull(activeTasks));
+ standbyTasks = Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks));
+ warmupTasks = Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks));
Review Comment:
Yes
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
Review Comment:
If we know that the client does not own any tasks (and this is the only case where we know it, because we do not store the owned tasks), we already have what the rejoin is aiming to achieve (no owned tasks), so there is no point in fencing.
This comes from the original current assignment builder.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
Review Comment:
Done
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
Review Comment:
Done
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(
+ Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask
+ ) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ Set<Integer> currentPartitions = currentAssignment.getOrDefault(subtopologyId, Collections.emptySet());
+ Set<Integer> targetPartitions = targetAssignment.getOrDefault(subtopologyId, Collections.emptySet());
+
+ // Result Assigned Partitions = Current Partitions ∩ Target Partitions
+ // i.e. we remove all partitions from the current assignment that are not in the target
+ // assignment
+ Set<Integer> resultAssignedPartitions = new HashSet<>(currentPartitions);
+ resultAssignedPartitions.retainAll(targetPartitions);
+
+ // Result Partitions Pending Revocation = Current Partitions - Result Assigned Partitions
+ // i.e. we will ask the member to revoke all partitions in its current assignment that
+ // are not in the target assignment
+ Set<Integer> resultPartitionsPendingRevocation = new HashSet<>(currentPartitions);
+ resultPartitionsPendingRevocation.removeAll(resultAssignedPartitions);
+
+ // Result Partitions Pending Assignment = Target Partitions - Result Assigned Partitions - Unreleased Partitions
+ // i.e. we will ask the member to assign all partitions in its target assignment,
+ // except those that are already assigned, and those that are unrelead
Review Comment:
Done
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
Review Comment:
Done
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
Review Comment:
Good comment! Interesting idea, but I think you are missing something. See, whether `member.tasksPendingRevocation` was defined in a previous heartbeat, which may have had a different target assignment, than the latest assignment. So here, all we know is that we can transition to the next target assignment, because we revoked all tasks that we need to according to that previous target assignment. This doesn't imply much about what will go on in `buildNewMember`, because we are using a new target assignment, which may require other tasks to be revoked. So if we find out, that we actually need to revoke more tasks to reach the target assignment, we will use the result of `memember.meberEpoch() + 1`.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(
+ Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask
+ ) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ Set<Integer> currentPartitions = currentAssignment.getOrDefault(subtopologyId, Collections.emptySet());
Review Comment:
I did that on purpose. We'll otherwise call (String, Integer) a task and (Integer) a task. I found that very confusing. Also in the `TaskId` class, we call the integer the partition.
But if you think my solution is more confusing, we can also call (String, Integer) and Integer both a task. We'll get naming collisions - `resultAssignedTasks` for all subtopologies vs. `resultAssignedTasks` for one subtopology. I appended `ForThisSubtopology` to all the collections to make it compile.
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru merged PR #18476:
URL: https://github.com/apache/kafka/pull/18476
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "bbejeck (via GitHub)" <gi...@apache.org>.
bbejeck commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1926077357
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,702 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
Review Comment:
Minor nit on the test overall - in `build()` for case `STABLE` the test always has the `memberEpoch` not equal the `targetAssignmentEpoch` thus never hitting the `return member` branch and line 293 `resultAssignedTasksForThisSubtopology` is never empty; not sure if either of them are potentially significant.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/StreamsGroupMember.java:
##########
@@ -432,32 +364,9 @@ private static List<StreamsGroupDescribeResponseData.TaskIds> taskIdsFromMap(
}
/**
- * @return True if the two provided members have different assigned active tasks.
- */
- public static boolean hasAssignedActiveTasksChanged(
- StreamsGroupMember member1,
- StreamsGroupMember member2
- ) {
- return !member1.assignedActiveTasks().equals(member2.assignedActiveTasks());
- }
-
- /**
- * @return True if the two provided members have different assigned active tasks.
- */
- public static boolean hasAssignedStandbyTasksChanged(
- StreamsGroupMember member1,
- StreamsGroupMember member2
- ) {
- return !member1.assignedStandbyTasks().equals(member2.assignedStandbyTasks());
- }
-
- /**
- * @return True if the two provided members have different assigned active tasks.
+ * @return True if the two provided members have different assigned tasks.
*/
- public static boolean hasAssignedWarmupTasksChanged(
- StreamsGroupMember member1,
- StreamsGroupMember member2
- ) {
- return !member1.assignedWarmupTasks().equals(member2.assignedWarmupTasks());
+ public static boolean hasAssignedTasksChanged(StreamsGroupMember member1, StreamsGroupMember member2) {
Review Comment:
This method is never used.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,451 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(int targetAssignmentEpoch,
+ TaskTuple targetAssignment) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the process ID of a subtopology ID /
+ * partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(BiFunction<String, Integer, String> currentActiveTaskProcessId) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the process IDs of a subtopology
+ * ID / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the process IDs of a subtopology ID
+ * / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(TaskTuple ownedAssignment) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ hasUnreleasedTasks |= computeAssignmentDifferenceForOneSubtopology(
+ subtopologyId,
+ currentAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ targetAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ resultAssignedTasks,
+ resultTasksPendingRevocation,
+ resultTasksPendingAssignment,
+ isUnreleasedTask
+ );
+ }
+ return hasUnreleasedTasks;
+ }
+
+ private static boolean computeAssignmentDifferenceForOneSubtopology(final String subtopologyId,
+ final Set<Integer> currentTasksForThisSubtopology,
+ final Set<Integer> targetTasksForThisSubtopology,
+ final Map<String, Set<Integer>> resultAssignedTasks,
+ final Map<String, Set<Integer>> resultTasksPendingRevocation,
+ final Map<String, Set<Integer>> resultTasksPendingAssignment,
+ final BiPredicate<String, Integer> isUnreleasedTask) {
+ // Result Assigned Tasks = Current Tasks ∩ Target Tasks
+ // i.e. we remove all tasks from the current assignment that are not in the target
+ // assignment
+ Set<Integer> resultAssignedTasksForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultAssignedTasksForThisSubtopology.retainAll(targetTasksForThisSubtopology);
+
+ // Result Tasks Pending Revocation = Current Tasks - Result Assigned Tasks
+ // i.e. we will ask the member to revoke all tasks in its current assignment that
+ // are not in the target assignment
+ Set<Integer> resultTasksPendingRevocationForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultTasksPendingRevocationForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+
+ // Result Tasks Pending Assignment = Target Tasks - Result Assigned Tasks - Unreleased Tasks
+ // i.e. we will ask the member to assign all tasks in its target assignment,
+ // except those that are already assigned, and those that are unreleased
+ Set<Integer> resultTasksPendingAssignmentForThisSubtopology = new HashSet<>(targetTasksForThisSubtopology);
+ resultTasksPendingAssignmentForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+ boolean hasUnreleasedTasks = resultTasksPendingAssignmentForThisSubtopology.removeIf(taskId ->
+ isUnreleasedTask.test(subtopologyId, taskId)
+ );
+
+ if (!resultAssignedTasksForThisSubtopology.isEmpty()) {
+ resultAssignedTasks.put(subtopologyId, resultAssignedTasksForThisSubtopology);
+ }
+
+ if (!resultTasksPendingRevocationForThisSubtopology.isEmpty()) {
+ resultTasksPendingRevocation.put(subtopologyId, resultTasksPendingRevocationForThisSubtopology);
+ }
+
+ if (!resultTasksPendingAssignmentForThisSubtopology.isEmpty()) {
+ resultTasksPendingAssignment.put(subtopologyId, resultTasksPendingAssignmentForThisSubtopology);
+ }
+
+ return hasUnreleasedTasks;
+ }
+
+ /**
+ * Computes the next assignment.
+ *
+ * @param memberEpoch The epoch of the member to use. This may be different from
+ * the epoch in {@link CurrentAssignmentBuilder#member}.
+ * @param memberAssignedTasks The assigned tasks of the member to use.
+ * @return A new StreamsGroupMember.
+ */
+ private StreamsGroupMember computeNextAssignment(int memberEpoch,
+ TaskTuple memberAssignedTasks) {
+ Map<String, Set<Integer>> newActiveAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingAssignment = new HashMap<>();
+
+ boolean hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.activeTasks(),
+ targetAssignment.activeTasks(),
+ newActiveAssignedTasks,
+ newActiveTasksPendingRevocation,
+ newActiveTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ currentActiveTaskProcessId.apply(subtopologyId, partitionId) != null ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ );
+
+ hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.standbyTasks(),
+ targetAssignment.standbyTasks(),
+ newStandbyAssignedTasks,
+ newStandbyTasksPendingRevocation,
+ newStandbyTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ Objects.equals(currentActiveTaskProcessId.apply(subtopologyId, partitionId),
+ member.processId()) ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ ) || hasUnreleasedTasks;
+
+ hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.warmupTasks(),
+ targetAssignment.warmupTasks(),
+ newWarmupAssignedTasks,
+ newWarmupTasksPendingRevocation,
+ newWarmupTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ Objects.equals(currentActiveTaskProcessId.apply(subtopologyId, partitionId),
+ member.processId()) ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ ) || hasUnreleasedTasks;
+
+ return buildNewMember(
+ memberEpoch,
+ new TaskTuple(
+ newActiveTasksPendingRevocation,
+ newStandbyTasksPendingRevocation,
+ newWarmupTasksPendingRevocation
+ ),
+ new TaskTuple(
+ newActiveAssignedTasks,
+ newStandbyAssignedTasks,
+ newWarmupAssignedTasks
+ ),
+ new TaskTuple(
+ newActiveTasksPendingAssignment,
+ newStandbyTasksPendingAssignment,
+ newWarmupTasksPendingAssignment
+ ),
+ hasUnreleasedTasks
+ );
+ }
+
+ private StreamsGroupMember buildNewMember(final int memberEpoch,
+ final TaskTuple newTasksPendingRevocation,
+ final TaskTuple newAssignedTasks,
+ final TaskTuple newTasksPendingAssignment,
+ final boolean hasUnreleasedTasks) {
+
+ final boolean hasTasksToBeRevoked =
+ (!newTasksPendingRevocation.isEmpty())
+ && (ownedTasks.isEmpty() || ownedTasks.get().containsAny(newTasksPendingRevocation));
+
+ if (hasTasksToBeRevoked) {
+ // If there are tasks to be revoked, the member remains in its current
+ // epoch and requests the revocation of those tasks. It transitions to
+ // the UNREVOKED_TASKS state to wait until the client acknowledges the
+ // revocation of the tasks.
+ return new StreamsGroupMember.Builder(member)
+ .setState(MemberState.UNREVOKED_TASKS)
+ .updateMemberEpoch(memberEpoch)
+ .setAssignedTasks(newAssignedTasks)
+ .setTasksPendingRevocation(newTasksPendingRevocation)
+ .build();
+ } else if (!newTasksPendingAssignment.isEmpty()) {
+ // If there are tasks to be assigned, the member transitions to the
+ // target epoch and requests the assignment of those tasks. Note that
+ // the tasks are directly added to the assigned tasks set. The
+ // member transitions to the STABLE state or to the UNRELEASED_TASKS
+ // state depending on whether there are unreleased tasks or not.
+ MemberState newState =
+ hasUnreleasedTasks
+ ? MemberState.UNRELEASED_TASKS
+ : MemberState.STABLE;
+ return new StreamsGroupMember.Builder(member)
+ .setState(newState)
+ .updateMemberEpoch(targetAssignmentEpoch)
+ .setAssignedTasks(newAssignedTasks.merge(newTasksPendingAssignment))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+ } else if (hasUnreleasedTasks) {
+ // If there are no tasks to be revoked nor to be assigned but some
+ // tasks are not available yet, the member transitions to the target
+ // epoch, to the UNRELEASED_TASKS state and waits.
+ return new StreamsGroupMember.Builder(member)
+ .setState(MemberState.UNRELEASED_TASKS)
Review Comment:
super nit: `UNRELASED_TASKS` implies this member is the one to release them, but if it's waiting for another member to release should there be a state `WAITING_FOR_UNRELEASED` or something similar.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/StreamsGroupMember.java:
##########
@@ -76,22 +66,12 @@ public record StreamsGroupMember(String memberId,
String processId,
Optional<StreamsGroupMemberMetadataValue.Endpoint> userEndpoint,
Map<String, String> clientTags,
- Map<String, Set<Integer>> assignedActiveTasks,
- Map<String, Set<Integer>> assignedStandbyTasks,
- Map<String, Set<Integer>> assignedWarmupTasks,
- Map<String, Set<Integer>> activeTasksPendingRevocation,
- Map<String, Set<Integer>> standbyTasksPendingRevocation,
- Map<String, Set<Integer>> warmupTasksPendingRevocation) {
+ TaskTuple assignedTasks,
+ TaskTuple tasksPendingRevocation) {
Review Comment:
nice cleanup going with the `TaskTuple`
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,451 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(int targetAssignmentEpoch,
+ TaskTuple targetAssignment) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the process ID of a subtopology ID /
+ * partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(BiFunction<String, Integer, String> currentActiveTaskProcessId) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the process IDs of a subtopology
+ * ID / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the process IDs of a subtopology ID
+ * / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(TaskTuple ownedAssignment) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ hasUnreleasedTasks |= computeAssignmentDifferenceForOneSubtopology(
+ subtopologyId,
+ currentAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ targetAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ resultAssignedTasks,
+ resultTasksPendingRevocation,
+ resultTasksPendingAssignment,
+ isUnreleasedTask
+ );
+ }
+ return hasUnreleasedTasks;
+ }
+
+ private static boolean computeAssignmentDifferenceForOneSubtopology(final String subtopologyId,
+ final Set<Integer> currentTasksForThisSubtopology,
+ final Set<Integer> targetTasksForThisSubtopology,
+ final Map<String, Set<Integer>> resultAssignedTasks,
+ final Map<String, Set<Integer>> resultTasksPendingRevocation,
+ final Map<String, Set<Integer>> resultTasksPendingAssignment,
+ final BiPredicate<String, Integer> isUnreleasedTask) {
+ // Result Assigned Tasks = Current Tasks ∩ Target Tasks
+ // i.e. we remove all tasks from the current assignment that are not in the target
+ // assignment
+ Set<Integer> resultAssignedTasksForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultAssignedTasksForThisSubtopology.retainAll(targetTasksForThisSubtopology);
+
+ // Result Tasks Pending Revocation = Current Tasks - Result Assigned Tasks
+ // i.e. we will ask the member to revoke all tasks in its current assignment that
+ // are not in the target assignment
+ Set<Integer> resultTasksPendingRevocationForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultTasksPendingRevocationForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+
+ // Result Tasks Pending Assignment = Target Tasks - Result Assigned Tasks - Unreleased Tasks
+ // i.e. we will ask the member to assign all tasks in its target assignment,
+ // except those that are already assigned, and those that are unreleased
+ Set<Integer> resultTasksPendingAssignmentForThisSubtopology = new HashSet<>(targetTasksForThisSubtopology);
+ resultTasksPendingAssignmentForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+ boolean hasUnreleasedTasks = resultTasksPendingAssignmentForThisSubtopology.removeIf(taskId ->
+ isUnreleasedTask.test(subtopologyId, taskId)
+ );
+
+ if (!resultAssignedTasksForThisSubtopology.isEmpty()) {
+ resultAssignedTasks.put(subtopologyId, resultAssignedTasksForThisSubtopology);
+ }
+
+ if (!resultTasksPendingRevocationForThisSubtopology.isEmpty()) {
+ resultTasksPendingRevocation.put(subtopologyId, resultTasksPendingRevocationForThisSubtopology);
+ }
+
+ if (!resultTasksPendingAssignmentForThisSubtopology.isEmpty()) {
+ resultTasksPendingAssignment.put(subtopologyId, resultTasksPendingAssignmentForThisSubtopology);
+ }
+
+ return hasUnreleasedTasks;
+ }
+
+ /**
+ * Computes the next assignment.
+ *
+ * @param memberEpoch The epoch of the member to use. This may be different from
+ * the epoch in {@link CurrentAssignmentBuilder#member}.
+ * @param memberAssignedTasks The assigned tasks of the member to use.
+ * @return A new StreamsGroupMember.
+ */
+ private StreamsGroupMember computeNextAssignment(int memberEpoch,
+ TaskTuple memberAssignedTasks) {
+ Map<String, Set<Integer>> newActiveAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingAssignment = new HashMap<>();
+
+ boolean hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.activeTasks(),
+ targetAssignment.activeTasks(),
+ newActiveAssignedTasks,
+ newActiveTasksPendingRevocation,
+ newActiveTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ currentActiveTaskProcessId.apply(subtopologyId, partitionId) != null ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ );
+
+ hasUnreleasedTasks = computeAssignmentDifference(
Review Comment:
nit: would it make sense to have a separate `boolean` for each type of unreleased tasks?
EDIT: I see how it's being used below so I get the one boolean approach, but I'm thinking would it make sense to alert the member to which tasks must be released
EDIT:EDIT: OK I see how you are doing exactly that with the `TaskTuple`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1926583303
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,451 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(int targetAssignmentEpoch,
+ TaskTuple targetAssignment) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the process ID of a subtopology ID /
+ * partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(BiFunction<String, Integer, String> currentActiveTaskProcessId) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the process IDs of a subtopology
+ * ID / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the process IDs of a subtopology ID
+ * / partition ID pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(TaskTuple ownedAssignment) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ hasUnreleasedTasks |= computeAssignmentDifferenceForOneSubtopology(
+ subtopologyId,
+ currentAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ targetAssignment.getOrDefault(subtopologyId, Collections.emptySet()),
+ resultAssignedTasks,
+ resultTasksPendingRevocation,
+ resultTasksPendingAssignment,
+ isUnreleasedTask
+ );
+ }
+ return hasUnreleasedTasks;
+ }
+
+ private static boolean computeAssignmentDifferenceForOneSubtopology(final String subtopologyId,
+ final Set<Integer> currentTasksForThisSubtopology,
+ final Set<Integer> targetTasksForThisSubtopology,
+ final Map<String, Set<Integer>> resultAssignedTasks,
+ final Map<String, Set<Integer>> resultTasksPendingRevocation,
+ final Map<String, Set<Integer>> resultTasksPendingAssignment,
+ final BiPredicate<String, Integer> isUnreleasedTask) {
+ // Result Assigned Tasks = Current Tasks ∩ Target Tasks
+ // i.e. we remove all tasks from the current assignment that are not in the target
+ // assignment
+ Set<Integer> resultAssignedTasksForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultAssignedTasksForThisSubtopology.retainAll(targetTasksForThisSubtopology);
+
+ // Result Tasks Pending Revocation = Current Tasks - Result Assigned Tasks
+ // i.e. we will ask the member to revoke all tasks in its current assignment that
+ // are not in the target assignment
+ Set<Integer> resultTasksPendingRevocationForThisSubtopology = new HashSet<>(currentTasksForThisSubtopology);
+ resultTasksPendingRevocationForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+
+ // Result Tasks Pending Assignment = Target Tasks - Result Assigned Tasks - Unreleased Tasks
+ // i.e. we will ask the member to assign all tasks in its target assignment,
+ // except those that are already assigned, and those that are unreleased
+ Set<Integer> resultTasksPendingAssignmentForThisSubtopology = new HashSet<>(targetTasksForThisSubtopology);
+ resultTasksPendingAssignmentForThisSubtopology.removeAll(resultAssignedTasksForThisSubtopology);
+ boolean hasUnreleasedTasks = resultTasksPendingAssignmentForThisSubtopology.removeIf(taskId ->
+ isUnreleasedTask.test(subtopologyId, taskId)
+ );
+
+ if (!resultAssignedTasksForThisSubtopology.isEmpty()) {
+ resultAssignedTasks.put(subtopologyId, resultAssignedTasksForThisSubtopology);
+ }
+
+ if (!resultTasksPendingRevocationForThisSubtopology.isEmpty()) {
+ resultTasksPendingRevocation.put(subtopologyId, resultTasksPendingRevocationForThisSubtopology);
+ }
+
+ if (!resultTasksPendingAssignmentForThisSubtopology.isEmpty()) {
+ resultTasksPendingAssignment.put(subtopologyId, resultTasksPendingAssignmentForThisSubtopology);
+ }
+
+ return hasUnreleasedTasks;
+ }
+
+ /**
+ * Computes the next assignment.
+ *
+ * @param memberEpoch The epoch of the member to use. This may be different from
+ * the epoch in {@link CurrentAssignmentBuilder#member}.
+ * @param memberAssignedTasks The assigned tasks of the member to use.
+ * @return A new StreamsGroupMember.
+ */
+ private StreamsGroupMember computeNextAssignment(int memberEpoch,
+ TaskTuple memberAssignedTasks) {
+ Map<String, Set<Integer>> newActiveAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newActiveTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newStandbyTasksPendingAssignment = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupAssignedTasks = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingRevocation = new HashMap<>();
+ Map<String, Set<Integer>> newWarmupTasksPendingAssignment = new HashMap<>();
+
+ boolean hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.activeTasks(),
+ targetAssignment.activeTasks(),
+ newActiveAssignedTasks,
+ newActiveTasksPendingRevocation,
+ newActiveTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ currentActiveTaskProcessId.apply(subtopologyId, partitionId) != null ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ );
+
+ hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.standbyTasks(),
+ targetAssignment.standbyTasks(),
+ newStandbyAssignedTasks,
+ newStandbyTasksPendingRevocation,
+ newStandbyTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ Objects.equals(currentActiveTaskProcessId.apply(subtopologyId, partitionId),
+ member.processId()) ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ ) || hasUnreleasedTasks;
+
+ hasUnreleasedTasks = computeAssignmentDifference(
+ memberAssignedTasks.warmupTasks(),
+ targetAssignment.warmupTasks(),
+ newWarmupAssignedTasks,
+ newWarmupTasksPendingRevocation,
+ newWarmupTasksPendingAssignment,
+ (subtopologyId, partitionId) ->
+ Objects.equals(currentActiveTaskProcessId.apply(subtopologyId, partitionId),
+ member.processId()) ||
+ currentStandbyTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId()) ||
+ currentWarmupTaskProcessIds.apply(subtopologyId, partitionId)
+ .contains(member.processId())
+ ) || hasUnreleasedTasks;
+
+ return buildNewMember(
+ memberEpoch,
+ new TaskTuple(
+ newActiveTasksPendingRevocation,
+ newStandbyTasksPendingRevocation,
+ newWarmupTasksPendingRevocation
+ ),
+ new TaskTuple(
+ newActiveAssignedTasks,
+ newStandbyAssignedTasks,
+ newWarmupAssignedTasks
+ ),
+ new TaskTuple(
+ newActiveTasksPendingAssignment,
+ newStandbyTasksPendingAssignment,
+ newWarmupTasksPendingAssignment
+ ),
+ hasUnreleasedTasks
+ );
+ }
+
+ private StreamsGroupMember buildNewMember(final int memberEpoch,
+ final TaskTuple newTasksPendingRevocation,
+ final TaskTuple newAssignedTasks,
+ final TaskTuple newTasksPendingAssignment,
+ final boolean hasUnreleasedTasks) {
+
+ final boolean hasTasksToBeRevoked =
+ (!newTasksPendingRevocation.isEmpty())
+ && (ownedTasks.isEmpty() || ownedTasks.get().containsAny(newTasksPendingRevocation));
+
+ if (hasTasksToBeRevoked) {
+ // If there are tasks to be revoked, the member remains in its current
+ // epoch and requests the revocation of those tasks. It transitions to
+ // the UNREVOKED_TASKS state to wait until the client acknowledges the
+ // revocation of the tasks.
+ return new StreamsGroupMember.Builder(member)
+ .setState(MemberState.UNREVOKED_TASKS)
+ .updateMemberEpoch(memberEpoch)
+ .setAssignedTasks(newAssignedTasks)
+ .setTasksPendingRevocation(newTasksPendingRevocation)
+ .build();
+ } else if (!newTasksPendingAssignment.isEmpty()) {
+ // If there are tasks to be assigned, the member transitions to the
+ // target epoch and requests the assignment of those tasks. Note that
+ // the tasks are directly added to the assigned tasks set. The
+ // member transitions to the STABLE state or to the UNRELEASED_TASKS
+ // state depending on whether there are unreleased tasks or not.
+ MemberState newState =
+ hasUnreleasedTasks
+ ? MemberState.UNRELEASED_TASKS
+ : MemberState.STABLE;
+ return new StreamsGroupMember.Builder(member)
+ .setState(newState)
+ .updateMemberEpoch(targetAssignmentEpoch)
+ .setAssignedTasks(newAssignedTasks.merge(newTasksPendingAssignment))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+ } else if (hasUnreleasedTasks) {
+ // If there are no tasks to be revoked nor to be assigned but some
+ // tasks are not available yet, the member transitions to the target
+ // epoch, to the UNRELEASED_TASKS state and waits.
+ return new StreamsGroupMember.Builder(member)
+ .setState(MemberState.UNRELEASED_TASKS)
Review Comment:
It's a bit different. `UNREVOKED_TASKS` is used if the member needs to revoke tasks. `UNRELEASED_TASKS` is used if the member waits for other members to revoke tasks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "cadonna (via GitHub)" <gi...@apache.org>.
cadonna commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1912971326
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
Review Comment:
What about simply `Tasks`?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
Review Comment:
Here and elsewhere:
```suggestion
public CurrentAssignmentBuilder withTargetAssignment(int targetAssignmentEpoch,
TaskTuple targetAssignment) {
```
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
+ Map<String, Set<Integer>> standbyTasks,
+ Map<String, Set<Integer>> warmupTasks) {
+
+ public TaskTuple {
+ activeTasks = Collections.unmodifiableMap(Objects.requireNonNull(activeTasks));
+ standbyTasks = Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks));
+ warmupTasks = Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks));
+ }
+
+ /**
+ * An empty task tuple.
+ */
+ public static final TaskTuple EMPTY = new TaskTuple(
+ Collections.emptyMap(),
+ Collections.emptyMap(),
+ Collections.emptyMap()
+ );
+
+ /**
+ * @return true if all collections in the tuple are empty.
+ */
+ public boolean isEmpty() {
+ return activeTasks.isEmpty() && standbyTasks.isEmpty() && warmupTasks.isEmpty();
+ }
+
+ /**
+ * Merges this task tuple with another task tuple.
+ *
+ * @param other The other task tuple.
+ * @return A new task tuple, containing all active tasks, standby tasks and warm-up tasks from both tuples.
+ */
+ public TaskTuple merge(TaskTuple other) {
+ Map<String, Set<Integer>> mergedActiveTasks = merge(activeTasks, other.activeTasks);
+ Map<String, Set<Integer>> mergedStandbyTasks = merge(standbyTasks, other.standbyTasks);
+ Map<String, Set<Integer>> mergedWarmupTasks = merge(warmupTasks, other.warmupTasks);
+ return new TaskTuple(mergedActiveTasks, mergedStandbyTasks, mergedWarmupTasks);
+ }
+
+ private Map<String, Set<Integer>> merge(final Map<String, Set<Integer>> tasks1, final Map<String, Set<Integer>> tasks2) {
Review Comment:
This could be static to emphasize that this object is not changed.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
+ Map<String, Set<Integer>> standbyTasks,
+ Map<String, Set<Integer>> warmupTasks) {
+
+ public TaskTuple {
+ activeTasks = Collections.unmodifiableMap(Objects.requireNonNull(activeTasks));
+ standbyTasks = Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks));
+ warmupTasks = Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks));
Review Comment:
Why making the maps unmodifiable but not the sets? To avoid the iterations?
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/TaskAssignmentTestUtil.java:
##########
@@ -17,37 +17,41 @@
package org.apache.kafka.coordinator.group.streams;
import java.util.AbstractMap;
-import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
-import java.util.Objects;
import java.util.Set;
public class TaskAssignmentTestUtil {
- public static Assignment mkAssignment(final Map<String, Set<Integer>> activeTasks,
- final Map<String, Set<Integer>> standbyTasks,
- final Map<String, Set<Integer>> warmupTasks) {
- return new Assignment(
- Collections.unmodifiableMap(Objects.requireNonNull(activeTasks)),
- Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks)),
- Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks))
- );
+ public enum TaskRole {
+ ACTIVE,
+ STANDBY,
+ WARMUP
+ }
+
+ @SafeVarargs
+ public static TaskTuple mkTaskTuple(TaskRole taskRole, Map.Entry<String, Set<Integer>>... entries) {
+ return switch (taskRole) {
+ case ACTIVE -> new TaskTuple(mkTasksPerSubtopology(entries), new HashMap<>(), new HashMap<>());
+ case STANDBY -> new TaskTuple(new HashMap<>(), mkTasksPerSubtopology(entries), new HashMap<>());
+ case WARMUP -> new TaskTuple(new HashMap<>(), new HashMap<>(), mkTasksPerSubtopology(entries));
+ };
}
- public static Map.Entry<String, Set<Integer>> mkTasks(String subtopologyId,
- Integer... tasks) {
+ public static Map.Entry<String, Set<Integer>> mkTasks(
+ String subtopologyId,
+ Integer... tasks
+ ) {
Review Comment:
Why this change? The style was better before.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
Review Comment:
Do we need to verify for the invariant that a member epoch is always <= target assignment epoch?
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
Review Comment:
This could also be a class constant that can be re-used in other parts.
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
Review Comment:
Those could be class constants.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
Review Comment:
Why is the case owned tasks data structure exists but is empty treated differently?
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
Review Comment:
Should the target assignment epoch here be formulated by means of the member epoch to document the invariant member epoch <= target assignment epoch?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(
+ Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask
+ ) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ Set<Integer> currentPartitions = currentAssignment.getOrDefault(subtopologyId, Collections.emptySet());
+ Set<Integer> targetPartitions = targetAssignment.getOrDefault(subtopologyId, Collections.emptySet());
+
+ // Result Assigned Partitions = Current Partitions ∩ Target Partitions
+ // i.e. we remove all partitions from the current assignment that are not in the target
+ // assignment
+ Set<Integer> resultAssignedPartitions = new HashSet<>(currentPartitions);
+ resultAssignedPartitions.retainAll(targetPartitions);
+
+ // Result Partitions Pending Revocation = Current Partitions - Result Assigned Partitions
+ // i.e. we will ask the member to revoke all partitions in its current assignment that
+ // are not in the target assignment
+ Set<Integer> resultPartitionsPendingRevocation = new HashSet<>(currentPartitions);
+ resultPartitionsPendingRevocation.removeAll(resultAssignedPartitions);
+
+ // Result Partitions Pending Assignment = Target Partitions - Result Assigned Partitions - Unreleased Partitions
+ // i.e. we will ask the member to assign all partitions in its target assignment,
+ // except those that are already assigned, and those that are unrelead
Review Comment:
```suggestion
// except those that are already assigned, and those that are unreleased
```
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
Review Comment:
Could it be that this member epoch + 1 is never used because in `buildNewMember()` the member epoch is set to the target assignment epoch in this case?
If yes, I believe we need to refactor a bit this code to make it more readable.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
Review Comment:
Also member ID does not sound right. Isn't it the process ID which is the ID of the Streams client that can contain more than one member?
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
Review Comment:
Should those values depend on previous values?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
+ throw new FencedMemberEpochException(
+ "The streams group member is in a unknown state. "
+ + "The member must abandon all its tasks and rejoin.");
+ }
+
+ return computeNextAssignment(
+ targetAssignmentEpoch,
+ member.assignedTasks()
+ );
+ }
+
+ return member;
+ }
+
+ /**
+ * Takes the current currentAssignment and the targetAssignment, and generates three
+ * collections:
+ *
+ * - the resultAssignedTasks: the tasks that are assigned in both the current and target
+ * assignments.
+ * - the resultTasksPendingRevocation: the tasks that are assigned in the current
+ * assignment but not in the target assignment.
+ * - the resultTasksPendingAssignment: the tasks that are assigned in the target assignment but
+ * not in the current assignment, and can be assigned currently (i.e., they are not owned by
+ * another member, as defined by the `isUnreleasedTask` predicate).
+ */
+ private boolean computeAssignmentDifference(
+ Map<String, Set<Integer>> currentAssignment,
+ Map<String, Set<Integer>> targetAssignment,
+ Map<String, Set<Integer>> resultAssignedTasks,
+ Map<String, Set<Integer>> resultTasksPendingRevocation,
+ Map<String, Set<Integer>> resultTasksPendingAssignment,
+ BiPredicate<String, Integer> isUnreleasedTask
+ ) {
+ boolean hasUnreleasedTasks = false;
+
+ Set<String> allSubtopologyIds = new HashSet<>(targetAssignment.keySet());
+ allSubtopologyIds.addAll(currentAssignment.keySet());
+
+ for (String subtopologyId : allSubtopologyIds) {
+ Set<Integer> currentPartitions = currentAssignment.getOrDefault(subtopologyId, Collections.emptySet());
Review Comment:
Could you please rename the variables to use `Tasks` instead of `Partitions`? With `Partitions` it is a bit confusing.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
Review Comment:
Here and elsewhere
```suggestion
* @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology ID /
```
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStableWithNewTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 4)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 3)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasksWithOwnedTasksNotHavingRevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, __) ->
+ subtopologyId2.equals(subtopologyId) ? "process" : null
+ )
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testRemainsInUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ CurrentAssignmentBuilder currentAssignmentBuilder = new CurrentAssignmentBuilder(
+ member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet());
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(null)
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5, 6)))
+ .build()
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 5)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
Review Comment:
Shouldn't this be non-empty? The member is in `UNREVOKED_TASKS`. It should have pending tasks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1922328205
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
Review Comment:
Good comment! Interesting idea, but I think you are missing something. See, whether `member.tasksPendingRevocation` was defined in a previous heartbeat, which may have had a different target assignment than the latest target assignment. So here, all we know is that we can transition to the next target epoch, because we revoked all tasks that we need to according to that previous target assignment. This doesn't imply much about what will go on in `buildNewMember`, because we are using a new target assignment, which may require other tasks to be revoked. For example, if we find out that we actually need to revoke more tasks to reach the latest target assignment, we will use the result of `member.memberEpoch() + 1` and transition to UNREVOKED_TASKS.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-18324: Add CurrentAssignmentBuilder [kafka]
Posted by "cadonna (via GitHub)" <gi...@apache.org>.
cadonna commented on code in PR #18476:
URL: https://github.com/apache/kafka/pull/18476#discussion_r1925612798
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
Review Comment:
I do not agree. IMO, distinguishing a `10` from an `11` is much harder than `memberEpoch` from `memberEpoch + 1` or `10` from `10 + 1` if you like that more.
That about reading. In my comment, I was more referring to documenting that we specifically test the case where the member epoch of the updated member is greater than the member epoch of the original member and not that the epoch are just different.
I just wanted to explain where I am coming from. I am fine if you want to keep it as it is.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
+ member.assignedTasks()
+ );
+
+ case UNRELEASED_TASKS:
+ // When the member is in the UNRELEASED_TASKS, we reconcile the
+ // member towards the latest target assignment. This will assign any
+ // of the unreleased tasks when they become available.
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+
+ case UNKNOWN:
+ // We could only end up in this state if a new state is added in the
+ // future and the group coordinator is downgraded. In this case, the
+ // best option is to fence the member to force it to rejoin the group
+ // without any tasks and to reconcile it again from scratch.
+ if ((ownedTasks.isEmpty() || !ownedTasks.get().isEmpty())) {
Review Comment:
Got it!
If I remember correctly that was also my understanding, but I wanted confirmation.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
Review Comment:
I tend to agree. I would call it `TasksTuple`, though, since each position within the tuple consists of multiple tasks. But I leave it to you if you want to change that.
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
Review Comment:
You could at least use variable names like `memberEpoch` and `nextMemberEpoch` or similar to make it clearer. Just an idea! Also this I leave it to you if you want to keep it.
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java:
##########
@@ -0,0 +1,744 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.util.Collections;
+
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTaskTuple;
+import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class CurrentAssignmentBuilderTest {
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member =
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(
+ mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(
+ taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToStableWithNewTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 4)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 3)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 4),
+ mkTasks(subtopologyId2, 3, 4, 7)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> "process")
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testStableToUnreleasedTasksWithOwnedTasksNotHavingRevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 4)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3, 5)))
+ .withCurrentActiveTaskProcessId((subtopologyId, __) ->
+ subtopologyId2.equals(subtopologyId) ? "process" : null
+ )
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNRELEASED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2),
+ mkTasks(subtopologyId2, 3)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToStable(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(11, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Collections.emptySet())
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.STABLE)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testRemainsInUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ CurrentAssignmentBuilder currentAssignmentBuilder = new CurrentAssignmentBuilder(
+ member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withCurrentStandbyTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet())
+ .withCurrentWarmupTaskProcessIds(
+ (subtopologyId, partitionId) -> Collections.emptySet());
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(null)
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build()
+ );
+
+ assertEquals(
+ member,
+ currentAssignmentBuilder
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 4, 5, 6)))
+ .build()
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnrevokedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 1),
+ mkTasks(subtopologyId2, 4)))
+ .build();
+
+ StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+ .withTargetAssignment(12, mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
+ .withOwnedAssignment(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .build();
+
+ assertEquals(
+ new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 3),
+ mkTasks(subtopologyId2, 6)))
+ .setTasksPendingRevocation(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2),
+ mkTasks(subtopologyId2, 5)))
+ .build(),
+ updatedMember
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(TaskRole.class)
+ public void testUnrevokedTasksToUnreleasedTasks(TaskRole taskRole) {
+ String subtopologyId1 = Uuid.randomUuid().toString();
+ String subtopologyId2 = Uuid.randomUuid().toString();
+
+ StreamsGroupMember member = new StreamsGroupMember.Builder("member")
+ .setState(MemberState.UNREVOKED_TASKS)
+ .setProcessId("process")
+ .setMemberEpoch(11)
+ .setPreviousMemberEpoch(10)
+ .setAssignedTasks(mkTaskTuple(taskRole,
+ mkTasks(subtopologyId1, 2, 3),
+ mkTasks(subtopologyId2, 5, 6)))
+ .setTasksPendingRevocation(TaskTuple.EMPTY)
Review Comment:
Did you forget to add this change? This line is still `.setTasksPendingRevocation(TaskTuple.EMPTY)`.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/TaskTuple.java:
##########
@@ -0,0 +1,135 @@
+/*
+ * 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;
+
+import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * An immutable tuple containing active, standby and warm-up tasks.
+ *
+ * @param activeTasks Active tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param standbyTasks Standby tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ * @param warmupTasks Warm-up tasks.
+ * The key of the map is the subtopology ID and the value is the set of partition IDs.
+ */
+public record TaskTuple(Map<String, Set<Integer>> activeTasks,
+ Map<String, Set<Integer>> standbyTasks,
+ Map<String, Set<Integer>> warmupTasks) {
+
+ public TaskTuple {
+ activeTasks = Collections.unmodifiableMap(Objects.requireNonNull(activeTasks));
+ standbyTasks = Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks));
+ warmupTasks = Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks));
Review Comment:
OK
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
Review Comment:
Good!
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java:
##########
@@ -0,0 +1,447 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.errors.FencedMemberEpochException;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+
+/**
+ * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the streams group protocol. Given the current state of a
+ * member and a desired or target assignment state, the state machine takes the necessary steps to converge them.
+ */
+public class CurrentAssignmentBuilder {
+
+ /**
+ * The streams group member which is reconciled.
+ */
+ private final StreamsGroupMember member;
+
+ /**
+ * The target assignment epoch.
+ */
+ private int targetAssignmentEpoch;
+
+ /**
+ * The target assignment.
+ */
+ private TaskTuple targetAssignment;
+
+ /**
+ * A function which returns the current process ID of an active task or null if the active task
+ * is not assigned. The current process ID is the process ID of the current owner.
+ */
+ private BiFunction<String, Integer, String> currentActiveTaskProcessId;
+
+ /**
+ * A function which returns the current process IDs of a standby task or null if the standby
+ * task is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds;
+
+ /**
+ * A function which returns the current process IDs of a warmup task or null if the warmup task
+ * is not assigned. The current process IDs are the process IDs of all current owners.
+ */
+ private BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds;
+
+ /**
+ * The tasks owned by the member. This may be provided by the member in the StreamsGroupHeartbeat request.
+ */
+ private Optional<TaskTuple> ownedTasks = Optional.empty();
+
+ /**
+ * Constructs the CurrentAssignmentBuilder based on the current state of the provided streams group member.
+ *
+ * @param member The streams group member that must be reconciled.
+ */
+ public CurrentAssignmentBuilder(StreamsGroupMember member) {
+ this.member = Objects.requireNonNull(member);
+ }
+
+ /**
+ * Sets the target assignment epoch and the target assignment that the streams group member must be reconciled to.
+ *
+ * @param targetAssignmentEpoch The target assignment epoch.
+ * @param targetAssignment The target assignment.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withTargetAssignment(
+ int targetAssignmentEpoch,
+ TaskTuple targetAssignment
+ ) {
+ this.targetAssignmentEpoch = targetAssignmentEpoch;
+ this.targetAssignment = Objects.requireNonNull(targetAssignment);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process ID of an active task. This is
+ * used by the state machine to determine if an active task is free or still used by another
+ * member, and if there is still a task on a specific process that is not yet revoked.
+ *
+ * @param currentActiveTaskProcessId A BiFunction which gets the memberId of a subtopology id /
+ * partition id pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentActiveTaskProcessId(
+ BiFunction<String, Integer, String> currentActiveTaskProcessId
+ ) {
+ this.currentActiveTaskProcessId = Objects.requireNonNull(currentActiveTaskProcessId);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a standby task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentStandbyTaskProcessIds A BiFunction which gets the memberIds of a subtopology
+ * ids / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentStandbyTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentStandbyTaskProcessIds
+ ) {
+ this.currentStandbyTaskProcessIds = Objects.requireNonNull(currentStandbyTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets a BiFunction which allows to retrieve the current process IDs of a warmup task. This is
+ * used by the state machine to determine if there is still a task on a specific process that is
+ * not yet revoked.
+ *
+ * @param currentWarmupTaskProcessIds A BiFunction which gets the memberIds of a subtopology ids
+ * / partition ids pair.
+ * @return This object.
+ */
+ public CurrentAssignmentBuilder withCurrentWarmupTaskProcessIds(
+ BiFunction<String, Integer, Set<String>> currentWarmupTaskProcessIds
+ ) {
+ this.currentWarmupTaskProcessIds = Objects.requireNonNull(currentWarmupTaskProcessIds);
+ return this;
+ }
+
+ /**
+ * Sets the tasks currently owned by the member. This comes directly from the last StreamsGroupHeartbeat request. This is used to
+ * determine if the member has revoked the necessary tasks. Passing null into this function means that the member did not provide
+ * its owned tasks in this heartbeat.
+ *
+ * @param ownedAssignment A collection of active, standby and warm-up tasks
+ * @return This object.
+ */
+ protected CurrentAssignmentBuilder withOwnedAssignment(
+ TaskTuple ownedAssignment
+ ) {
+ this.ownedTasks = Optional.ofNullable(ownedAssignment);
+ return this;
+ }
+
+ /**
+ * Builds the next state for the member or keep the current one if it is not possible to move forward with the current state.
+ *
+ * @return A new StreamsGroupMember or the current one.
+ */
+ public StreamsGroupMember build() {
+ switch (member.state()) {
+ case STABLE:
+ // When the member is in the STABLE state, we verify if a newer
+ // epoch (or target assignment) is available. If it is, we can
+ // reconcile the member towards it. Otherwise, we return.
+ if (member.memberEpoch() != targetAssignmentEpoch) {
+ return computeNextAssignment(
+ member.memberEpoch(),
+ member.assignedTasks()
+ );
+ } else {
+ return member;
+ }
+
+ case UNREVOKED_TASKS:
+ // When the member is in the UNREVOKED_TASKS state, we wait
+ // until the member has revoked the necessary tasks. They are
+ // considered revoked when they are not anymore reported in the
+ // owned tasks set in the StreamsGroupHeartbeat API.
+
+ // If the member provides its owned tasks, we verify if it still
+ // owns any of the revoked tasks. If it did not provide it's
+ // owned tasks, or we still own some of the revoked tasks, we
+ // cannot progress.
+ if (
+ ownedTasks.isEmpty() || ownedTasks.get().containsAny(member.tasksPendingRevocation())
+ ) {
+ return member;
+ }
+
+ // When the member has revoked all the pending tasks, it can
+ // transition to the next epoch (current + 1) and we can reconcile
+ // its state towards the latest target assignment.
+ return computeNextAssignment(
+ member.memberEpoch() + 1,
Review Comment:
Ah, that makes sense! I really missed something.
Thanks for the explanation!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org