From c819c6e8daeb52cbe8f31200b89e0f4f0064fe01 Mon Sep 17 00:00:00 2001 From: Gabriella Fu <2234862823@qq.com> Date: Wed, 23 Sep 2026 15:40:42 -0400 Subject: [PATCH] add testbed for steams assignor --- .../assignor/AssignmentInvariants.java | 104 +++ .../streams/assignor/AssignmentMetrics.java | 457 ++++++++++ .../assignor/StickyTaskAssignorFuzzTest.java | 99 +++ .../streams/assignor/TaskAssignorTestbed.java | 781 ++++++++++++++++++ 4 files changed, 1441 insertions(+) create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentInvariants.java create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentMetrics.java create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorFuzzTest.java create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/TaskAssignorTestbed.java diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentInvariants.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentInvariants.java new file mode 100644 index 0000000000000..857e1fed736ea --- /dev/null +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentInvariants.java @@ -0,0 +1,104 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.coordinator.group.streams.assignor; + +import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.api.streams.assignor.MemberAssignment; +import org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.Scenario; +import org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.Topology; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.toTaskIds; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The properties every valid assignment must satisfy, independent of the assignor: the assignment covers exactly + * the group members, every active task has exactly one owner, no task has more than {@code numStandbyReplicas} + * standbys, stateless tasks have no standby, and no process holds two copies of the same task. + */ +final class AssignmentInvariants { + + private AssignmentInvariants() { + } + + static void assertValid(final Scenario scenario, final GroupAssignment result) { + assertMembersMatch(scenario, result); + + final Map> activeOwners = new HashMap<>(); + final Map> standbyOwners = new HashMap<>(); + collectOwners(scenario, result, activeOwners, standbyOwners); + + assertEachActiveTaskOwnedOnce(scenario.topology, activeOwners); + assertStandbyBound(scenario, standbyOwners); + } + + private static void assertMembersMatch(final Scenario scenario, final GroupAssignment result) { + assertEquals(scenario.memberIds(), result.members().keySet(), "assignment must cover exactly the group members"); + } + + /** + * Fills the owner maps, asserting on the way that every task is known, that standbys are only assigned for + * stateful tasks and that no process holds a task twice. + */ + private static void collectOwners( + final Scenario scenario, + final GroupAssignment result, + final Map> activeOwners, + final Map> standbyOwners + ) { + final Topology topology = scenario.topology; + final Map> tasksPerProcess = new HashMap<>(); + for (final Map.Entry entry : result.members().entrySet()) { + final String memberId = entry.getKey(); + final String processId = scenario.processOf(memberId); + final Set processTasks = tasksPerProcess.computeIfAbsent(processId, id -> new HashSet<>()); + for (final TaskId task : toTaskIds(entry.getValue().activeTasks())) { + assertTrue(topology.tasks().contains(task), "unknown active task " + task + " on " + memberId); + activeOwners.computeIfAbsent(task, t -> new ArrayList<>()).add(memberId); + assertTrue(processTasks.add(task), "process " + processId + " holds task " + task + " twice"); + } + for (final TaskId task : toTaskIds(entry.getValue().standbyTasks())) { + assertTrue(topology.statefulTasks().contains(task), "standby for stateless or unknown task " + task + " on " + memberId); + standbyOwners.computeIfAbsent(task, t -> new ArrayList<>()).add(memberId); + assertTrue(processTasks.add(task), "process " + processId + " holds task " + task + " twice"); + } + } + } + + private static void assertEachActiveTaskOwnedOnce(final Topology topology, final Map> activeOwners) { + for (final TaskId task : topology.tasks()) { + final List owners = activeOwners.getOrDefault(task, List.of()); + assertEquals(1, owners.size(), "active task " + task + " must have exactly one owner but has " + owners); + } + } + + private static void assertStandbyBound(final Scenario scenario, final Map> standbyOwners) { + standbyOwners.forEach((task, owners) -> + assertTrue( + owners.size() <= scenario.numStandbyReplicas, + "task " + task + " has " + owners.size() + " standbys, above the configured " + scenario.numStandbyReplicas + ": " + owners + ) + ); + } +} diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentMetrics.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentMetrics.java new file mode 100644 index 0000000000000..687894d9147e2 --- /dev/null +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/AssignmentMetrics.java @@ -0,0 +1,457 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.coordinator.group.streams.assignor; + +import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.api.streams.assignor.MemberAssignment; +import org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.ProcessSpec; +import org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.Scenario; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.OptionalDouble; +import java.util.OptionalInt; +import java.util.Set; +import java.util.TreeMap; + +import static org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.toTaskIds; + +/** + * Grades the converged assignment of a rebalance: balance across members and processes, stickiness and task + * movement relative to the assignment held before the rebalance, and rack diversity of the copies of each stateful + * task. Nothing here is asserted; the {@link Summary} over a run is printed so two versions of an assignor can be + * compared on the same scenarios. + */ +final class AssignmentMetrics { + + private AssignmentMetrics() { + } + + /** + * @param activeSpreadPerMember active tasks per member, max - min + * @param statefulActiveSpreadPerMember stateful active tasks per member, max - min + * @param subtopologyExcessSpread active tasks of one subtopology per member, max - min above the unavoidable spread, worst subtopology + * @param totalSpreadPerMember active + standby tasks per member, max - min + * @param processLoadSpread (active + standby) / members per process, max - min + * @param standbyLoadSpread standby tasks / members per process, max - min + * @param processStickiness fraction of the active tasks held before the rebalance that stayed on their process + * @param memberStickiness fraction of the active tasks held before the rebalance that stayed on their member, if it is still in the group + * @param stateMoves stateful task copies placed on a process that held no copy before + * @param stateReuse fraction of the tasks reported from state directories that were placed on that process + * @param copiesLost stateful tasks whose every copy sat on a process the event removed, for multi-process failures + * @param diversityPerTag rack diversity per configured tag, relative to the best spread of that tag alone + */ + record Metrics( + int activeSpreadPerMember, + int statefulActiveSpreadPerMember, + int subtopologyExcessSpread, + int totalSpreadPerMember, + double processLoadSpread, + double standbyLoadSpread, + OptionalDouble processStickiness, + OptionalDouble memberStickiness, + OptionalInt stateMoves, + OptionalDouble stateReuse, + OptionalInt copiesLost, + Map diversityPerTag + ) { + } + + /** + * @param before what the members held before the event that triggered the rebalance + * @param restoredTasks the tasks each process reported from its state directories when the rebalance started + */ + static Metrics compute( + final Scenario scenario, + final GroupAssignment result, + final Scenario.Baseline before, + final Map> restoredTasks + ) { + final Map activePerMember = new HashMap<>(); + final Map statefulActivePerMember = new HashMap<>(); + final Map> activePerSubtopologyPerMember = new HashMap<>(); + final Map totalPerMember = new HashMap<>(); + final Map standbyPerProcess = new HashMap<>(); + final Map activeMemberOf = new HashMap<>(); + final Map activeProcessOf = new HashMap<>(); + final Map> ownerProcessesOf = new HashMap<>(); + for (final String processId : scenario.processes.keySet()) { + standbyPerProcess.put(processId, 0); + } + for (final Map.Entry entry : result.members().entrySet()) { + final String memberId = entry.getKey(); + final String processId = scenario.processOf(memberId); + final Set active = toTaskIds(entry.getValue().activeTasks()); + final Set standby = toTaskIds(entry.getValue().standbyTasks()); + int stateful = 0; + for (final TaskId task : active) { + if (scenario.topology.statefulTasks().contains(task)) { + stateful++; + } + activeMemberOf.put(task, memberId); + activeProcessOf.put(task, processId); + ownerProcessesOf.computeIfAbsent(task, t -> new HashSet<>()).add(processId); + activePerSubtopologyPerMember.computeIfAbsent(task.subtopologyId(), s -> new HashMap<>()).merge(memberId, 1, Integer::sum); + } + for (final TaskId task : standby) { + ownerProcessesOf.computeIfAbsent(task, t -> new HashSet<>()).add(processId); + } + activePerMember.put(memberId, active.size()); + statefulActivePerMember.put(memberId, stateful); + totalPerMember.put(memberId, active.size() + standby.size()); + standbyPerProcess.merge(processId, standby.size(), Integer::sum); + } + + double maxLoad = 0; + double minLoad = Double.MAX_VALUE; + double maxStandbyLoad = 0; + double minStandbyLoad = Double.MAX_VALUE; + for (final Map.Entry entry : scenario.processes.entrySet()) { + final ProcessSpec process = entry.getValue(); + int tasks = 0; + for (final String memberId : process.members) { + tasks += totalPerMember.get(memberId); + } + final double load = (double) tasks / process.members.size(); + maxLoad = Math.max(maxLoad, load); + minLoad = Math.min(minLoad, load); + final double standbyLoad = (double) standbyPerProcess.get(entry.getKey()) / process.members.size(); + maxStandbyLoad = Math.max(maxStandbyLoad, standbyLoad); + minStandbyLoad = Math.min(minStandbyLoad, standbyLoad); + } + + // Members without a task of the subtopology count as 0. A spread of 1 is unavoidable when the partitions + // do not divide evenly over the members, so only the excess over that is graded. + final int members = result.members().size(); + int subtopologyExcessSpread = 0; + for (final Map perMember : activePerSubtopologyPerMember.values()) { + final int min = perMember.size() < members ? 0 : perMember.values().stream().min(Integer::compare).orElseThrow(); + final int max = perMember.values().stream().max(Integer::compare).orElseThrow(); + final int partitions = perMember.values().stream().mapToInt(Integer::intValue).sum(); + final int unavoidable = partitions % members == 0 ? 0 : 1; + subtopologyExcessSpread = Math.max(subtopologyExcessSpread, max - min - unavoidable); + } + + final Movement movement = movement(scenario, before, activeMemberOf, activeProcessOf, ownerProcessesOf); + return new Metrics( + spread(activePerMember.values()), + spread(statefulActivePerMember.values()), + subtopologyExcessSpread, + spread(totalPerMember.values()), + maxLoad - minLoad, + maxStandbyLoad - minStandbyLoad, + movement.processStickiness(), + movement.memberStickiness(), + movement.stateMoves(), + stateReuse(restoredTasks, ownerProcessesOf), + copiesLost(scenario, before), + diversityPerTag(scenario, ownerProcessesOf) + ); + } + + private record Movement( + OptionalDouble processStickiness, + OptionalDouble memberStickiness, + OptionalInt stateMoves + ) { + } + + /** + * Compares against the assignment held before the event, over the tasks whose process is still in the group. + * Tasks of a process that left had to move, so they are not counted; tasks of a restarted process are, since + * its state directories still hold them. Member stickiness further skips tasks whose member left, since the + * restarted process comes back with new member ids. Tasks of a removed subtopology are not counted either. + */ + private static Movement movement( + final Scenario scenario, + final Scenario.Baseline before, + final Map activeMemberOf, + final Map activeProcessOf, + final Map> ownerProcessesOf + ) { + if (before.assignment().isEmpty()) { + return new Movement(OptionalDouble.empty(), OptionalDouble.empty(), OptionalInt.empty()); + } + int processCandidates = 0; + int stuckOnProcess = 0; + int memberCandidates = 0; + int stuckOnMember = 0; + for (final Map.Entry entry : before.assignment().entrySet()) { + final String previousMember = entry.getKey(); + final String previousProcess = before.processOfMember().get(previousMember); + if (!scenario.processes.containsKey(previousProcess)) { + continue; + } + final boolean memberStillInGroup = scenario.memberIds().contains(previousMember); + for (final TaskId task : toTaskIds(entry.getValue().activeTasks())) { + if (!scenario.topology.tasks().contains(task)) { + continue; + } + processCandidates++; + if (previousProcess.equals(activeProcessOf.get(task))) { + stuckOnProcess++; + } + if (memberStillInGroup) { + memberCandidates++; + if (previousMember.equals(activeMemberOf.get(task))) { + stuckOnMember++; + } + } + } + } + return new Movement( + processCandidates == 0 ? OptionalDouble.empty() : OptionalDouble.of((double) stuckOnProcess / processCandidates), + memberCandidates == 0 ? OptionalDouble.empty() : OptionalDouble.of((double) stuckOnMember / memberCandidates), + OptionalInt.of(stateMoves(scenario, before, ownerProcessesOf)) + ); + } + + /** + * Copies of a stateful task, active or standby, dropped from a process still in the group and placed on a + * process that held no copy before, so the state has to be rebuilt there. Swapping active and standby between + * two owners is not a move, nor is a dropped copy with no new one elsewhere (fewer standbys configured). + */ + private static int stateMoves( + final Scenario scenario, + final Scenario.Baseline before, + final Map> ownerProcessesOf + ) { + final Map> previousOwnerProcessesOf = new HashMap<>(); + for (final Map.Entry entry : before.assignment().entrySet()) { + final String previousProcess = before.processOfMember().get(entry.getKey()); + if (!scenario.processes.containsKey(previousProcess)) { + continue; + } + final Set previousTasks = toTaskIds(entry.getValue().activeTasks()); + previousTasks.addAll(toTaskIds(entry.getValue().standbyTasks())); + previousTasks.retainAll(scenario.topology.statefulTasks()); + for (final TaskId task : previousTasks) { + previousOwnerProcessesOf.computeIfAbsent(task, t -> new HashSet<>()).add(previousProcess); + } + } + int stateMoves = 0; + for (final Map.Entry> entry : previousOwnerProcessesOf.entrySet()) { + final Set current = ownerProcessesOf.getOrDefault(entry.getKey(), Set.of()); + final Set dropped = new HashSet<>(entry.getValue()); + dropped.removeAll(current); + final Set added = new HashSet<>(current); + added.removeAll(entry.getValue()); + stateMoves += Math.min(dropped.size(), added.size()); + } + return stateMoves; + } + + /** + * Of the tasks the processes reported from their state directories when the rebalance started, the fraction + * that ended up on the reporting process as active or standby. Nothing is graded when nothing was reported. + */ + private static OptionalDouble stateReuse( + final Map> restoredTasks, + final Map> ownerProcessesOf + ) { + int reported = 0; + int reused = 0; + for (final Map.Entry> entry : restoredTasks.entrySet()) { + for (final TaskId task : entry.getValue()) { + reported++; + if (ownerProcessesOf.getOrDefault(task, Set.of()).contains(entry.getKey())) { + reused++; + } + } + } + return reported == 0 ? OptionalDouble.empty() : OptionalDouble.of((double) reused / reported); + } + + /** + * Stateful tasks that lost every copy, active and standbys, to the processes the event removed: the state has + * to be rebuilt from the changelog. What rack-aware standby placement is meant to prevent. This grades the + * placement held before the event. Graded only when several processes left at once, as in a zone failure, and + * standbys are configured; a single process leaving cannot take every copy once standbys are placed. + */ + private static OptionalInt copiesLost(final Scenario scenario, final Scenario.Baseline before) { + final Set gone = new HashSet<>(before.processOfMember().values()); + gone.removeAll(scenario.processes.keySet()); + if (gone.size() < 2 || scenario.numStandbyReplicas == 0) { + return OptionalInt.empty(); + } + final Map survives = new HashMap<>(); + for (final Map.Entry entry : before.assignment().entrySet()) { + final boolean alive = !gone.contains(before.processOfMember().get(entry.getKey())); + for (final TaskId task : toTaskIds(entry.getValue().activeTasks())) { + if (scenario.topology.statefulTasks().contains(task)) { + survives.merge(task, alive, Boolean::logicalOr); + } + } + for (final TaskId task : toTaskIds(entry.getValue().standbyTasks())) { + survives.merge(task, alive, Boolean::logicalOr); + } + } + return OptionalInt.of((int) survives.values().stream().filter(alive -> !alive).count()); + } + + /** + * For each stateful task and configured tag: distinct tag values among the processes holding a copy of the + * task, divided by the best achievable for that tag alone (the configured copies, capped by the distinct values + * present in the group). Tags are bounded separately, so 1.0 on every tag at once may be unreachable; compare + * assignors against each other rather than against 1.0. Averaged per tag over all tasks; a tag with at most one + * value in the group is left out, as is the host tag whose value is unique per process, and nothing is graded + * without standbys since every task then has a single copy. + */ + private static Map diversityPerTag(final Scenario scenario, final Map> ownerProcessesOf) { + final Map diversityPerTag = new LinkedHashMap<>(); + if (scenario.tagKeys.isEmpty() || scenario.numStandbyReplicas == 0) { + return diversityPerTag; + } + final Map> valuesInGroup = new HashMap<>(); + for (final ProcessSpec process : scenario.processes.values()) { + process.tags.forEach((key, value) -> valuesInGroup.computeIfAbsent(key, k -> new HashSet<>()).add(value)); + } + for (final String key : scenario.tagKeys) { + final int available = valuesInGroup.getOrDefault(key, Set.of()).size(); + if (available <= 1 || key.equals(TaskAssignorTestbed.HOST_TAG)) { + continue; + } + int tasks = 0; + double score = 0; + for (final TaskId task : scenario.topology.statefulTasks()) { + final Set owners = ownerProcessesOf.getOrDefault(task, Set.of()); + final Set ownerValues = new HashSet<>(); + for (final String processId : owners) { + final String value = scenario.processes.get(processId).tags.get(key); + if (value != null) { + ownerValues.add(value); + } + } + tasks++; + score += (double) ownerValues.size() / Math.min(1 + scenario.numStandbyReplicas, available); + } + if (tasks > 0) { + diversityPerTag.put(key, score / tasks); + } + } + return diversityPerTag; + } + + private static int spread(final Iterable values) { + int max = Integer.MIN_VALUE; + int min = Integer.MAX_VALUE; + for (final int value : values) { + max = Math.max(max, value); + min = Math.min(min, value); + } + return max == Integer.MIN_VALUE ? 0 : max - min; + } + + /** Aggregates the metrics of every rebalance in a run into one table. */ + static final class Summary { + private final String title; + private final Stat activeSpread = new Stat(); + private final Stat statefulActiveSpread = new Stat(); + private final Stat subtopologyExcessSpread = new Stat(); + private final Stat totalSpread = new Stat(); + private final Stat processLoadSpread = new Stat(); + private final Stat standbyLoadSpread = new Stat(); + private final Stat processStickiness = new Stat(); + private final Stat memberStickiness = new Stat(); + private final Stat stateMoves = new Stat(); + private final Stat stateReuse = new Stat(); + private final Stat copiesLost = new Stat(); + private final Map diversityPerTag = new TreeMap<>(); + private final Stat convergenceIterations = new Stat(); + private int rebalances; + private int notConverged; + + Summary(final String title) { + this.title = title; + } + + void add(final Metrics metrics) { + rebalances++; + activeSpread.add(metrics.activeSpreadPerMember); + statefulActiveSpread.add(metrics.statefulActiveSpreadPerMember); + subtopologyExcessSpread.add(metrics.subtopologyExcessSpread); + totalSpread.add(metrics.totalSpreadPerMember); + processLoadSpread.add(metrics.processLoadSpread); + standbyLoadSpread.add(metrics.standbyLoadSpread); + metrics.processStickiness.ifPresent(processStickiness::add); + metrics.memberStickiness.ifPresent(memberStickiness::add); + metrics.stateMoves.ifPresent(stateMoves::add); + metrics.stateReuse.ifPresent(stateReuse::add); + metrics.copiesLost.ifPresent(copiesLost::add); + metrics.diversityPerTag.forEach((key, value) -> diversityPerTag.computeIfAbsent(key, k -> new Stat()).add(value)); + } + + /** Records how many assignor runs a rebalance needed until the assignment was stable. */ + void addConvergence(final int iterations) { + convergenceIterations.add(iterations); + } + + /** Records a rebalance whose assignment was still changing when the iteration limit was reached. */ + void addNotConverged() { + notConverged++; + } + + @Override + public String toString() { + final StringBuilder builder = new StringBuilder("Fuzz metrics: ").append(title).append('\n') + .append(String.format(" %-44s %8s %8s %8s %8s%n", "metric", "avg", "min", "max", "n")); + row(builder, "active tasks per member (max-min)", activeSpread); + row(builder, "stateful active tasks per member (max-min)", statefulActiveSpread); + row(builder, "subtopology spread above unavoidable", subtopologyExcessSpread); + row(builder, "total tasks per member (max-min)", totalSpread); + row(builder, "process load (max-min)", processLoadSpread); + row(builder, "standby load per process (max-min)", standbyLoadSpread); + row(builder, "process stickiness", processStickiness); + row(builder, "member stickiness", memberStickiness); + row(builder, "stateful copies moved per rebalance", stateMoves); + row(builder, "restored state reused (1.0 = all)", stateReuse); + row(builder, "all copies lost in multi-process failure", copiesLost); + diversityPerTag.forEach((key, stat) -> row(builder, "rack diversity [" + key + "] (vs per-tag best)", stat)); + row(builder, "convergence iterations", convergenceIterations); + builder.append(String.format(" %-44s %d of %d rebalances%n", "not converged within limit", notConverged, rebalances)); + return builder.toString(); + } + + private static void row(final StringBuilder builder, final String name, final Stat stat) { + builder.append(" ").append(String.format("%-44s ", name)).append(stat).append('\n'); + } + } + + static final class Stat { + private int count; + private double sum; + private double min = Double.MAX_VALUE; + private double max = -Double.MAX_VALUE; + + void add(final double value) { + count++; + sum += value; + min = Math.min(min, value); + max = Math.max(max, value); + } + + @Override + public String toString() { + if (count == 0) { + return String.format("%8s %8s %8s %8d", "n/a", "n/a", "n/a", 0); + } + return String.format("%8.3f %8.3f %8.3f %8d", sum / count, min, max, count); + } + } +} diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorFuzzTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorFuzzTest.java new file mode 100644 index 0000000000000..94bc93b4504fe --- /dev/null +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorFuzzTest.java @@ -0,0 +1,99 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.coordinator.group.streams.assignor; + +import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.api.streams.assignor.MemberAssignment; +import org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.Profile; +import org.apache.kafka.coordinator.group.streams.assignor.TaskAssignorTestbed.Scenario; + +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Runs {@link StickyTaskAssignor} through the {@link TaskAssignorTestbed} and adds the checks specific to it: + * no member exceeds the active task quota, every stateful task gets as many standbys as the processes allow, and + * rack-aware tags do not change the active assignment. + */ +public class StickyTaskAssignorFuzzTest { + + private final StickyTaskAssignor assignor = new StickyTaskAssignor(); + private final TaskAssignorTestbed testbed = new TaskAssignorTestbed( + assignor, + List.of(this::verifyActiveQuota, this::verifyStandbyCount, this::verifyTagsDoNotChangeActiveAssignment) + ); + + @Test + public void shouldConvergeToValidAssignmentsForSmallGroupsUnderRandomEvents() { + testbed.run(Profile.SMALL); + } + + @Test + public void shouldConvergeToValidAssignmentsForLargeGroupsUnderRandomEvents() { + testbed.run(Profile.LARGE); + } + + /** The fill phase always picks the least loaded process, so no member ends above {@code ceil(tasks / members)}. */ + private void verifyActiveQuota(final Scenario scenario, final GroupAssignment result, final boolean last) { + final int members = scenario.memberIds().size(); + final int quota = (scenario.topology.tasks().size() + members - 1) / members; + for (final Map.Entry entry : result.members().entrySet()) { + final int active = TaskAssignorTestbed.toTaskIds(entry.getValue().activeTasks()).size(); + assertTrue(active <= quota, entry.getKey() + " holds " + active + " active tasks, above the quota of " + quota); + } + } + + /** Standbys are placed on the least loaded process without the task, so they only fall short when processes run out. */ + private void verifyStandbyCount(final Scenario scenario, final GroupAssignment result, final boolean last) { + final int expectedStandbys = Math.min(scenario.numStandbyReplicas, scenario.processes.size() - 1); + final Map standbyCounts = new HashMap<>(); + for (final MemberAssignment assignment : result.members().values()) { + for (final TaskId task : TaskAssignorTestbed.toTaskIds(assignment.standbyTasks())) { + standbyCounts.merge(task, 1, Integer::sum); + } + } + for (final TaskId task : scenario.topology.statefulTasks()) { + final int actual = standbyCounts.getOrDefault(task, 0); + assertEquals(expectedStandbys, actual, "stateful task " + task + " must have " + expectedStandbys + " standbys but has " + actual); + } + } + + /** + * Rack awareness only places standbys, so the same group without any rack-aware tags must produce the same + * active assignment. Catches a rack-aware change leaking into the active steps or the shared quota bookkeeping. + * Costs a second assignment, so it runs only on the last result of each rebalance. + */ + private void verifyTagsDoNotChangeActiveAssignment(final Scenario scenario, final GroupAssignment result, final boolean last) { + if (!last || scenario.tagKeys.isEmpty()) { + return; + } + final GroupAssignment withoutTags = assignor.assign(scenario.groupSpec(List.of()), scenario.topology); + for (final Map.Entry entry : result.members().entrySet()) { + assertEquals( + withoutTags.members().get(entry.getKey()).activeTasks(), + entry.getValue().activeTasks(), + "active tasks of " + entry.getKey() + " differ between the assignment with and without rack-aware tags" + ); + } + } +} diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/TaskAssignorTestbed.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/TaskAssignorTestbed.java new file mode 100644 index 0000000000000..eb11d5723e168 --- /dev/null +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/TaskAssignorTestbed.java @@ -0,0 +1,781 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.coordinator.group.streams.assignor; + +import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.api.streams.assignor.MemberAssignment; +import org.apache.kafka.coordinator.group.api.streams.assignor.TaskAssignor; +import org.apache.kafka.coordinator.group.api.streams.assignor.TopologyDescriber; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Optional; +import java.util.Random; +import java.util.Set; +import java.util.TreeMap; +import java.util.TreeSet; + +/** + * Randomized testbed for {@link TaskAssignor} implementations. Each scenario generates a random topology and group + * from a seed, runs the assignor until the assignment is stable, then applies random rebalance events (processes + * joining, leaving, restarting or changing tags, all processes of one tag value leaving at once, members joining + * or leaving, the standby count changing, subtopologies growing, appearing or disappearing) and converges again. + * Every assignment is checked against {@link AssignmentInvariants} and the assignor-specific checks given by the + * caller; the converged result of each rebalance is graded by {@link AssignmentMetrics}, and the summary is printed + * so two versions of an assignor can be compared on the same scenarios. + *

+ * Runs are deterministic: all scenarios derive from {@link #DEFAULT_BASE_SEED}. Set + * {@code STREAMS_ASSIGNOR_FUZZ_BASE_SEED=} to run a different set of scenarios, or + * {@code STREAMS_ASSIGNOR_FUZZ_SEED=} to replay the single scenario named by a failure. From an IDE the + * system properties {@code streams.assignor.fuzz.base.seed} and {@code streams.assignor.fuzz.seed} work as well; + * Gradle does not forward {@code -D} to the test JVM. + */ +final class TaskAssignorTestbed { + + /** + * An assignor-specific check run on every assignment, in addition to the invariants. {@code last} is true for + * the final assignment of a rebalance, the converged one or the last attempt, so expensive checks can run once. + */ + @FunctionalInterface + interface AssignmentCheck { + void check(Scenario scenario, GroupAssignment result, boolean last); + } + + static final long DEFAULT_BASE_SEED = 42L; + static final List TAG_KEYS = List.of("region", "zone", "rack"); + static final String HOST_TAG = "host"; + + private static final String SEED_PROPERTY = "streams.assignor.fuzz.seed"; + private static final String SEED_ENVIRONMENT_VARIABLE = "STREAMS_ASSIGNOR_FUZZ_SEED"; + private static final String BASE_SEED_PROPERTY = "streams.assignor.fuzz.base.seed"; + private static final String BASE_SEED_ENVIRONMENT_VARIABLE = "STREAMS_ASSIGNOR_FUZZ_BASE_SEED"; + private static final int MAX_EVENTS = 8; + private static final int MAX_CONVERGENCE_ITERATIONS = 10; + private static final long RESTORED_OFFSET = 100L; + /** Mean of {@link Profile#LARGE}'s membersPerProcess draw, used to size the group from the task count. */ + private static final int EXPECTED_MEMBERS_PER_LARGE_PROCESS = 9; + + /** + * How large the generated groups are. Each bound is drawn per scenario (or per process) from the given random. + */ + enum Profile { + /** + * Hits the logical corners: fewer processes than copies, tag values exhausted, quota boundaries, load ties. + * Small enough that a failure prints every assignment. + */ + SMALL { + int scenarios() { + return 300; + } + + boolean fullHistory() { + return true; + } + + int processes(final Random random, final int tasks) { + return random.nextInt(8) + 1; + } + + int membersPerProcess(final Random random) { + return random.nextInt(4) + 1; + } + + int subtopologies(final Random random) { + return random.nextInt(6) + 1; + } + + int partitions(final Random random) { + return random.nextInt(8) + 1; + } + + boolean stateful(final Random random) { + return random.nextBoolean(); + } + + int tagKeyCount(final Random random) { + return random.nextInt(TAG_KEYS.size() + 1); + } + + int valuesPerTag(final Random random) { + return random.nextInt(3) + 2; + } + + boolean hostTag(final Random random) { + return false; + } + + int standbyReplicas(final Random random) { + return random.nextInt(3); + } + }, + /** + * Production shape: tens to over a hundred processes with a few large ones, sized so that members hold a + * few active tasks each and the quota leaves room for balancing decisions; one or two tags with about three + * values each so tag dimensions rarely run out, standbys mostly 1. Occasionally adds a host tag whose value + * is unique per process. + */ + LARGE { + int scenarios() { + return 30; + } + + boolean fullHistory() { + return false; + } + + int processes(final Random random, final int tasks) { + final int draw = random.nextInt(10); + final int tasksPerMember = draw < 2 ? 2 : draw < 5 ? 3 : draw < 8 ? 4 : draw < 9 ? 6 : 8; + return Math.max(2, tasks / tasksPerMember / EXPECTED_MEMBERS_PER_LARGE_PROCESS); + } + + int membersPerProcess(final Random random) { + return random.nextInt(10) == 0 ? random.nextInt(33) + 32 : random.nextInt(8) + 1; + } + + int subtopologies(final Random random) { + return random.nextInt(17) + 4; + } + + int partitions(final Random random) { + return random.nextInt(113) + 16; + } + + boolean stateful(final Random random) { + return random.nextInt(10) < 6; + } + + int tagKeyCount(final Random random) { + return random.nextInt(2) + 1; + } + + int valuesPerTag(final Random random) { + final int draw = random.nextInt(10); + return draw < 6 ? 3 : draw < 8 ? 2 : 4; + } + + boolean hostTag(final Random random) { + return random.nextInt(5) == 0; + } + + int standbyReplicas(final Random random) { + final int draw = random.nextInt(10); + return draw < 2 ? 0 : draw < 8 ? 1 : 2; + } + }; + + abstract int scenarios(); + + /** Whether the history records every assignment, or only the metrics of each converged result. */ + abstract boolean fullHistory(); + + /** How many processes to generate for a topology of the given size. */ + abstract int processes(Random random, int tasks); + + abstract int membersPerProcess(Random random); + + abstract int subtopologies(Random random); + + abstract int partitions(Random random); + + abstract boolean stateful(Random random); + + abstract int tagKeyCount(Random random); + + abstract int valuesPerTag(Random random); + + abstract boolean hostTag(Random random); + + abstract int standbyReplicas(Random random); + } + + private final TaskAssignor assignor; + private final List checks; + + TaskAssignorTestbed(final TaskAssignor assignor, final List checks) { + this.assignor = assignor; + this.checks = List.copyOf(checks); + } + + /** + * Runs all scenarios of the profile and prints two summaries: the first assignment of each group, computed from + * an empty assignment, and the rebalances after the random events, computed from the previous assignment. + */ + void run(final Profile profile) { + final Long fixedSeed = readSeed(SEED_PROPERTY, SEED_ENVIRONMENT_VARIABLE); + final Long fixedBaseSeed = readSeed(BASE_SEED_PROPERTY, BASE_SEED_ENVIRONMENT_VARIABLE); + final long baseSeed; + if (fixedSeed != null) { + baseSeed = fixedSeed; + } else { + // Offset per profile so the runs do not replay the same scenarios. + final long base = fixedBaseSeed != null ? fixedBaseSeed : DEFAULT_BASE_SEED; + baseSeed = base + profile.ordinal() * 1_000_000L; + } + final int scenarios = fixedSeed != null ? 1 : profile.scenarios(); + final String title = assignor.name() + ", " + profile + ", " + scenarios + " scenarios, base seed " + baseSeed; + final AssignmentMetrics.Summary fromEmpty = new AssignmentMetrics.Summary(title + ", from empty assignment"); + final AssignmentMetrics.Summary fromPrevious = new AssignmentMetrics.Summary(title + ", from previous assignment"); + for (int i = 0; i < scenarios; i++) { + runScenario(baseSeed + i, profile, fromEmpty, fromPrevious); + } + System.out.println(fromEmpty); + System.out.println(fromPrevious); + } + + private static Long readSeed(final String property, final String environmentVariable) { + final Long fromProperty = Long.getLong(property); + if (fromProperty != null) { + return fromProperty; + } + final String fromEnvironment = System.getenv(environmentVariable); + return fromEnvironment == null ? null : Long.valueOf(fromEnvironment); + } + + private void runScenario( + final long seed, + final Profile profile, + final AssignmentMetrics.Summary fromEmpty, + final AssignmentMetrics.Summary fromPrevious + ) { + final Random random = new Random(seed); + final Scenario scenario = Scenario.generate(random, profile); + try { + converge(scenario, fromEmpty, scenario.baseline()); + final int events = random.nextInt(MAX_EVENTS) + 1; + for (int i = 0; i < events; i++) { + final Scenario.Baseline before = scenario.baseline(); + scenario.applyRandomEvent(random); + converge(scenario, fromPrevious, before); + } + } catch (final AssertionError | RuntimeException e) { + throw new AssertionError( + "Fuzz scenario failed. Reproduce with " + SEED_ENVIRONMENT_VARIABLE + "=" + seed + " (or -D" + SEED_PROPERTY + "=" + seed + + ") in the " + profile + " test\n" + scenario.history, + e + ); + } + } + + /** + * Runs the assignor, feeding each result back as the previous assignment, until the result stops changing or + * {@link #MAX_CONVERGENCE_ITERATIONS} is reached. Not reaching a fixed point is recorded, not failed: the + * coordinator runs the assignor once per rebalance, so it only means the next rebalance will move tasks again. + * {@code before} is the state before the event that triggered this rebalance; the last result is graded against it. + */ + private void converge(final Scenario scenario, final AssignmentMetrics.Summary summary, final Scenario.Baseline before) { + // Reported in the first heartbeat after a restart and gone once the members hold an assignment again. + final Map> restoredTasks = scenario.restoredTasksByProcess(); + for (int iteration = 1; iteration <= MAX_CONVERGENCE_ITERATIONS; iteration++) { + final GroupAssignment result = assignor.assign(scenario.groupSpec(), scenario.topology); + final boolean stable = result.members().equals(scenario.previousAssignment); + final boolean last = stable || iteration == MAX_CONVERGENCE_ITERATIONS; + try { + AssignmentInvariants.assertValid(scenario, result); + for (final AssignmentCheck check : checks) { + check.check(scenario, result, last); + } + } catch (final AssertionError e) { + scenario.history.append(" iteration ").append(iteration).append(" input: ").append(format(scenario.previousAssignment)).append('\n') + .append(" iteration ").append(iteration).append(" FAILED: ").append(format(result.members())).append('\n'); + throw e; + } + if (scenario.profile.fullHistory()) { + scenario.history.append(" iteration ").append(iteration).append(": ").append(format(result.members())).append('\n'); + } + scenario.feedBack(result); + if (last) { + final AssignmentMetrics.Metrics metrics = AssignmentMetrics.compute(scenario, result, before, restoredTasks); + summary.add(metrics); + if (stable) { + summary.addConvergence(iteration); + scenario.history.append(" converged after ").append(iteration).append(": ").append(metrics).append('\n'); + } else { + summary.addNotConverged(); + scenario.history.append(" not converged within ").append(iteration).append(": ").append(metrics).append('\n'); + } + return; + } + } + } + + static Set toTaskIds(final Map> tasks) { + final Set taskIds = new HashSet<>(); + tasks.forEach((subtopology, partitions) -> partitions.forEach(partition -> taskIds.add(new TaskId(subtopology, partition)))); + return taskIds; + } + + static String format(final Map members) { + final StringBuilder builder = new StringBuilder(); + for (final Map.Entry entry : new TreeMap<>(members).entrySet()) { + builder.append(entry.getKey()) + .append(" A").append(new TreeSet<>(toTaskIds(entry.getValue().activeTasks()))) + .append(" S").append(new TreeSet<>(toTaskIds(entry.getValue().standbyTasks()))) + .append("; "); + } + return builder.toString(); + } + + // ---- Random scenario: topology, group and events ---- + + record Subtopology(String id, int partitions, boolean stateful) { + } + + record Topology(List specs, Set tasks, Set statefulTasks) implements TopologyDescriber { + + static Topology generate(final Random random, final Profile profile) { + final List specs = new ArrayList<>(); + final int count = profile.subtopologies(random); + for (int i = 0; i < count; i++) { + specs.add(new Subtopology("s" + i, profile.partitions(random), profile.stateful(random))); + } + return of(specs); + } + + static Topology of(final List specs) { + final Set tasks = new HashSet<>(); + final Set statefulTasks = new HashSet<>(); + for (final Subtopology subtopology : specs) { + for (int partition = 0; partition < subtopology.partitions; partition++) { + final TaskId task = new TaskId(subtopology.id, partition); + tasks.add(task); + if (subtopology.stateful) { + statefulTasks.add(task); + } + } + } + return new Topology(List.copyOf(specs), tasks, statefulTasks); + } + + @Override + public List subtopologies() { + return specs.stream().map(Subtopology::id).toList(); + } + + @Override + public int maxNumInputPartitions(final String subtopologyId) throws NoSuchElementException { + return find(subtopologyId).partitions; + } + + @Override + public boolean isStateful(final String subtopologyId) { + return find(subtopologyId).stateful; + } + + private Subtopology find(final String subtopologyId) { + return specs.stream() + .filter(subtopology -> subtopology.id.equals(subtopologyId)) + .findFirst() + .orElseThrow(); + } + + @Override + public String toString() { + return specs.toString(); + } + } + + static final class ProcessSpec { + final List members = new ArrayList<>(); + final Map tags; + + ProcessSpec(final Map tags) { + this.tags = tags; + } + + @Override + public String toString() { + return members + " " + tags; + } + } + + /** A group under test: its topology, processes and members, and the assignment they currently hold. */ + static final class Scenario { + final Profile profile; + final List tagKeys; + Topology topology; + final Map processes = new LinkedHashMap<>(); + final StringBuilder history = new StringBuilder(); + Map previousAssignment = Map.of(); + int numStandbyReplicas; + + private final int valuesPerTag; + private final Map memberToProcess = new HashMap<>(); + private final Map>> restoredOffsets = new HashMap<>(); + private int nextProcess; + private int nextGeneration; + private int nextSubtopology; + + /** + * What the members held before the event that triggered a rebalance, and on which process, so that + * stickiness can follow tasks across restarts. + */ + record Baseline(Map assignment, Map processOfMember) { + } + + private Scenario( + final Profile profile, + final Topology topology, + final List tagKeys, + final int valuesPerTag + ) { + this.profile = profile; + this.topology = topology; + this.tagKeys = tagKeys; + this.valuesPerTag = valuesPerTag; + this.nextSubtopology = topology.specs().size(); + } + + static Scenario generate(final Random random, final Profile profile) { + final Topology topology = Topology.generate(random, profile); + final List tagKeys = new ArrayList<>(TAG_KEYS.subList(0, profile.tagKeyCount(random))); + if (profile.hostTag(random)) { + tagKeys.add(HOST_TAG); + } + final Scenario scenario = new Scenario(profile, topology, List.copyOf(tagKeys), profile.valuesPerTag(random)); + scenario.numStandbyReplicas = profile.standbyReplicas(random); + final int processCount = profile.processes(random, topology.tasks().size()); + for (int i = 0; i < processCount; i++) { + scenario.addProcess(random); + } + scenario.history.append("topology ").append(topology) + .append(", standbyReplicas=").append(scenario.numStandbyReplicas) + .append(", tags=").append(tagKeys).append(" with ").append(scenario.valuesPerTag).append(" values") + .append(", processes ").append(scenario.processes).append('\n'); + return scenario; + } + + /** A rebalance event; returns false when it is not applicable to the group as it is, so another is drawn. */ + @FunctionalInterface + private interface Event { + boolean apply(Random random); + } + + private final List events = List.of( + this::addProcessEvent, + this::dropProcessEvent, + this::restartProcessEvent, + this::addMemberEvent, + this::dropMemberEvent, + this::retagProcess, + this::dropTagValue, + this::expandPartitions, + this::addSubtopology, + this::removeSubtopology, + this::changeStandbyReplicas + ); + + /** Applies one random event that changes the group; a drawn event that would change nothing is redrawn. */ + void applyRandomEvent(final Random random) { + boolean applied = false; + while (!applied) { + applied = events.get(random.nextInt(events.size())).apply(random); + } + } + + private boolean addProcessEvent(final Random random) { + final String processId = addProcess(random); + history.append("event: add process ").append(processId).append(' ').append(processes.get(processId)).append('\n'); + return true; + } + + private boolean dropProcessEvent(final Random random) { + if (processes.size() == 1) { + return false; + } + final String processId = randomProcess(random); + removeProcess(processId); + history.append("event: drop process ").append(processId).append('\n'); + return true; + } + + private boolean restartProcessEvent(final Random random) { + final String processId = randomProcess(random); + restartProcess(processId); + history.append("event: restart process ").append(processId).append(" -> ").append(processes.get(processId)).append('\n'); + return true; + } + + private boolean addMemberEvent(final Random random) { + final String memberId = addMember(randomProcess(random)); + history.append("event: add member ").append(memberId).append('\n'); + return true; + } + + private boolean dropMemberEvent(final Random random) { + final String processId = randomProcess(random); + final ProcessSpec process = processes.get(processId); + if (process.members.size() == 1) { + return false; + } + final String memberId = process.members.get(random.nextInt(process.members.size())); + removeMember(processId, memberId); + history.append("event: drop member ").append(memberId).append('\n'); + return true; + } + + private boolean retagProcess(final Random random) { + final String processId = randomProcess(random); + final String key = randomTagKey(random); + if (key == null) { + return false; + } + final String value = key + "-" + random.nextInt(valuesPerTag); + if (value.equals(processes.get(processId).tags.put(key, value))) { + return false; + } + history.append("event: retag process ").append(processId).append(" -> ").append(processes.get(processId)).append('\n'); + return true; + } + + private boolean changeStandbyReplicas(final Random random) { + final int standbyReplicas = profile.standbyReplicas(random); + if (standbyReplicas == numStandbyReplicas) { + return false; + } + numStandbyReplicas = standbyReplicas; + history.append("event: standbyReplicas=").append(numStandbyReplicas).append('\n'); + return true; + } + + private String addProcess(final Random random) { + final String processId = "p" + nextProcess++; + final Map tags = new HashMap<>(); + for (final String key : tagKeys) { + if (key.equals(HOST_TAG)) { + tags.put(key, processId); + } else if (random.nextInt(10) != 0) { + // A process occasionally misses a tag, as a client without that config would. + tags.put(key, key + "-" + random.nextInt(valuesPerTag)); + } + } + processes.put(processId, new ProcessSpec(tags)); + final int members = profile.membersPerProcess(random); + for (int i = 0; i < members; i++) { + addMember(processId); + } + return processId; + } + + private String addMember(final String processId) { + final ProcessSpec process = processes.get(processId); + final String memberId = processId + "-m" + process.members.size() + "g" + nextGeneration++; + process.members.add(memberId); + memberToProcess.put(memberId, processId); + return memberId; + } + + private void removeMember(final String processId, final String memberId) { + processes.get(processId).members.remove(memberId); + memberToProcess.remove(memberId); + restoredOffsets.remove(memberId); + final Map remaining = new HashMap<>(previousAssignment); + remaining.remove(memberId); + previousAssignment = remaining; + } + + private void removeProcess(final String processId) { + final ProcessSpec process = processes.remove(processId); + for (final String memberId : process.members) { + memberToProcess.remove(memberId); + restoredOffsets.remove(memberId); + } + final Map remaining = new HashMap<>(previousAssignment); + remaining.keySet().removeAll(process.members); + previousAssignment = remaining; + } + + /** + * The process comes back with fresh member ids and no target assignment, but its state directories still + * hold the stateful tasks it owned, which it reports as task offsets. + */ + private void restartProcess(final String processId) { + final ProcessSpec process = processes.get(processId); + final List oldMembers = new ArrayList<>(process.members); + final Map remaining = new HashMap<>(previousAssignment); + process.members.clear(); + for (final String oldMember : oldMembers) { + memberToProcess.remove(oldMember); + restoredOffsets.remove(oldMember); + final MemberAssignment previous = remaining.remove(oldMember); + final String newMember = addMember(processId); + if (previous != null) { + final Map> offsets = new HashMap<>(); + addOffsets(offsets, previous.activeTasks()); + addOffsets(offsets, previous.standbyTasks()); + restoredOffsets.put(newMember, offsets); + } + } + previousAssignment = remaining; + } + + Baseline baseline() { + return new Baseline(previousAssignment, Map.copyOf(memberToProcess)); + } + + /** The tasks each process currently reports from its state directories, as a restarted process does. */ + Map> restoredTasksByProcess() { + final Map> restoredTasksByProcess = new HashMap<>(); + restoredOffsets.forEach((memberId, offsets) -> + restoredTasksByProcess.computeIfAbsent(memberToProcess.get(memberId), p -> new HashSet<>()).addAll(toTaskIds(toPartitions(offsets))) + ); + return restoredTasksByProcess; + } + + private static Map> toPartitions(final Map> offsets) { + final Map> partitions = new HashMap<>(); + offsets.forEach((subtopology, perPartition) -> partitions.put(subtopology, perPartition.keySet())); + return partitions; + } + + /** Only stateful tasks leave state on disk, so only they are reported. */ + private void addOffsets(final Map> offsets, final Map> tasks) { + tasks.forEach((subtopology, partitions) -> partitions.forEach(partition -> { + if (topology.statefulTasks().contains(new TaskId(subtopology, partition))) { + offsets.computeIfAbsent(subtopology, s -> new HashMap<>()).put(partition, RESTORED_OFFSET); + } + })); + } + + private String randomProcess(final Random random) { + final List ids = new ArrayList<>(processes.keySet()); + return ids.get(random.nextInt(ids.size())); + } + + /** A configured tag key other than the per-process host tag, or null if there is none. */ + private String randomTagKey(final Random random) { + final List keys = tagKeys.stream().filter(key -> !key.equals(HOST_TAG)).toList(); + return keys.isEmpty() ? null : keys.get(random.nextInt(keys.size())); + } + + /** + * Correlated failure: every process sharing one value of a tag leaves at once, as when a zone goes down. + * Not applicable without tags, or when the value covers the whole group. + */ + private boolean dropTagValue(final Random random) { + final String key = randomTagKey(random); + if (key == null) { + return false; + } + final List values = processes.values().stream().map(process -> process.tags.get(key)).filter(Objects::nonNull).distinct().toList(); + if (values.isEmpty()) { + return false; + } + final String value = values.get(random.nextInt(values.size())); + final List victims = processes.entrySet().stream() + .filter(entry -> value.equals(entry.getValue().tags.get(key))) + .map(Map.Entry::getKey) + .toList(); + if (victims.size() == processes.size()) { + return false; + } + victims.forEach(this::removeProcess); + history.append("event: drop all processes with ").append(key).append('=').append(value).append(' ').append(victims).append('\n'); + return true; + } + + /** An input topic gets more partitions: the subtopology grows by up to its current size. */ + private boolean expandPartitions(final Random random) { + final List specs = new ArrayList<>(topology.specs()); + final int index = random.nextInt(specs.size()); + final Subtopology old = specs.get(index); + specs.set(index, new Subtopology(old.id, old.partitions + 1 + random.nextInt(old.partitions), old.stateful)); + topology = Topology.of(specs); + history.append("event: expand ").append(old.id).append(" from ").append(old.partitions).append(" to ").append(specs.get(index).partitions).append(" partitions\n"); + return true; + } + + /** A topology update adds a subtopology; its tasks have no previous owner. */ + private boolean addSubtopology(final Random random) { + final List specs = new ArrayList<>(topology.specs()); + final Subtopology added = new Subtopology("s" + nextSubtopology++, profile.partitions(random), profile.stateful(random)); + specs.add(added); + topology = Topology.of(specs); + history.append("event: add subtopology ").append(added).append('\n'); + return true; + } + + /** + * A topology update removes a subtopology. The members still report its tasks as owned until the next + * assignment, so the assignor sees tasks that no longer exist. + */ + private boolean removeSubtopology(final Random random) { + if (topology.specs().size() == 1) { + return false; + } + final List specs = new ArrayList<>(topology.specs()); + final Subtopology removed = specs.remove(random.nextInt(specs.size())); + topology = Topology.of(specs); + history.append("event: remove subtopology ").append(removed.id).append('\n'); + return true; + } + + GroupSpecImpl groupSpec() { + return groupSpec(tagKeys); + } + + /** The group with the given rack-aware tags configured; the members still carry their client tags. */ + GroupSpecImpl groupSpec(final List rackAwareAssignmentTags) { + final Map members = new HashMap<>(); + processes.forEach((processId, process) -> { + for (final String memberId : process.members) { + final MemberAssignment previous = previousAssignment.get(memberId); + members.put(memberId, new MemberMetadataAndStateImpl( + Optional.empty(), + Optional.empty(), + processId, + process.tags, + previous == null ? Map.of() : previous.activeTasks(), + previous == null ? Map.of() : previous.standbyTasks(), + Map.of(), + restoredOffsets.getOrDefault(memberId, Map.of()), + Map.of() + )); + } + }); + return new GroupSpecImpl( + members, + AssignmentConfigsImpl.DEFAULT + .withNumStandbyReplicas(numStandbyReplicas) + .withRackAwareAssignmentTags(rackAwareAssignmentTags) + ); + } + + void feedBack(final GroupAssignment result) { + previousAssignment = new HashMap<>(result.members()); + restoredOffsets.clear(); + } + + Set memberIds() { + return memberToProcess.keySet(); + } + + String processOf(final String memberId) { + final String processId = memberToProcess.get(memberId); + if (processId == null) { + throw new AssertionError("assignment names unknown member " + memberId); + } + return processId; + } + } +}