lucasbru commented on code in PR #23556: URL: https://github.com/apache/kafka/pull/23556#discussion_r4132780453
########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/TaskAssignorTestbed.java: ########## @@ -0,0 +1,737 @@ +/* + * 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.Optional; +import java.util.OptionalDouble; +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 converges, then applies random changes (a process joining, + * leaving or restarting, a member leaving, a subtopology being added) 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 assignor-specific graders given by the caller, and + * the summary is printed so two assignors, or two versions of one, can be compared on the same scenarios. When the + * scenario has client tags, the same input is also assigned with {@code rackAwareAssignmentTags} empty, and the + * tag-related rows of that tag-blind result are reported next to the real ones. + * <p> + * Runs are deterministic: all scenarios derive from {@link #DEFAULT_BASE_SEED}. Set + * {@code STREAMS_ASSIGNOR_FUZZ_BASE_SEED=<seed>} to run a different set of scenarios, or + * {@code STREAMS_ASSIGNOR_FUZZ_SEED=<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); + } + + /** + * Assignor-specific metrics graded on the converged result of every rebalance, next to {@link AssignmentMetrics}. + * Every call returns the same rows in the same order; an empty value is not graded for this rebalance. + * {@code reportedTasks} are the stateful tasks each process reported offsets for when the rebalance started, + * as a restarted process does for the state on its disk. + */ + @FunctionalInterface + interface AssignmentGrader { + Map<String, OptionalDouble> grade( + Scenario scenario, + GroupAssignment result, + Scenario.Baseline before, + Map<String, Set<TaskId>> reportedTasks + ); + } + + static final long DEFAULT_BASE_SEED = 42L; + static final List<String> 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_CHANGES = 8; + private static final int MAX_CONVERGENCE_ITERATIONS = 10; + private static final long REPORTED_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<AssignmentCheck> checks; + private final List<AssignmentGrader> graders; + + TaskAssignorTestbed(final TaskAssignor assignor, final List<AssignmentCheck> checks, final List<AssignmentGrader> graders) { + this.assignor = assignor; + this.checks = List.copyOf(checks); + final List<AssignmentGrader> allGraders = new ArrayList<>(); + allGraders.add(AssignmentMetrics::grade); + allGraders.addAll(graders); + this.graders = List.copyOf(allGraders); + } + + /** + * Runs all scenarios of the profile and prints two summaries: the initial assignment of each group, computed from + * an empty assignment, and the rebalances after the random changes, 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); + } Review Comment: Two things worth adding here, I think: 1. Timing. There's no reference implementation in this module to compare against, so I wouldn't gate CI on perf, that tends to be flaky, but it'd be useful to at least measure and print per-rebalance assignment time in the summary, so a future accidental complexity regression at least shows up in the output. 2. Right now this only runs as a JUnit test with the full SMALL/LARGE profiles, and the summary is print-only, nothing here is asserted except the assignor-specific checks. Could be worth splitting into two entry points: a `main` that runs the full profiles and prints the summary tables for manual comparison, and a separate fast test with a fixed seed and a small scenario count that asserts fixed bounds on a couple of the metrics, so a real quality regression actually fails the build instead of just printing a worse number in CI output that nobody reads. -- 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]
