You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by lu...@apache.org on 2026/10/09 08:45:48 UTC
(kafka) branch trunk updated: KAFKA-21026: Add JMH benchmark for the streams assignment refiner (#23616)
This is an automated email from the ASF dual-hosted git repository.
lucasbru pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new 4151a818477 KAFKA-21026: Add JMH benchmark for the streams assignment refiner (#23616)
4151a818477 is described below
commit 4151a8184774464ae49f029baa9dad51c9c780d5
Author: Gavin Wang <z2...@uwaterloo.ca>
AuthorDate: Fri Oct 9 04:45:39 2026 -0400
KAFKA-21026: Add JMH benchmark for the streams assignment refiner (#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.
Reviewers: Lucas Brutschy <lb...@confluent.io>
---
.../jmh/assignor/AssignmentRefinerBenchmark.java | 175 ++++++++++++
.../assignor/StreamsAssignorBenchmarkUtils.java | 315 +++++++++++++++++++++
2 files changed, 490 insertions(+)
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignmentRefinerBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignmentRefinerBenchmark.java
new file mode 100644
index 00000000000..0ad76cf18c2
--- /dev/null
+++ b/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) {
+ Map<String, TasksTuple> refinedAssignment = refiner.refine(
+ members,
+ targetAssignment,
+ taskOffsets,
+ subtopologyMap,
+ numWarmupReplicas,
+ ACCEPTABLE_RECOVERY_LAG
+ );
+ blackhole.consume(AssignmentRefiner.preservesActiveTaskCount(targetAssignment, refinedAssignment));
+ blackhole.consume(refinedAssignment);
+ }
+}
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmarkUtils.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmarkUtils.java
index 6f85a5d25f9..2db14c7e415 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmarkUtils.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmarkUtils.java
@@ -18,23 +18,68 @@ package org.apache.kafka.jmh.assignor;
import org.apache.kafka.coordinator.group.api.streams.assignor.AssignmentConfigs;
import org.apache.kafka.coordinator.group.api.streams.assignor.GroupSpec;
+import org.apache.kafka.coordinator.group.streams.MemberTaskOffsets;
import org.apache.kafka.coordinator.group.streams.StreamsGroupMember;
+import org.apache.kafka.coordinator.group.streams.TasksTuple;
+import org.apache.kafka.coordinator.group.streams.TasksTupleWithEpochs;
import org.apache.kafka.coordinator.group.streams.assignor.GroupSpecImpl;
import org.apache.kafka.coordinator.group.streams.assignor.MemberMetadataAndStateImpl;
+import org.apache.kafka.coordinator.group.streams.assignor.TaskId;
import org.apache.kafka.coordinator.group.streams.topics.ConfiguredInternalTopic;
import org.apache.kafka.coordinator.group.streams.topics.ConfiguredSubtopology;
+import java.util.ArrayList;
+import java.util.Comparator;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.SortedMap;
+import java.util.SortedSet;
import java.util.TreeMap;
+import java.util.TreeSet;
public class StreamsAssignorBenchmarkUtils {
+ /**
+ * The end offset every reported changelog is at. Far enough from zero that a lagging offset stays positive for
+ * any realistic acceptable recovery lag.
+ */
+ private static final long CHANGELOG_END_OFFSET = 1L << 40;
+
+ /**
+ * What the members report about the state they restore, which is what the refiner's decisions hinge on. A
+ * member reports offsets for its standby and warm-up tasks, and for an active task only while it is still
+ * restoring it.
+ */
+ public enum LagPicture {
+ /** Every standby and warm-up task is within the acceptable recovery lag. */
+ ALL_CAUGHT_UP,
+ /** Every standby and warm-up task has restored part of its state, but lags beyond the acceptable recovery lag. */
+ ALL_LAGGING,
+ /** No member reports any offsets, as right after a coordinator failover. */
+ NOT_REPORTED,
+ /** Standby and warm-up tasks take turns being caught up, lagging, and not reported. */
+ MIXED,
+ /**
+ * Like ALL_LAGGING, and every member that a task migrates to is also restoring one of its own active tasks,
+ * which stalls every warm-up task on that member. Only a member that keeps a stateful active task of its target
+ * can report one as restoring, so where each member has a single active task, as with 1000 members and 10
+ * subtopologies of 100 partitions, this is the same as ALL_LAGGING.
+ */
+ RESTORING_DESTINATIONS
+ }
+
+ /**
+ * The turns the copies of a MIXED lag picture take, in order.
+ */
+ private static final LagPicture[] MIXED_TURNS = {
+ LagPicture.ALL_CAUGHT_UP, LagPicture.ALL_LAGGING, LagPicture.NOT_REPORTED
+ };
+
/**
* Creates a GroupSpec from the given StreamsGroupMembers.
*
@@ -181,4 +226,274 @@ public class StreamsAssignorBenchmarkUtils {
}
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) {
+ 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);
+ }
+ }
+
+ memberTaskOffsets.put(memberId, new MemberTaskOffsets(taskOffsets, taskEndOffsets));
+ }
+
+ return memberTaskOffsets;
+ }
+
+ private static void reportOffsets(
+ Map<String, Map<Integer, Long>> taskOffsets,
+ Map<String, Map<Integer, Long>> taskEndOffsets,
+ TaskId task,
+ long offset
+ ) {
+ taskOffsets.computeIfAbsent(task.subtopologyId(), id -> new HashMap<>()).put(task.partition(), offset);
+ taskEndOffsets.computeIfAbsent(task.subtopologyId(), id -> new HashMap<>()).put(task.partition(), CHANGELOG_END_OFFSET);
+ }
+
+ private static SortedSet<TaskId> statefulTasks(
+ Map<String, Set<Integer>> tasks,
+ SortedMap<String, ConfiguredSubtopology> subtopologyMap
+ ) {
+ SortedSet<TaskId> statefulTasks = new TreeSet<>();
+ tasks.forEach((subtopologyId, partitions) -> {
+ if (isStateful(subtopologyMap, subtopologyId)) {
+ partitions.forEach(partition -> statefulTasks.add(new TaskId(subtopologyId, partition)));
+ }
+ });
+ return statefulTasks;
+ }
+
+ private static boolean isStateful(SortedMap<String, ConfiguredSubtopology> subtopologyMap, String subtopologyId) {
+ ConfiguredSubtopology subtopology = subtopologyMap.get(subtopologyId);
+ return subtopology != null && !subtopology.stateChangelogTopics().isEmpty();
+ }
+
+ private static void addActive(
+ Map<String, Map<String, Map<Integer, Integer>>> activeTasks,
+ String memberId,
+ TaskId task,
+ int epoch
+ ) {
+ activeTasks.computeIfAbsent(memberId, id -> new HashMap<>())
+ .computeIfAbsent(task.subtopologyId(), id -> new HashMap<>())
+ .put(task.partition(), epoch);
+ }
+
+ private static void removeActive(
+ Map<String, Map<String, Map<Integer, Integer>>> activeTasks,
+ String memberId,
+ TaskId task
+ ) {
+ Map<String, Map<Integer, Integer>> memberTasks = activeTasks.get(memberId);
+ Map<Integer, Integer> partitions = memberTasks.get(task.subtopologyId());
+ partitions.remove(task.partition());
+ if (partitions.isEmpty()) {
+ memberTasks.remove(task.subtopologyId());
+ }
+ }
+
+ private static void addTask(Map<String, Map<String, Set<Integer>>> tasks, String memberId, TaskId task) {
+ tasks.computeIfAbsent(memberId, id -> new HashMap<>())
+ .computeIfAbsent(task.subtopologyId(), id -> new HashSet<>())
+ .add(task.partition());
+ }
+
+ private static void removeTask(Map<String, Map<String, Set<Integer>>> tasks, String memberId, TaskId task) {
+ Map<String, Set<Integer>> memberTasks = tasks.get(memberId);
+ Set<Integer> partitions = memberTasks.get(task.subtopologyId());
+ partitions.remove(task.partition());
+ if (partitions.isEmpty()) {
+ memberTasks.remove(task.subtopologyId());
+ }
+ }
}