You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by "suzhiking (via GitHub)" <gi...@apache.org> on 2026/09/29 05:37:19 UTC
[PR] KAFKA-21026: Add JMH benchmark for the streams assignment refiner [kafka]
suzhiking opened a new pull request, #23616:
URL: https://github.com/apache/kafka/pull/23616
The streams group coordinator refines the task assignor's target
assignment into an intermediate assignment that holds stateful task
migrations back behind warm-up tasks (KAFKA-20665). The refinement runs
on the heartbeat path: for a stable group, `streamsGroupHeartbeat`
derives it for the whole group on every heartbeat, so its cost scales
with each group's heartbeat rate. This adds a JMH benchmark for it.
`AssignmentRefinerBenchmark` measures one `AssignmentRefiner#refine`
call followed by the `preservesActiveTaskCount` check the coordinator
applies to the result, for `NoOpAssignmentRefiner` (the broker default)
and `AssignmentRefinerImpl`. It calls the public `refine`, so it needs
no change as the refiner evolves. The group is set up as a refinement
step finds it: the target assignment comes from the sticky assignor, a
fraction of the stateful active tasks still runs on a member of another
process, the warm-up budget is already spent on migrations under way,
and the members report restore progress according to a lag picture (all
caught up, all lagging, not reported, mixed, or destinations that are
themselves restoring an active task). The parameters cover member count,
members per process, subtopology and partition count, standby replicas,
the warm-up budget, the migrating fraction and the lag picture. The
default grid is trimmed to values that bracket the behaviour; the rest
are reachable with `-p`.
`StreamsAssignorBenchmarkUtils` gains the two fixture builders this
needs: `divergeAssignment`, which derives the members' current
assignment from the target, and `createMemberTaskOffsets`, which builds
the offsets and end offsets the members report in their heartbeats.
Testing: `./gradlew :jmh-benchmarks:compileJava checkstyleMain
spotbugsMain` passes. A probe over the fixtures confirmed that no
process ever holds two copies of a task, that the active task count is
preserved, and that each lag picture drives a different refiner
decision (grant, keep, plant, park). A run against trunk with the
refiner wired in (288 refiner rows, 5 warm-up and 5 measured iterations
of 1 s) shows that a stable group with nothing to move is the dominant
cost: about 4.2-4.8 ms and 5 MB per call at 5,000 stateful tasks with
one standby, 58-111 times the no-op refiner, largely independent of
member count. Moving half of the stateful tasks adds 50-100% on top, and
the warm-up budget (2 vs. 100) makes no measurable difference.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-21026: Add JMH benchmark for the streams assignment refiner [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #23616:
URL: https://github.com/apache/kafka/pull/23616#discussion_r4205222687
##########
jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignmentRefinerBenchmark.java:
##########
@@ -0,0 +1,175 @@
+/*
+ * 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.jmh.assignor;
+
+import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
+import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment;
+import org.apache.kafka.coordinator.group.streams.AssignmentRefiner;
+import org.apache.kafka.coordinator.group.streams.AssignmentRefinerImpl;
+import org.apache.kafka.coordinator.group.streams.MemberTaskOffsets;
+import org.apache.kafka.coordinator.group.streams.NoOpAssignmentRefiner;
+import org.apache.kafka.coordinator.group.streams.StreamsGroupMember;
+import org.apache.kafka.coordinator.group.streams.TasksTuple;
+import org.apache.kafka.coordinator.group.streams.TopologyMetadata;
+import org.apache.kafka.coordinator.group.streams.assignor.AssignmentConfigsImpl;
+import org.apache.kafka.coordinator.group.streams.assignor.StickyTaskAssignor;
+import org.apache.kafka.coordinator.group.streams.topics.ConfiguredSubtopology;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.SortedMap;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Measures one derivation of the streams group's intermediate assignment: what the group coordinator runs when it
+ * refines the target assignment into the warm-up steps the members converge through, together with the check it
+ * applies to the result before using it.
+ * <p>
+ * The group is set up as a refinement step finds it: the target assignment comes from the sticky assignor, a fraction
+ * of the stateful active tasks still runs on a member of another process, the warm-up budget is spent on migrations
+ * already under way, and the members report their restore progress as the lag picture says.
+ * <p>
+ * The full grid is large; narrow it with -p, for example -p memberCount=1000 -p lagPicture=MIXED.
+ */
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 5, time = 1)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+public class AssignmentRefinerBenchmark {
+
+ /**
+ * The default of acceptable.recovery.lag.
+ */
+ private static final long ACCEPTABLE_RECOVERY_LAG = 10_000L;
+
+ @Param({"100", "1000"})
+ private int memberCount;
+
+ @Param({"1", "50"})
+ private int membersPerProcess;
+
+ @Param({"10", "100"})
+ private int subtopologyCount;
+
+ @Param({"100"})
+ private int partitionCount;
+
+ @Param({"0", "1"})
+ private int standbyReplicas;
+
+ @Param({"2", "100"})
+ private int numWarmupReplicas;
+
+ @Param({"0.0", "0.01", "0.5"})
+ private double migratingTaskFraction;
+
+ @Param({"ALL_CAUGHT_UP", "ALL_LAGGING", "NOT_REPORTED", "MIXED", "RESTORING_DESTINATIONS"})
+ private StreamsAssignorBenchmarkUtils.LagPicture lagPicture;
+
+ private final AssignmentRefiner noOpRefiner = new NoOpAssignmentRefiner();
+
+ private final AssignmentRefiner refinerImpl = new AssignmentRefinerImpl();
+
+ private Map<String, StreamsGroupMember> members;
+
+ private Map<String, TasksTuple> targetAssignment;
+
+ private Map<String, MemberTaskOffsets> taskOffsets;
+
+ private SortedMap<String, ConfiguredSubtopology> subtopologyMap;
+
+ @Setup(Level.Trial)
+ public void setup() {
+ List<String> allTopicNames = AssignorBenchmarkUtils.createTopicNames(subtopologyCount);
+ subtopologyMap = StreamsAssignorBenchmarkUtils.createSubtopologyMap(partitionCount, allTopicNames);
+ CoordinatorMetadataImage metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionCount);
+
+ Map<String, StreamsGroupMember> newMembers = StreamsAssignorBenchmarkUtils.createStreamsMembers(memberCount, membersPerProcess);
+ GroupAssignment groupAssignment = new StickyTaskAssignor().assign(
+ StreamsAssignorBenchmarkUtils.createGroupSpec(
+ newMembers,
+ AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(standbyReplicas),
+ Map.of()
+ ),
+ new TopologyMetadata(metadataImage, subtopologyMap)
+ );
+
+ targetAssignment = new HashMap<>();
+ groupAssignment.members().forEach((memberId, memberAssignment) -> targetAssignment.put(
+ memberId,
+ new TasksTuple(memberAssignment.activeTasks(), memberAssignment.standbyTasks(), Map.of())
+ ));
+
+ members = StreamsAssignorBenchmarkUtils.divergeAssignment(
+ newMembers,
+ targetAssignment,
+ subtopologyMap,
+ migratingTaskFraction,
+ numWarmupReplicas
+ );
+ taskOffsets = StreamsAssignorBenchmarkUtils.createMemberTaskOffsets(
+ members,
+ targetAssignment,
+ subtopologyMap,
+ lagPicture,
+ ACCEPTABLE_RECOVERY_LAG
+ );
+ }
+
+ @Benchmark
+ @Threads(1)
+ public void noOpRefiner(Blackhole blackhole) {
+ refine(noOpRefiner, blackhole);
+ }
+
+ @Benchmark
+ @Threads(1)
+ public void refinerImpl(Blackhole blackhole) {
+ refine(refinerImpl, blackhole);
+ }
+
+ private void refine(AssignmentRefiner refiner, Blackhole blackhole) {
Review Comment:
This measures one refine call, which is the right unit, but the description reads it as a per-call cost. In streamsGroupHeartbeat a stable group runs the refinement on every heartbeat without using the per-epoch cache (that cache only serves the non-stable path and new target assignments), so a group costs about memberCount / heartbeatInterval times the per-call time. With 1000 members at the default 5 s interval that is about 200 calls/s, so at the 4.2-4.8 ms you measured about 0.9 s of coordinator time per second, on a thread shared with every other group of the shard. The per-call time does not depend on member count, so the group cost grows linearly with members. I think that is the main result of this benchmark and it would be good to state it in the description. The inputs of a stable group only change when a member reports new task offsets (every 60 s by default) or the epoch is bumped, so most of these calls recompute an identical result. It might be worth skipping the der
ivation when nothing changed, and this benchmark would be a good way to check such a change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-21026: Add JMH benchmark for the streams assignment refiner [kafka]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #23616:
URL: https://github.com/apache/kafka/pull/23616#issuecomment-6030246950
A label of 'needs-attention' was automatically added to this PR in order to raise the
attention of the committers. Once this issue has been triaged, the `triage` label
should be removed to prevent this automation from happening again.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-21026: Add JMH benchmark for the streams assignment refiner [kafka]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #23616:
URL: https://github.com/apache/kafka/pull/23616#discussion_r4205062474
##########
jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmarkUtils.java:
##########
@@ -181,4 +224,275 @@ public static SortedMap<String, ConfiguredSubtopology> createSubtopologyMap(
}
return subtopologyMap;
}
+
+ /**
+ * Creates the members' current assignment as the target assignment with a fraction of the stateful active tasks
+ * still running where they were before the target assignment moved them, which is the work the refiner has to
+ * do. A group whose current assignment equals its target leaves the refiner nothing to decide.
+ * <p>
+ * Every chosen task runs on a member of another process than its target owner, preferring one that holds no copy
+ * of the task, round-robin over that process's members. When every other process holds a standby of the task,
+ * the active and one standby swap places instead, so that the target owner holds the standby. The chosen tasks
+ * are spread evenly over the stateful tasks in task order. The target's standby tasks are kept otherwise.
+ * <p>
+ * A refinement step after the first one finds the warm-up budget spent on migrations already under way, so
+ * inFlightWarmups of the migrating tasks whose target owner holds nothing for them get a warm-up task on it,
+ * spread evenly over those tasks.
+ *
+ * @param members The members, which must span at least two processes if any task migrates.
+ * @param targetAssignment The target assignment, keyed by member ID.
+ * @param subtopologyMap The subtopologies; only the tasks of stateful ones migrate.
+ * @param migratingTaskFraction The fraction of the stateful active tasks not on their target owner yet.
+ * @param inFlightWarmups The number of migrating tasks already warming up on their target owner.
+ *
+ * @return The members, with their current assignment set.
+ */
+ public static Map<String, StreamsGroupMember> divergeAssignment(
+ Map<String, StreamsGroupMember> members,
+ Map<String, TasksTuple> targetAssignment,
+ SortedMap<String, ConfiguredSubtopology> subtopologyMap,
+ double migratingTaskFraction,
+ int inFlightWarmups
+ ) {
+ Map<String, Map<String, Map<Integer, Integer>>> activeTasks = new HashMap<>();
+ Map<String, Map<String, Set<Integer>>> standbyTasks = new HashMap<>();
+ Map<String, Map<String, Set<Integer>>> warmupTasks = new HashMap<>();
+ TreeMap<TaskId, String> statefulActiveOwners = new TreeMap<>();
+ Map<TaskId, SortedSet<String>> standbyHolders = new HashMap<>();
+
+ for (Map.Entry<String, TasksTuple> entry : targetAssignment.entrySet()) {
+ String memberId = entry.getKey();
+ int memberEpoch = members.get(memberId).memberEpoch();
+ entry.getValue().activeTasks().forEach((subtopologyId, partitions) -> partitions.forEach(partition -> {
+ addActive(activeTasks, memberId, new TaskId(subtopologyId, partition), memberEpoch);
+ if (isStateful(subtopologyMap, subtopologyId)) {
+ statefulActiveOwners.put(new TaskId(subtopologyId, partition), memberId);
+ }
+ }));
+ entry.getValue().standbyTasks().forEach((subtopologyId, partitions) -> partitions.forEach(partition -> {
+ addTask(standbyTasks, memberId, new TaskId(subtopologyId, partition));
+ standbyHolders.computeIfAbsent(new TaskId(subtopologyId, partition), id -> new TreeSet<>()).add(memberId);
+ }));
+ }
+
+ List<String> processIds = new ArrayList<>();
+ Map<String, List<String>> membersByProcess = new TreeMap<>();
+ members.values().stream().sorted(Comparator.comparing(StreamsGroupMember::memberId)).forEach(member ->
+ membersByProcess.computeIfAbsent(member.processId(), id -> new ArrayList<>()).add(member.memberId()));
+ processIds.addAll(membersByProcess.keySet());
+ if (migratingTaskFraction > 0 && processIds.size() < 2) {
+ throw new IllegalArgumentException("A task can only migrate between two processes, but the members run in "
+ + processIds.size() + ".");
+ }
+ Map<String, Integer> nextMemberOfProcess = new HashMap<>();
+
+ List<Map.Entry<TaskId, String>> warmupCandidates = new ArrayList<>();
+ int taskIndex = 0;
+ for (Map.Entry<TaskId, String> entry : statefulActiveOwners.entrySet()) {
+ boolean migrating = Math.floor((taskIndex + 1) * migratingTaskFraction) > Math.floor(taskIndex * migratingTaskFraction);
+ taskIndex++;
+ if (!migrating) {
+ continue;
+ }
+
+ TaskId task = entry.getKey();
+ String targetOwner = entry.getValue();
+ int epoch = activeTasks.get(targetOwner).get(task.subtopologyId()).get(task.partition());
+ removeActive(activeTasks, targetOwner, task);
+
+ Optional<String> freeProcess = processWithoutCopy(members, processIds, targetOwner, standbyHolders.get(task));
+ if (freeProcess.isPresent()) {
+ List<String> processMembers = membersByProcess.get(freeProcess.get());
+ int next = nextMemberOfProcess.merge(freeProcess.get(), 1, Integer::sum) - 1;
+ addActive(activeTasks, processMembers.get(next % processMembers.size()), task, epoch);
+ warmupCandidates.add(entry);
+ } else {
+ String standbyHolder = standbyHolders.get(task).first();
+ removeTask(standbyTasks, standbyHolder, task);
+ addActive(activeTasks, standbyHolder, task, epoch);
+ addTask(standbyTasks, targetOwner, task);
+ }
+ }
+
+ int warmups = Math.min(inFlightWarmups, warmupCandidates.size());
+ for (int warmup = 0; warmup < warmups; warmup++) {
+ Map.Entry<TaskId, String> candidate =
+ warmupCandidates.get((int) ((long) warmup * warmupCandidates.size() / warmups));
+ addTask(warmupTasks, candidate.getValue(), candidate.getKey());
+ }
+
+ Map<String, StreamsGroupMember> divergedMembers = new HashMap<>();
+ for (StreamsGroupMember member : members.values()) {
+ String memberId = member.memberId();
+ divergedMembers.put(memberId, new StreamsGroupMember.Builder(member)
+ .setAssignedTasks(new TasksTupleWithEpochs(
+ activeTasks.getOrDefault(memberId, Map.of()),
+ standbyTasks.getOrDefault(memberId, Map.of()),
+ warmupTasks.getOrDefault(memberId, Map.of())
+ ))
+ .build());
+ }
+ return divergedMembers;
+ }
+
+ /**
+ * The first process after the target owner's, in process ID order, that holds no standby of the task.
+ */
+ private static Optional<String> processWithoutCopy(
+ Map<String, StreamsGroupMember> members,
+ List<String> processIds,
+ String targetOwner,
+ Set<String> standbyHolders
+ ) {
+ Set<String> copyProcesses = new HashSet<>();
+ if (standbyHolders != null) {
+ standbyHolders.forEach(holder -> copyProcesses.add(members.get(holder).processId()));
+ }
+
+ int targetProcessIndex = processIds.indexOf(members.get(targetOwner).processId());
+ for (int offset = 1; offset < processIds.size(); offset++) {
+ String processId = processIds.get((targetProcessIndex + offset) % processIds.size());
+ if (!copyProcesses.contains(processId)) {
+ return Optional.of(processId);
+ }
+ }
+ return Optional.empty();
+ }
+
+ /**
+ * Creates the changelog offsets and end offsets the members report in their heartbeats, from which the refiner
+ * tells whether a copy of a task is caught up.
+ * <p>
+ * Every member reports the stateful standby and warm-up tasks of its current assignment, lagging or caught up as
+ * the lag picture says. Under RESTORING_DESTINATIONS, every member that the target assignment moves a stateful
+ * active task to also reports its first stateful active task that stays with it, which reads as restoring.
+ *
+ * @param members The members, with their current assignment set.
+ * @param targetAssignment The target assignment, keyed by member ID.
+ * @param subtopologyMap The subtopologies; only the tasks of stateful ones are reported.
+ * @param lagPicture What the members report.
+ * @param acceptableRecoveryLag The lag at or below which a task counts as caught up.
+ *
+ * @return The reported offsets, per member.
+ */
+ public static Map<String, MemberTaskOffsets> createMemberTaskOffsets(
+ Map<String, StreamsGroupMember> members,
+ Map<String, TasksTuple> targetAssignment,
+ SortedMap<String, ConfiguredSubtopology> subtopologyMap,
+ LagPicture lagPicture,
+ long acceptableRecoveryLag
+ ) {
+ if (acceptableRecoveryLag < 0 || acceptableRecoveryLag > CHANGELOG_END_OFFSET / 4) {
+ throw new IllegalArgumentException("Unsupported acceptable recovery lag " + acceptableRecoveryLag + ".");
+ }
+ long caughtUpOffset = CHANGELOG_END_OFFSET - acceptableRecoveryLag;
+ long laggingOffset = CHANGELOG_END_OFFSET - 2 * acceptableRecoveryLag - 1;
+
+ Map<String, MemberTaskOffsets> memberTaskOffsets = new HashMap<>();
+ int copyIndex = 0;
+ List<String> memberIds = new ArrayList<>(members.keySet());
+ memberIds.sort(null);
+ for (String memberId : memberIds) {
+ if (lagPicture == LagPicture.NOT_REPORTED) {
+ memberTaskOffsets.put(memberId, MemberTaskOffsets.EMPTY);
+ continue;
+ }
Review Comment:
`NOT_REPORTED` is documented as the post-failover case, where the coordinator's offsets map has no entry for any member, but this builds a map containing an `EMPTY` value for every member. `AssignmentRefinerImpl` uses `getOrDefault`, so decisions match, but the benchmark measures successful lookups instead of the absent-entry lookups used after failover. Omit these entries so this parameter represents the production input it claims to benchmark.
##########
jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmarkUtils.java:
##########
@@ -181,4 +224,275 @@ public static SortedMap<String, ConfiguredSubtopology> createSubtopologyMap(
}
return subtopologyMap;
}
+
+ /**
+ * Creates the members' current assignment as the target assignment with a fraction of the stateful active tasks
+ * still running where they were before the target assignment moved them, which is the work the refiner has to
+ * do. A group whose current assignment equals its target leaves the refiner nothing to decide.
+ * <p>
+ * Every chosen task runs on a member of another process than its target owner, preferring one that holds no copy
+ * of the task, round-robin over that process's members. When every other process holds a standby of the task,
+ * the active and one standby swap places instead, so that the target owner holds the standby. The chosen tasks
+ * are spread evenly over the stateful tasks in task order. The target's standby tasks are kept otherwise.
+ * <p>
+ * A refinement step after the first one finds the warm-up budget spent on migrations already under way, so
+ * inFlightWarmups of the migrating tasks whose target owner holds nothing for them get a warm-up task on it,
+ * spread evenly over those tasks.
+ *
+ * @param members The members, which must span at least two processes if any task migrates.
+ * @param targetAssignment The target assignment, keyed by member ID.
+ * @param subtopologyMap The subtopologies; only the tasks of stateful ones migrate.
+ * @param migratingTaskFraction The fraction of the stateful active tasks not on their target owner yet.
+ * @param inFlightWarmups The number of migrating tasks already warming up on their target owner.
+ *
+ * @return The members, with their current assignment set.
+ */
+ public static Map<String, StreamsGroupMember> divergeAssignment(
+ Map<String, StreamsGroupMember> members,
+ Map<String, TasksTuple> targetAssignment,
+ SortedMap<String, ConfiguredSubtopology> subtopologyMap,
+ double migratingTaskFraction,
+ int inFlightWarmups
+ ) {
+ Map<String, Map<String, Map<Integer, Integer>>> activeTasks = new HashMap<>();
+ Map<String, Map<String, Set<Integer>>> standbyTasks = new HashMap<>();
+ Map<String, Map<String, Set<Integer>>> warmupTasks = new HashMap<>();
+ TreeMap<TaskId, String> statefulActiveOwners = new TreeMap<>();
+ Map<TaskId, SortedSet<String>> standbyHolders = new HashMap<>();
+
+ for (Map.Entry<String, TasksTuple> entry : targetAssignment.entrySet()) {
+ String memberId = entry.getKey();
+ int memberEpoch = members.get(memberId).memberEpoch();
+ entry.getValue().activeTasks().forEach((subtopologyId, partitions) -> partitions.forEach(partition -> {
+ addActive(activeTasks, memberId, new TaskId(subtopologyId, partition), memberEpoch);
+ if (isStateful(subtopologyMap, subtopologyId)) {
+ statefulActiveOwners.put(new TaskId(subtopologyId, partition), memberId);
+ }
+ }));
+ entry.getValue().standbyTasks().forEach((subtopologyId, partitions) -> partitions.forEach(partition -> {
+ addTask(standbyTasks, memberId, new TaskId(subtopologyId, partition));
+ standbyHolders.computeIfAbsent(new TaskId(subtopologyId, partition), id -> new TreeSet<>()).add(memberId);
+ }));
+ }
+
+ List<String> processIds = new ArrayList<>();
+ Map<String, List<String>> membersByProcess = new TreeMap<>();
+ members.values().stream().sorted(Comparator.comparing(StreamsGroupMember::memberId)).forEach(member ->
+ membersByProcess.computeIfAbsent(member.processId(), id -> new ArrayList<>()).add(member.memberId()));
+ processIds.addAll(membersByProcess.keySet());
+ if (migratingTaskFraction > 0 && processIds.size() < 2) {
+ throw new IllegalArgumentException("A task can only migrate between two processes, but the members run in "
+ + processIds.size() + ".");
+ }
+ Map<String, Integer> nextMemberOfProcess = new HashMap<>();
+
+ List<Map.Entry<TaskId, String>> warmupCandidates = new ArrayList<>();
+ int taskIndex = 0;
+ for (Map.Entry<TaskId, String> entry : statefulActiveOwners.entrySet()) {
+ boolean migrating = Math.floor((taskIndex + 1) * migratingTaskFraction) > Math.floor(taskIndex * migratingTaskFraction);
+ taskIndex++;
+ if (!migrating) {
+ continue;
+ }
+
+ TaskId task = entry.getKey();
+ String targetOwner = entry.getValue();
+ int epoch = activeTasks.get(targetOwner).get(task.subtopologyId()).get(task.partition());
+ removeActive(activeTasks, targetOwner, task);
+
+ Optional<String> freeProcess = processWithoutCopy(members, processIds, targetOwner, standbyHolders.get(task));
+ if (freeProcess.isPresent()) {
+ List<String> processMembers = membersByProcess.get(freeProcess.get());
+ int next = nextMemberOfProcess.merge(freeProcess.get(), 1, Integer::sum) - 1;
+ addActive(activeTasks, processMembers.get(next % processMembers.size()), task, epoch);
+ warmupCandidates.add(entry);
+ } else {
+ String standbyHolder = standbyHolders.get(task).first();
+ removeTask(standbyTasks, standbyHolder, task);
+ addActive(activeTasks, standbyHolder, task, epoch);
+ addTask(standbyTasks, targetOwner, task);
+ }
+ }
+
+ int warmups = Math.min(inFlightWarmups, warmupCandidates.size());
+ for (int warmup = 0; warmup < warmups; warmup++) {
+ Map.Entry<TaskId, String> candidate =
+ warmupCandidates.get((int) ((long) warmup * warmupCandidates.size() / warmups));
+ addTask(warmupTasks, candidate.getValue(), candidate.getKey());
+ }
+
+ Map<String, StreamsGroupMember> divergedMembers = new HashMap<>();
+ for (StreamsGroupMember member : members.values()) {
+ String memberId = member.memberId();
+ divergedMembers.put(memberId, new StreamsGroupMember.Builder(member)
+ .setAssignedTasks(new TasksTupleWithEpochs(
+ activeTasks.getOrDefault(memberId, Map.of()),
+ standbyTasks.getOrDefault(memberId, Map.of()),
+ warmupTasks.getOrDefault(memberId, Map.of())
+ ))
+ .build());
+ }
+ return divergedMembers;
+ }
+
+ /**
+ * The first process after the target owner's, in process ID order, that holds no standby of the task.
+ */
+ private static Optional<String> processWithoutCopy(
+ Map<String, StreamsGroupMember> members,
+ List<String> processIds,
+ String targetOwner,
+ Set<String> standbyHolders
+ ) {
+ Set<String> copyProcesses = new HashSet<>();
+ if (standbyHolders != null) {
+ standbyHolders.forEach(holder -> copyProcesses.add(members.get(holder).processId()));
+ }
+
+ int targetProcessIndex = processIds.indexOf(members.get(targetOwner).processId());
+ for (int offset = 1; offset < processIds.size(); offset++) {
+ String processId = processIds.get((targetProcessIndex + offset) % processIds.size());
+ if (!copyProcesses.contains(processId)) {
+ return Optional.of(processId);
+ }
+ }
+ return Optional.empty();
+ }
+
+ /**
+ * Creates the changelog offsets and end offsets the members report in their heartbeats, from which the refiner
+ * tells whether a copy of a task is caught up.
+ * <p>
+ * Every member reports the stateful standby and warm-up tasks of its current assignment, lagging or caught up as
+ * the lag picture says. Under RESTORING_DESTINATIONS, every member that the target assignment moves a stateful
+ * active task to also reports its first stateful active task that stays with it, which reads as restoring.
+ *
+ * @param members The members, with their current assignment set.
+ * @param targetAssignment The target assignment, keyed by member ID.
+ * @param subtopologyMap The subtopologies; only the tasks of stateful ones are reported.
+ * @param lagPicture What the members report.
+ * @param acceptableRecoveryLag The lag at or below which a task counts as caught up.
+ *
+ * @return The reported offsets, per member.
+ */
+ public static Map<String, MemberTaskOffsets> createMemberTaskOffsets(
+ Map<String, StreamsGroupMember> members,
+ Map<String, TasksTuple> targetAssignment,
+ SortedMap<String, ConfiguredSubtopology> subtopologyMap,
+ LagPicture lagPicture,
+ long acceptableRecoveryLag
+ ) {
+ if (acceptableRecoveryLag < 0 || acceptableRecoveryLag > CHANGELOG_END_OFFSET / 4) {
+ throw new IllegalArgumentException("Unsupported acceptable recovery lag " + acceptableRecoveryLag + ".");
+ }
+ long caughtUpOffset = CHANGELOG_END_OFFSET - acceptableRecoveryLag;
+ long laggingOffset = CHANGELOG_END_OFFSET - 2 * acceptableRecoveryLag - 1;
+
+ Map<String, MemberTaskOffsets> memberTaskOffsets = new HashMap<>();
+ int copyIndex = 0;
+ List<String> memberIds = new ArrayList<>(members.keySet());
+ memberIds.sort(null);
+ for (String memberId : memberIds) {
+ if (lagPicture == LagPicture.NOT_REPORTED) {
+ memberTaskOffsets.put(memberId, MemberTaskOffsets.EMPTY);
+ continue;
+ }
+
+ TasksTupleWithEpochs assignedTasks = members.get(memberId).assignedTasks();
+ Map<String, Map<Integer, Long>> taskOffsets = new HashMap<>();
+ Map<String, Map<Integer, Long>> taskEndOffsets = new HashMap<>();
+
+ SortedSet<TaskId> copies = statefulTasks(assignedTasks.standbyTasks(), subtopologyMap);
+ copies.addAll(statefulTasks(assignedTasks.warmupTasks(), subtopologyMap));
+ for (TaskId task : copies) {
+ LagPicture reported = lagPicture == LagPicture.MIXED ? MIXED_TURNS[copyIndex++ % MIXED_TURNS.length] : lagPicture;
+ if (reported != LagPicture.NOT_REPORTED) {
+ reportOffsets(taskOffsets, taskEndOffsets, task,
+ reported == LagPicture.ALL_CAUGHT_UP ? caughtUpOffset : laggingOffset);
+ }
+ }
+
+ if (lagPicture == LagPicture.RESTORING_DESTINATIONS) {
+ SortedSet<TaskId> currentActives = statefulTasks(assignedTasks.activeTasks(), subtopologyMap);
+ SortedSet<TaskId> targetActives = statefulTasks(
+ targetAssignment.getOrDefault(memberId, TasksTuple.EMPTY).activeTasks(), subtopologyMap);
+ boolean destination = !currentActives.containsAll(targetActives);
+ currentActives.retainAll(targetActives);
+ if (destination && !currentActives.isEmpty()) {
+ reportOffsets(taskOffsets, taskEndOffsets, currentActives.first(), laggingOffset);
+ }
Review Comment:
This guard means `RESTORING_DESTINATIONS` does not actually create its advertised condition for part of the default grid. With `memberCount=1000`, `subtopologyCount=10`, and `partitionCount=100`, there are exactly 1,000 active tasks and the sticky assignor's active quota is one per member. Once a selected stateful task is moved away, its target owner has no active target task left, so `currentActives` is empty and no restoring active is reported. Those rows therefore collapse to the `ALL_LAGGING` scenario instead of exercising the refiner's stalled-destination path. The fixture or parameter grid should guarantee a retained stateful active on every migration destination.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-21026: Add JMH benchmark for the streams assignment refiner [kafka]
Posted by "suzhiking (via GitHub)" <gi...@apache.org>.
suzhiking commented on code in PR #23616:
URL: https://github.com/apache/kafka/pull/23616#discussion_r4215450870
##########
jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignmentRefinerBenchmark.java:
##########
@@ -0,0 +1,175 @@
+/*
+ * 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.jmh.assignor;
+
+import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
+import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment;
+import org.apache.kafka.coordinator.group.streams.AssignmentRefiner;
+import org.apache.kafka.coordinator.group.streams.AssignmentRefinerImpl;
+import org.apache.kafka.coordinator.group.streams.MemberTaskOffsets;
+import org.apache.kafka.coordinator.group.streams.NoOpAssignmentRefiner;
+import org.apache.kafka.coordinator.group.streams.StreamsGroupMember;
+import org.apache.kafka.coordinator.group.streams.TasksTuple;
+import org.apache.kafka.coordinator.group.streams.TopologyMetadata;
+import org.apache.kafka.coordinator.group.streams.assignor.AssignmentConfigsImpl;
+import org.apache.kafka.coordinator.group.streams.assignor.StickyTaskAssignor;
+import org.apache.kafka.coordinator.group.streams.topics.ConfiguredSubtopology;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.SortedMap;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Measures one derivation of the streams group's intermediate assignment: what the group coordinator runs when it
+ * refines the target assignment into the warm-up steps the members converge through, together with the check it
+ * applies to the result before using it.
+ * <p>
+ * The group is set up as a refinement step finds it: the target assignment comes from the sticky assignor, a fraction
+ * of the stateful active tasks still runs on a member of another process, the warm-up budget is spent on migrations
+ * already under way, and the members report their restore progress as the lag picture says.
+ * <p>
+ * The full grid is large; narrow it with -p, for example -p memberCount=1000 -p lagPicture=MIXED.
+ */
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 5, time = 1)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+public class AssignmentRefinerBenchmark {
+
+ /**
+ * The default of acceptable.recovery.lag.
+ */
+ private static final long ACCEPTABLE_RECOVERY_LAG = 10_000L;
+
+ @Param({"100", "1000"})
+ private int memberCount;
+
+ @Param({"1", "50"})
+ private int membersPerProcess;
+
+ @Param({"10", "100"})
+ private int subtopologyCount;
+
+ @Param({"100"})
+ private int partitionCount;
+
+ @Param({"0", "1"})
+ private int standbyReplicas;
+
+ @Param({"2", "100"})
+ private int numWarmupReplicas;
+
+ @Param({"0.0", "0.01", "0.5"})
+ private double migratingTaskFraction;
+
+ @Param({"ALL_CAUGHT_UP", "ALL_LAGGING", "NOT_REPORTED", "MIXED", "RESTORING_DESTINATIONS"})
+ private StreamsAssignorBenchmarkUtils.LagPicture lagPicture;
+
+ private final AssignmentRefiner noOpRefiner = new NoOpAssignmentRefiner();
+
+ private final AssignmentRefiner refinerImpl = new AssignmentRefinerImpl();
+
+ private Map<String, StreamsGroupMember> members;
+
+ private Map<String, TasksTuple> targetAssignment;
+
+ private Map<String, MemberTaskOffsets> taskOffsets;
+
+ private SortedMap<String, ConfiguredSubtopology> subtopologyMap;
+
+ @Setup(Level.Trial)
+ public void setup() {
+ List<String> allTopicNames = AssignorBenchmarkUtils.createTopicNames(subtopologyCount);
+ subtopologyMap = StreamsAssignorBenchmarkUtils.createSubtopologyMap(partitionCount, allTopicNames);
+ CoordinatorMetadataImage metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionCount);
+
+ Map<String, StreamsGroupMember> newMembers = StreamsAssignorBenchmarkUtils.createStreamsMembers(memberCount, membersPerProcess);
+ GroupAssignment groupAssignment = new StickyTaskAssignor().assign(
+ StreamsAssignorBenchmarkUtils.createGroupSpec(
+ newMembers,
+ AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(standbyReplicas),
+ Map.of()
+ ),
+ new TopologyMetadata(metadataImage, subtopologyMap)
+ );
+
+ targetAssignment = new HashMap<>();
+ groupAssignment.members().forEach((memberId, memberAssignment) -> targetAssignment.put(
+ memberId,
+ new TasksTuple(memberAssignment.activeTasks(), memberAssignment.standbyTasks(), Map.of())
+ ));
+
+ members = StreamsAssignorBenchmarkUtils.divergeAssignment(
+ newMembers,
+ targetAssignment,
+ subtopologyMap,
+ migratingTaskFraction,
+ numWarmupReplicas
+ );
+ taskOffsets = StreamsAssignorBenchmarkUtils.createMemberTaskOffsets(
+ members,
+ targetAssignment,
+ subtopologyMap,
+ lagPicture,
+ ACCEPTABLE_RECOVERY_LAG
+ );
+ }
+
+ @Benchmark
+ @Threads(1)
+ public void noOpRefiner(Blackhole blackhole) {
+ refine(noOpRefiner, blackhole);
+ }
+
+ @Benchmark
+ @Threads(1)
+ public void refinerImpl(Blackhole blackhole) {
+ refine(refinerImpl, blackhole);
+ }
+
+ private void refine(AssignmentRefiner refiner, Blackhole blackhole) {
Review Comment:
I think they are mostly correct with 2 nits:
1. It grows at least O(N^2 + N * S) where N is member count and S is stateful tasks count because refiner step iterate members and active tasks
2. This benchmark only measures refiner's speed but skip is in coordinator itself, so it could not measure how much time skip would save. That needs a separate heartbeat benchmark and I'll add it as follow-up
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-21026: Add JMH benchmark for the streams assignment refiner [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru merged PR #23616:
URL: https://github.com/apache/kafka/pull/23616
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org