lucasbru commented on code in PR #23325: URL: https://github.com/apache/kafka/pull/23325#discussion_r3934550117
########## group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerImpl.java: ########## @@ -0,0 +1,506 @@ +/* + * 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.streams.assignor.TaskId; +import org.apache.kafka.coordinator.group.streams.topics.ConfiguredSubtopology; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.SortedMap; +import java.util.TreeMap; +import java.util.function.Consumer; + +/** + * The {@link AssignmentRefiner} being built out to replace {@link NoOpAssignmentRefiner} as the broker's default + * once the derivation is complete. + * + * <p>{@link #refine} is still a stub -- it returns the target assignment unchanged, exactly like + * {@link NoOpAssignmentRefiner} -- while the derivation is built out incrementally across several changes. The + * methods below are its building blocks: indexing the current assignment, and deciding which migrations can + * complete immediately versus which have to stage behind a warm-up task. None of them are called from + * {@link #refine} yet. + */ +public class AssignmentRefinerImpl implements AssignmentRefiner { + + @Override + public Map<String, TasksTuple> refine( + Map<String, StreamsGroupMember> members, + Map<String, TasksTuple> targetAssignment, + Map<String, MemberTaskOffsets> taskOffsets, + SortedMap<String, ConfiguredSubtopology> subtopologies, + int numWarmupReplicas, + long acceptableRecoveryLag + ) { + return targetAssignment; + } + + /** + * Indexes the members' current assignment by task, so that the case analysis can look up what a task's situation is + * without scanning the group again for every task. This is a single pass over the members' task entries. + * + * <p>Only stateful tasks are indexed. A stateless task has no state to restore, so it is never staged and never + * consulted here; it simply flows through from the target assignment. + * + * @param members + * All members of the group. + * @param taskOffsets + * The latest changelog offsets/end-offsets reported by the members. + * @param subtopologies + * The resolved subtopologies, which tell whether a subtopology is stateful. + * @param acceptableRecoveryLag + * The lag at or below which a replica counts as caught up. + * + * @return The current assignment, indexed by task. + */ + static CurrentAssignmentIndex indexCurrentAssignment( + final Map<String, StreamsGroupMember> members, + final Map<String, MemberTaskOffsets> taskOffsets, + final SortedMap<String, ConfiguredSubtopology> subtopologies, + final long acceptableRecoveryLag + ) { + final Map<TaskId, String> activeOwner = new HashMap<>(); + final Map<TaskId, Set<String>> processesRevoking = new HashMap<>(); + final Map<TaskId, List<TaskCopy>> taskCopies = new HashMap<>(); + + for (final StreamsGroupMember member : members.values()) { + final MemberTaskOffsets offsets = taskOffsets.getOrDefault(member.memberId(), MemberTaskOffsets.EMPTY); + + forEachStatefulActiveTask( + member.assignedTasks().activeTasksWithEpochs(), + subtopologies, + task -> activeOwner.put(task, member.memberId()) + ); + + // The member has been told to give this task up and has stopped running it, so the member is + // deliberately not recorded as the task's active owner. Recording the member as the owner would make the + // case analysis try to keep the task there, undoing a hand-over that is already under way. + // + // The task does still occupy the member's process until the revocation completes, and that is what stops + // a standby of the same task being placed there. The block applies to the whole process, not just this + // one member, so the process is what gets recorded. + forEachStatefulActiveTask( + member.tasksPendingRevocation().activeTasksWithEpochs(), + subtopologies, + task -> processesRevoking.computeIfAbsent(task, __ -> new HashSet<>()).add(member.processId()) + ); + + forEachStatefulTask( + member.assignedTasks().standbyTasks(), + subtopologies, + task -> addTaskCopy(taskCopies, task, member, TaskRole.STANDBY, offsets, acceptableRecoveryLag) + ); + + forEachStatefulTask( + member.assignedTasks().warmupTasks(), + subtopologies, + task -> addTaskCopy(taskCopies, task, member, TaskRole.WARMUP, offsets, acceptableRecoveryLag) + ); + } + + return new CurrentAssignmentIndex(activeOwner, processesRevoking, taskCopies); + } + + /** + * Decides, for every stateful task whose active role is not already where the target assignment wants it, whether + * the migration has to be staged behind a warm-up task or can be completed in this step. + * + * <p>A task is <b>staged</b> only when all of the following hold: somebody runs it today, the target owner is + * somebody else, and there is a warming improvement left to achieve. Everything else completes right away, which + * needs no patch at all -- the target assignment already places the task on its target owner, and the previous + * owner's slice already omits it. That is why the two outcomes are so lopsided: staging is the exception, and the + * result is proportional to how far the current assignment has diverged from the target rather than to the group's + * size. + * + * @param currentAssignment + * The indexed current assignment, from {@link #indexCurrentAssignment}. + * @param targetAssignment + * All members' target assignments, as computed by the task assignor. + * @param members + * All members of the group, used to resolve which process a member runs in. + * @param subtopologies + * The resolved subtopologies, which tell whether a subtopology is stateful. + * + * @return What was decided for the tasks that are not already in place. + */ + static TaskDecisions analyzeTasks( + final CurrentAssignmentIndex currentAssignment, + final Map<String, TasksTuple> targetAssignment, + final Map<String, StreamsGroupMember> members, + final SortedMap<String, ConfiguredSubtopology> subtopologies + ) { + // Only a task the target assignment still contains needs a decision, so its owners alone drive the loop. A + // task the target assignment dropped -- after a topology change, say -- belongs in nobody's slice, and its + // current holders revoke it the ordinary way, so walking the current assignment's tasks too would only turn + // up tasks to skip. + // + // The map is sorted, which gives the canonical task order that makes a derivation reproducible and leaves the + // budget pass that follows a deterministic tie-break to fall back on. + final SortedMap<TaskId, String> targetOwners = statefulActiveOwners(targetAssignment, subtopologies); + + final List<StagedMigration> stagedMigrations = new ArrayList<>(); + final List<TaskGrant> grantedTasks = new ArrayList<>(); + + for (final Map.Entry<TaskId, String> targetOwnerByTask : targetOwners.entrySet()) { + final TaskId task = targetOwnerByTask.getKey(); + final String targetOwner = targetOwnerByTask.getValue(); + + final String currentOwner = currentAssignment.activeOwner().get(task); + if (targetOwner.equals(currentOwner)) { + // The task already runs where it belongs. + continue; + } + + // The target assignment can still name a member the group has already removed: it is only recomputed when + // the assignor runs again, which the assignment interval can defer, and a member can be fenced in the + // meantime. Such a member cannot restore anything, so nothing may be staged into it. + final StreamsGroupMember targetMember = members.get(targetOwner); + if (targetMember == null) { + if (currentOwner != null) { + // Keep the task where it runs until the assignor names a member that still exists. Leaving it out + // instead would make its current owner revoke it, so it would stop being processed for no gain. + stagedMigrations.add(new StagedMigration( + task, + currentOwner, + targetOwner, + Optional.empty(), + Optional.empty() + )); + } + // A task that nobody runs and whose target owner is gone is left to the next assignor run: granting it + // to a member that is no longer in the group would achieve nothing. + continue; + } + + final String targetProcessId = targetMember.processId(); + + // The task moves now, for either of two reasons. Nobody runs it -- it is new, its owner left, or a + // hand-over is in flight and the previous owner has already released it -- so there is no running task to + // protect and the target owner takes it even cold; preferring a warmer owner would be a placement + // decision, and placement is the assignor's job. Or somebody runs it but no achievable warming improvement + // remains. The current owner, when there is one, is necessarily a member of the group, because the index it + // comes from was built from the members themselves. + if (currentOwner == null + || isReady(currentAssignment, task, members.get(currentOwner).processId(), targetProcessId)) { + grantedTasks.add(new TaskGrant(task, targetOwner)); + continue; + } + + stagedMigrations.add(new StagedMigration( + task, + currentOwner, + targetOwner, + Optional.of(targetProcessId), + findCopyOnProcess(currentAssignment, task, targetProcessId) + )); + } + + return new TaskDecisions(List.copyOf(stagedMigrations), List.copyOf(grantedTasks)); + } + + /** + * Whether no achievable warming improvement remains for handing the task over to its target owner. + * + * <p>This is deliberately weaker than "the target owner has caught up". Warming up is only worth staging when it + * can actually shorten the hand-over, and there are two situations where it cannot -- one per clause of the + * predicate: + * <ul> + * <li><b>The task is moving between two members of one process.</b> A process must not hold the same task + * twice, so there is no way to warm the target owner up while the current owner still runs it. No condition + * on the state applies here: staging such a move would park it forever, so it has to count as ready.</li> + * <li><b>The target owner's process already holds a caught-up copy of the task.</b> Here <em>caught up</em> + * carries the weight: a copy that is still catching up leaves a genuine improvement to wait for, so the task + * stays staged, that copy keeps consuming, and a later step grants the task once the copy is hot. Planting a + * warm-up task on the target owner is no help either way, because its process would then hold the task + * twice. The copy sits either on the target owner itself, which promotes it in place, or on a sibling + * member, which has to release the task first so that the target owner can reopen it from the state + * directory.</li> + * </ul> + * + * <p>The two are not variants of one another, even though both turn on members sharing a process. The first is + * about the <em>current owner</em> sharing one with the target owner; the second about some <em>copy holder</em> + * doing so. They also cannot both apply: if the current owner is on the target owner's process then that process + * runs the task, so by one-task-per-process it holds no copy of the task for the second clause to find. The + * second therefore only ever decides a move that crosses process boundaries, and the order the two are tested in + * makes no difference to the outcome. + * + * <p>Only the in-place promotion is warm for every store type. The other two paths -- the move within one + * process, and the sibling releasing the task -- are warm only for a store that persists to disk, where the + * releasing member's clean close leaves a checkpoint behind for the incoming member to reopen from. <b>An + * in-memory store is rebuilt from the changelog in full:</b> its state lives on the releasing member's heap and + * is dropped when the task closes, and no hand-over of a running task between threads of one process exists to + * carry it across. Worse, the lag that made the task look ready was measured on the member that is about to + * close, so for an in-memory store it says nothing about what the incoming member then has to restore. This + * predicate cannot fix that; it would take a client-side cross-thread task hand-over. The broker cannot even see + * the difference, because the topology metadata carries changelog topics but not how a store is backed. + * + * <p>What bounds the damage is that a warm-up task the refiner plants always targets the target owner itself, so + * every migration the refiner stages resolves through the in-place promotion. The other paths arise only out of a + * layout the refiner inherited. + */ + private static boolean isReady( + final CurrentAssignmentIndex currentAssignment, + final TaskId task, + final String currentProcessId, + final String targetProcessId + ) { + if (targetProcessId.equals(currentProcessId)) { + return true; + } + return currentAssignment.taskCopies().getOrDefault(task, List.of()).stream() + .anyMatch(holder -> holder.processId().equals(targetProcessId) && holder.caughtUp()); + } + + /** + * Whether the member has restored the task closely enough to take it over as an active task. Mirrors the client's + * own predicate, so that both ends agree on when a warm-up task is caught up. + * + * <p>The lag is the distance between the reported end offset and the reported offset, and a lag that is not known + * is never within the threshold: an offset missing on either side, or capped at {@link Long#MAX_VALUE} to say that + * the restore has not started, counts as not caught up. A slightly negative lag does count, because the offset is a + * position while the end offset is the last offset, so a fully restored task reports a lag of -1. + * + * @param memberTaskOffsets + * The offsets the member reported, {@link MemberTaskOffsets#EMPTY} if it reported none. + * @param task + * The task to check. + * @param acceptableRecoveryLag + * The lag at or below which the task counts as caught up. + */ + static boolean isCaughtUp( + final MemberTaskOffsets memberTaskOffsets, + final TaskId task, + final long acceptableRecoveryLag + ) { + final Long offset = offsetOf(memberTaskOffsets.taskOffsets(), task); + final Long endOffset = offsetOf(memberTaskOffsets.taskEndOffsets(), task); + if (offset == null || endOffset == null || offset == Long.MAX_VALUE || endOffset == Long.MAX_VALUE) { + return false; + } + return endOffset - offset <= acceptableRecoveryLag; + } + + private static Long offsetOf(final Map<String, Map<Integer, Long>> offsets, final TaskId task) { + final Map<Integer, Long> byPartition = offsets.get(task.subtopologyId()); + return byPartition == null ? null : byPartition.get(task.partition()); + } + + /** + * The replica of the task that the given process already holds, if any. + * + * <p>There is at most one, so no tie-break between roles is needed: a process holds a given task in at most one + * role, on at most one of its members. The reconciler enforces that -- {@code isUnreleasedActiveTask}, + * {@code isUnreleasedStandbyTask} and {@code isUnreleasedWarmupTask} in {@link CurrentAssignmentBuilder} each + * block a role for as long as the process holds the task in any role. + */ + private static Optional<TaskCopy> findCopyOnProcess( + final CurrentAssignmentIndex currentAssignment, + final TaskId task, + final String processId + ) { + return currentAssignment.taskCopies().getOrDefault(task, List.of()).stream() + .filter(holder -> holder.processId().equals(processId)) + .findFirst(); + } + + /** + * Inverts the target assignment into a lookup from stateful task to the member that is to run it as an active task. + */ + private static SortedMap<TaskId, String> statefulActiveOwners( + final Map<String, TasksTuple> targetAssignment, + final SortedMap<String, ConfiguredSubtopology> subtopologies + ) { + final SortedMap<TaskId, String> owners = new TreeMap<>(); + targetAssignment.forEach((memberId, tasks) -> + forEachStatefulTask(tasks.activeTasks(), subtopologies, task -> owners.put(task, memberId))); + return owners; + } + + private static void addTaskCopy( + final Map<TaskId, List<TaskCopy>> taskCopies, + final TaskId task, + final StreamsGroupMember member, + final TaskRole role, + final MemberTaskOffsets offsets, + final long acceptableRecoveryLag + ) { + taskCopies.computeIfAbsent(task, __ -> new ArrayList<>()).add(new TaskCopy( + member.memberId(), + member.processId(), + role, + isCaughtUp(offsets, task, acceptableRecoveryLag) + )); + } + + private static void forEachStatefulActiveTask( + final Map<String, Map<Integer, Integer>> activeTasksWithEpochs, + final SortedMap<String, ConfiguredSubtopology> subtopologies, + final Consumer<TaskId> action + ) { + activeTasksWithEpochs.forEach((subtopologyId, partitionsWithEpochs) -> { + if (isStateful(subtopologies, subtopologyId)) { + partitionsWithEpochs.keySet() + .forEach(partitionId -> action.accept(new TaskId(subtopologyId, partitionId))); + } + }); + } + + private static void forEachStatefulTask( + final Map<String, Set<Integer>> tasks, + final SortedMap<String, ConfiguredSubtopology> subtopologies, + final Consumer<TaskId> action + ) { + tasks.forEach((subtopologyId, partitionIds) -> { + if (isStateful(subtopologies, subtopologyId)) { + partitionIds.forEach(partitionId -> action.accept(new TaskId(subtopologyId, partitionId))); + } + }); + } + + private static boolean isStateful( Review Comment: this is a bit of misnomer because you can have stateful tasks that are not bakced by a log. So I sometimes wonder if it should be "LoggedStateful" or something -- 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]
