lucasbru commented on code in PR #23588:
URL: https://github.com/apache/kafka/pull/23588#discussion_r4387877383
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -468,38 +485,59 @@ private static boolean hasUnfulfilledTaskQuota(
return process.memberToTaskCounts().get(member.memberId) <
localState.totalTasksPerMember;
}
+ /**
+ * The previous member to take a standby of {@code task}: the previous
active member, else the least-loaded previous
+ * standby member, each only while below the quota and on a process that
meets {@code prevMemberCondition}. Null if
+ * neither.
+ */
+ private static Member findPrevMemberForStandby(
+ final LocalState localState,
+ final TaskId task,
+ final Predicate<ProcessState> prevMemberCondition
+ ) {
+ // prev active task
+ final Member prevActiveMember =
localState.activeTaskToPrevMember.get(task);
+ if (prevActiveMember != null) {
+ final ProcessState prevActiveMemberProcessState =
localState.processIdToState.get(prevActiveMember.processId);
+ if (prevMemberCondition.test(prevActiveMemberProcessState) &&
!prevActiveMemberProcessState.hasTask(task) &&
hasUnfulfilledTaskQuota(localState, prevActiveMemberProcessState,
prevActiveMember)) {
+ return prevActiveMember;
+ }
+ }
+
+ // prev standby tasks
+ final Member prevStandbyMember =
findPrevMemberWithLeastLoad(localState,
localState.standbyTaskToPrevMember.get(task), Optional.of(task),
prevMemberCondition);
Review Comment:
If the prev standby member picked here is already at quota, this returns
null even when another prev standby member on a candidate process is still
below quota. `findPrevMemberWithLeastLoad` never looks at the member quota, and
only replaces its pick when both process and member load are strictly lower. In
step 1 this makes `placeRackAwareStandbys` return early, so the standby goes
cold in step 2.
Repro (fails with `expected: <[m2]> but was: <[m1b]>`):
```java
final Map<String, MemberMetadataAndStateImpl> members = mkMap(
mkEntry("a0", createTaggedMemberMetadata("P0", Map.of("rack", "A"),
Map.of("s", Set.of(0)), Map.of())),
mkEntry("m1", createTaggedMemberMetadataWithOffsets("P1", Map.of("rack",
"B"), Map.of("s", Set.of(1)), Map.of("s", Set.of(4, 0)), Map.of("s", Map.of(0,
9000L)))),
mkEntry("m1b", createTaggedMemberMetadata("P1", Map.of("rack", "B"),
Map.of("s", Set.of(2)), Map.of())),
mkEntry("m2", createTaggedMemberMetadataWithOffsets("P2", Map.of("rack",
"C"), Map.of("s", Set.of(3)), Map.of("s", Set.of(0)), Map.of("s", Map.of(0,
100L)))),
mkEntry("m2b", createTaggedMemberMetadata("P2", Map.of("rack", "C"),
Map.of("s", Set.of(4)), Map.of("s", Set.of(1))))
);
final GroupAssignment result = assignor.assign(
new GroupSpecImpl(members,
AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1).withRackAwareAssignmentTags(List.of("rack"))),
new TopologyDescriberImpl(5, true, List.of("s")));
assertEquals(Set.of("m2"), standbyHolders(result, "s", 0));
```
The tag-blind loop already had this before the PR, but there it just fell
through to least-loaded. Checking `hasUnfulfilledTaskQuota` inside
`findPrevMemberWithLeastLoad` (e.g. as part of the predicate) would fix both.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -132,17 +136,22 @@ private static LocalState initialize(final GroupSpec
groupSpec, final TopologyDe
localState.totalTasksPerMember =
computeTasksPerMember(localState.totalTasks,
localState.totalMembersWithTaskCapacity);
localState.processIdToState =
Utils.newHashMap(localState.totalMembersWithActiveTaskCapacity);
+ localState.processIdToClientTags =
Utils.newHashMap(localState.totalMembersWithActiveTaskCapacity);
localState.activeTaskToPrevMember =
Utils.newHashMap(localState.totalActiveTasks);
+ localState.statefulActiveTaskToProcess =
Utils.newHashMap(localState.totalStatefulActiveTasks);
// Standby-strength candidates per task, gathered in a single pass
over the members and ranked below.
final Map<TaskId, ArrayList<StandbyCandidate>> standbyCandidates = new
HashMap<>();
for (final String memberId : groupSpec.memberIds()) {
final MemberAssignmentState memberAssignmentState =
groupSpec.memberAssignmentState(memberId);
- final String processId =
groupSpec.memberMetadata(memberId).processId();
+ final MemberAssignmentMetadata memberMetadata =
groupSpec.memberMetadata(memberId);
+ final String processId = memberMetadata.processId();
final Member member = new Member(processId, memberId);
localState.processIdToState.computeIfAbsent(processId,
ProcessState::new)
.addMember(memberId);
+ // Client tags belong to the process, so every member of it
reports the same ones.
+ localState.processIdToClientTags.putIfAbsent(processId,
memberMetadata.clientTags());
Review Comment:
nit: this is an assumption rather than a guarantee. After an unclean
shutdown with changed `client.tag.*`, the old members (same processId from the
state dir) can stay in the group with the old tags until session timeout, and
then which tags win depends on `memberIds()` order. The classic assignor does
the same, so fine to keep, but maybe phrase the comment as an assumption.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -132,17 +136,22 @@ private static LocalState initialize(final GroupSpec
groupSpec, final TopologyDe
localState.totalTasksPerMember =
computeTasksPerMember(localState.totalTasks,
localState.totalMembersWithTaskCapacity);
localState.processIdToState =
Utils.newHashMap(localState.totalMembersWithActiveTaskCapacity);
+ localState.processIdToClientTags =
Utils.newHashMap(localState.totalMembersWithActiveTaskCapacity);
localState.activeTaskToPrevMember =
Utils.newHashMap(localState.totalActiveTasks);
+ localState.statefulActiveTaskToProcess =
Utils.newHashMap(localState.totalStatefulActiveTasks);
Review Comment:
Minor: this and `processIdToClientTags` are only read on the rack-aware
path. Could skip filling them when there are no tags or no standbys.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -433,7 +450,7 @@ private static Member findPrevMemberWithLeastLoad(
for (final Member member : members) {
final ProcessState processState =
localState.processIdToState.get(member.processId);
// A process that already owns a standby task (either as active or
standby) cannot take it again
- if (standbyTaskId.isPresent() &&
processState.hasTask(standbyTaskId.get())) {
+ if (standbyTaskId.isPresent() &&
processState.hasTask(standbyTaskId.get()) ||
!prevMemberCondition.test(processState)) {
Review Comment:
nit: could you add parentheses here, or split it into two `continue` checks?
It's correct by precedence but easy to misread.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -468,38 +485,59 @@ private static boolean hasUnfulfilledTaskQuota(
return process.memberToTaskCounts().get(member.memberId) <
localState.totalTasksPerMember;
}
+ /**
+ * The previous member to take a standby of {@code task}: the previous
active member, else the least-loaded previous
+ * standby member, each only while below the quota and on a process that
meets {@code prevMemberCondition}. Null if
+ * neither.
+ */
+ private static Member findPrevMemberForStandby(
+ final LocalState localState,
+ final TaskId task,
+ final Predicate<ProcessState> prevMemberCondition
+ ) {
+ // prev active task
+ final Member prevActiveMember =
localState.activeTaskToPrevMember.get(task);
+ if (prevActiveMember != null) {
+ final ProcessState prevActiveMemberProcessState =
localState.processIdToState.get(prevActiveMember.processId);
+ if (prevMemberCondition.test(prevActiveMemberProcessState) &&
!prevActiveMemberProcessState.hasTask(task) &&
hasUnfulfilledTaskQuota(localState, prevActiveMemberProcessState,
prevActiveMember)) {
+ return prevActiveMember;
+ }
+ }
+
+ // prev standby tasks
+ final Member prevStandbyMember =
findPrevMemberWithLeastLoad(localState,
localState.standbyTaskToPrevMember.get(task), Optional.of(task),
prevMemberCondition);
+ if (prevStandbyMember != null && hasUnfulfilledTaskQuota(localState,
localState.processIdToState.get(prevStandbyMember.processId),
prevStandbyMember)) {
+ return prevStandbyMember;
+ }
+ return null;
+ }
+
+ /**
+ * Assigns the standby tasks. Rack diversity ranks above stickiness: when
{@code rack.aware.assignment.tags} is
+ * set, steps 1 and 2 place every standby that still makes its task more
rack-diverse, using stickiness only to
+ * break ties among the most diverse processes; steps 3 and 4 are the
tag-blind assignment of the rest. Each pair
+ * runs its sticky step for all tasks before its least-loaded step, so
that a pick that leaves its previous members
+ * cannot take the room a later task needs to stay with one of its own.
Steps 3 and 4 run after step 2, so a
Review Comment:
Here's a case where this trade-off costs a standby move without buying any
diversity: zone tag, 1 standby, 4 tasks, A/B/C in z1 and D alone in z2 (A:
active 1, standby 0; B: active 0; C: active 3; D: active 2). Task 3's standby
fills D. Task 2's standby can go to A, B or C at equal load and takes A, which
fills it. Task 0's prev standby on A is lost and it ends up cold on C, although
B or C would have served task 2 just as well.
Would it be worth breaking step-2 ties away from processes that a
not-yet-placed prev standby holder still needs? Not a blocker, since diversity
is never lost.
##########
jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmark.java:
##########
@@ -120,6 +145,9 @@ public enum ReportedOffsets {
@Param({"NONE", "OWNED_AND_DORMANT"})
private ReportedOffsets reportedOffsets;
+ @Param({"NONE", "CLUSTER_AND_ZONE", "CLUSTER_ZONE_AND_RACK",
"CLUSTER_ZONE_AND_HOST", "REGION_CLUSTER_ZONE_RACK_AND_HOST"})
Review Comment:
This grows the matrix from 256 to 1920 combinations, and roughly half of
them measure the same thing: BALANCED never reads the tags, and STICKY with
`standbyReplicas=0` skips the standby path. Maybe add a note like the one for
`reportedOffsets`, or narrow the defaults?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -525,6 +564,100 @@ private static void assignStandby(final LocalState
localState, final LinkedList<
}
}
+ /**
+ * Steps 1 and 2 of {@link #assignStandby}: places the standbys that make
their task more rack-diverse and returns
+ * the processes they went to per task.
+ */
+ private static Map<TaskId, List<ProcessState>>
assignRackAwareStandbys(final LocalState localState, final List<TaskId>
standbyTasks) {
+ // A process has room while one of its members is below the per-member
quota. Loads and room change only when a
+ // standby is added below, and each change is reported to the picker.
+ final RackAwareStandbyPicker<ProcessState> picker = new
RackAwareStandbyPicker<>(
+ localState.rackAwareAssignmentTags,
+ localState.processIdToState.values(),
+ process ->
localState.processIdToClientTags.get(process.processId()),
+ ProcessState::load,
+ process -> hasRoom(localState, process)
+ );
+ final Map<TaskId, List<ProcessState>> rackAwareStandbys = new
HashMap<>();
+
+ // 1. re-assigning standby tasks in a new rack to clients that
previously had the same task (as active or standby)
+ for (final TaskId task : standbyTasks) {
+ final List<ProcessState> standbys = new ArrayList<>();
+ rackAwareStandbys.put(task, standbys);
+ placeRackAwareStandbys(localState, picker, task, standbys, true);
+ }
+
+ // 2. assigning remaining standby tasks in a new rack to the least
loaded client
+ for (final TaskId task : standbyTasks) {
+ placeRackAwareStandbys(localState, picker, task,
rackAwareStandbys.get(task), false);
+ }
+ return rackAwareStandbys;
+ }
+
+ /**
+ * Places standbys of {@code task} while they make it more rack-diverse
and adds their processes to {@code standbys},
+ * which holds those placed by an earlier call. Among the candidates of a
pick, a standby goes to the previous active
+ * member, else the least-loaded previous standby member, each while below
the quota, else to the least-loaded
+ * candidate. With {@code stickyOnly}, stops at the first standby with no
previous member to go to.
+ */
+ private static void placeRackAwareStandbys(
+ final LocalState localState,
+ final RackAwareStandbyPicker<ProcessState> picker,
+ final TaskId task,
+ final List<ProcessState> standbys,
+ final boolean stickyOnly
+ ) {
+ final List<ProcessState> holders = new ArrayList<>(1 +
standbys.size());
+ holders.add(localState.statefulActiveTaskToProcess.get(task));
+ holders.addAll(standbys);
+ picker.startTask(holders);
+
+ while (standbys.size() < localState.numStandbyReplicas) {
+ final RackAwareStandbyPicker.Pick<ProcessState> pick =
picker.pick();
+ if (pick == null) {
+ return;
+ }
+
+ // prev active or standby member on a candidate process, else the
least-loaded candidate
+ final Member prevMember = findPrevMemberForStandby(localState,
task, candidate -> picker.isCandidate(candidate, pick.conditions()));
Review Comment:
nit: none of the new tests covers a sticky pick where the prev member is at
quota but another member of its process still has room, so `isCandidate` passes
but `hasUnfulfilledTaskQuota` fails. Would be good to have one, given the issue
in `findPrevMemberForStandby`.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
