mjsax commented on code in PR #23620: URL: https://github.com/apache/kafka/pull/23620#discussion_r4393480247
########## streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/StreamsGroupWarmupTaskLifecycleIntegrationTest.java: ########## @@ -0,0 +1,724 @@ +/* + * 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.streams.integration; + +import org.apache.kafka.clients.admin.Admin; +import org.apache.kafka.clients.admin.AdminClientConfig; +import org.apache.kafka.clients.admin.OffsetSpec; +import org.apache.kafka.clients.admin.StreamsGroupDescription; +import org.apache.kafka.clients.admin.StreamsGroupMemberAssignment; +import org.apache.kafka.clients.admin.StreamsGroupMemberDescription; +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.consumer.ConsumerRecords; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.clients.producer.Producer; +import org.apache.kafka.common.GroupState; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.config.TopicConfig; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.common.serialization.StringSerializer; +import org.apache.kafka.common.utils.Utils; +import org.apache.kafka.coordinator.group.GroupCoordinatorConfig; +import org.apache.kafka.coordinator.group.streams.AssignmentRefinerImpl; +import org.apache.kafka.streams.GroupProtocol; +import org.apache.kafka.streams.KafkaClientSupplier; +import org.apache.kafka.streams.KafkaStreams; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.StreamsBuilder; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster; +import org.apache.kafka.streams.integration.utils.IntegrationTestUtils; +import org.apache.kafka.streams.kstream.Materialized; +import org.apache.kafka.streams.processor.StandbyUpdateListener; +import org.apache.kafka.streams.processor.StateRestoreListener; +import org.apache.kafka.streams.processor.TaskId; +import org.apache.kafka.streams.processor.internals.DefaultKafkaClientSupplier; +import org.apache.kafka.streams.state.Stores; +import org.apache.kafka.test.TestUtils; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInfo; +import org.junit.jupiter.api.Timeout; + +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.apache.kafka.common.utils.Utils.mkEntry; +import static org.apache.kafka.common.utils.Utils.mkMap; +import static org.apache.kafka.common.utils.Utils.mkObjectProperties; +import static org.apache.kafka.common.utils.Utils.mkProperties; +import static org.apache.kafka.streams.utils.TestUtils.safeUniqueTestName; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Covers the warm-up task lifecycle under the streams group protocol (KIP-1071) end to end, with real Kafka Streams + * clients against the broker-side {@link AssignmentRefinerImpl}: the warm-up budget, a warm-up task given up because + * the target assignment no longer wants it, the offset reporting that drives promotion, the promotion gate on + * {@code acceptable.recovery.lag}, and warm-up tasks being disabled. The scale-out case itself is covered by + * {@link StreamsGroupWarmupTaskIntegrationTest}. + * + * <p>All stores are in memory, so a task that is closed rather than recycled has to restore its changelog from + * scratch, which the restore listeners record. Where a test has to hold a warm-up task behind, a {@link RestoreGate} + * withholds the records of the instance's restore consumer. A gate that is closed from the start would leave the + * warm-up task with nothing restored and hence no offset to report, so a gate lets a small budget of records through + * first. + */ +@Timeout(600) +@Tag("integration") +public class StreamsGroupWarmupTaskLifecycleIntegrationTest { + private static final Properties BROKER_CONFIG = mkProperties(mkMap( + mkEntry(GroupCoordinatorConfig.STREAMS_GROUP_ASSIGNMENT_REFINER_CLASS_CONFIG, AssignmentRefinerImpl.class.getName()), + mkEntry(GroupCoordinatorConfig.STREAMS_GROUP_ACCEPTABLE_RECOVERY_LAG_CONFIG, "0") + )); + + public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(1, BROKER_CONFIG); + + private static final String SUBTOPOLOGY_ID = "0"; + private static final int HEARTBEAT_INTERVAL_MS = 500; + private static final long WAIT_MS = 60_000L; + private static final int RECORDS_PER_PARTITION = 200; + private static final int RESTORE_BUDGET = 20; + + private static Admin admin; + + private final List<Instance> instances = new ArrayList<>(); + private String appId; + private String inputTopic; + private String storeName; + private String changelogTopic; + private int numPartitions; + + @BeforeAll + public static void startCluster() throws IOException { + CLUSTER.start(); + admin = Admin.create(Map.of(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers())); + } + + @AfterAll + public static void closeCluster() { + admin.close(); + CLUSTER.stop(); + } + + @AfterEach + public void tearDown() { + for (final Instance instance : instances) { + instance.gate.open(); + instance.streams.close(Duration.ofSeconds(30)); + } + instances.clear(); + } + + @Test + public void shouldReportTheOffsetsOfAWarmupTaskAsSoonAsItIsCaughtUp(final TestInfo testInfo) throws Exception { + setUp(testInfo, 2); + final Instance a = startOwnerOfAllTasks(); + // Far longer than the test waits, so the interval never elapses and cannot be what reports the offsets. + CLUSTER.setGroupTaskOffsetInterval(appId, 300_000); + + final Instance b = startInstance("b", RESTORE_BUDGET); + final int task = awaitStagedMigration(a, b); + + // The warm-up task catches up only now, long after the offsets it reported when it was assigned. Only the + // client reporting a caught-up warm-up task on the next heartbeat lets the broker promote it in time. + b.gate.open(); + waitForGroup("the warm-up task " + task + " to be promoted to active on b", group -> + activeTasks(group, b).equals(Set.of(task)) && warmupTasks(group, b).isEmpty()); + } + + @Test + public void shouldGateWarmupPromotionOnAcceptableRecoveryLag(final TestInfo testInfo) throws Exception { + setUp(testInfo, 2); + final long acceptableRecoveryLag = RESTORE_BUDGET / 4; + CLUSTER.setGroupAcceptableRecoveryLag(appId, acceptableRecoveryLag); + final Instance a = startOwnerOfAllTasks(); + + final Instance b = startInstance("b", RESTORE_BUDGET); + final int task = awaitStagedMigration(a, b); + // The client reports Long.MAX_VALUE for an offset it does not know yet, which the broker treats as an unknown + // lag. The task would be withheld for that reason alone, so wait for both offsets to be real. + waitForGroup("the broker to learn how far behind the warm-up task " + task + " is", group -> { + final StreamsGroupMemberDescription member = member(group, b); + if (member == null) { + return false; + } + final long offset = reportedOffset(member.taskOffsets(), task); + final long endOffset = reportedOffset(member.taskEndOffsets(), task); + return offset >= RESTORE_BUDGET && offset != Long.MAX_VALUE + && endOffset != -1L && endOffset != Long.MAX_VALUE + && endOffset - offset > acceptableRecoveryLag; + }); + + // The lag is known and above the threshold, so the task stays where it is. + assertHolds("the warm-up task " + task + " to be withheld from b", Duration.ofSeconds(5), group -> + warmupTasks(group, b).equals(Set.of(task)) && activeTasks(group, a).contains(task)); + + // The gate still holds the warm-up task at the same lag, so it is only the raised threshold that promotes it. + CLUSTER.setGroupAcceptableRecoveryLag(appId, 1_000_000L); + waitForGroup("the warm-up task " + task + " to be promoted to active on b", group -> + activeTasks(group, b).equals(Set.of(task)) && !activeTasks(group, a).contains(task)); + } + + @Test + public void shouldWarmUpOneMigrationAtATimeWhenTwoMembersJoinAtOnce(final TestInfo testInfo) throws Exception { + setUp(testInfo, 3); + CLUSTER.setGroupNumWarmupReplicas(appId, 1); + final Instance a = startOwnerOfAllTasks(); + + final Instance b = startInstance("b", RESTORE_BUDGET); + final Instance c = startInstance("c", RESTORE_BUDGET); + // Both joiners are to take over a task, but with a budget of one warm-up task only one of the migrations is + // warmed up, and the other waits with its task still running on a. + waitForGroup("one of the two migrations to be warmed up, and the other to wait", group -> isStable(group, 3) + && targetActiveTasks(group, b).size() == 1 + && targetActiveTasks(group, c).size() == 1 + && activeTasks(group, a).size() == 3 + && totalWarmupTasks(group) == 1); + assertHolds("the waiting migration not to be warmed up while the budget is spent", Duration.ofSeconds(5), group -> + totalWarmupTasks(group) == 1 && activeTasks(group, a).size() == 3); + + b.gate.open(); + c.gate.open(); + final AtomicInteger maxWarmupTasksWhenStable = new AtomicInteger(); + waitForGroup("both migrations to complete", group -> { + if (isStable(group, 3)) { + maxWarmupTasksWhenStable.accumulateAndGet(totalWarmupTasks(group), Math::max); + } + return isStable(group, 3) + && activeTasks(group, a).size() == 1 + && activeTasks(group, b).equals(targetActiveTasks(group, b)) + && activeTasks(group, c).equals(targetActiveTasks(group, c)) + && totalWarmupTasks(group) == 0; + }); + assertTrue(maxWarmupTasksWhenStable.get() <= 1, "the group ran " + maxWarmupTasksWhenStable.get() + " warm-up tasks at once"); + // Neither joiner restored an active task, so both took over a warmed-up task: the waiting migration was warmed + // up once the first one freed the budget, rather than moving cold. The group is stable as soon as the joiners + // own their tasks, before a cold restore would even begin, so wait for each restore to end before checking. + final StreamsGroupDescription converged = describeGroup(); + awaitNoActiveRestore(b, single(activeTasks(converged, b))); + awaitNoActiveRestore(c, single(activeTasks(converged, c))); + } + + @Test + public void shouldGiveUpAWarmupTaskTheTargetAssignmentNoLongerWantsOnScaleIn(final TestInfo testInfo) throws Exception { + setUp(testInfo, 4); + final Instance a = startInstance("a"); + final Instance b = startInstance("b"); + waitForGroup("each instance to own two active tasks", group -> isStable(group, 2) + && activeTasks(group, a).size() == 2 && activeTasks(group, b).size() == 2 && totalWarmupTasks(group) == 0); + produceAndAwaitChangelog(); + awaitNoActiveTaskReportedAsRestoring(); + + final Instance c = startInstance("c", RESTORE_BUDGET); + final AtomicReference<Integer> warmupTask = new AtomicReference<>(); + waitForGroup("c to warm up one task", group -> { + if (!isStable(group, 3) || !activeTasks(group, c).isEmpty() || warmupTasks(group, c).size() != 1) { + return false; + } + warmupTask.set(single(warmupTasks(group, c))); + return true; + }); + final int task = warmupTask.get(); + final StreamsGroupDescription beforeScaleIn = describeGroup(); + final Instance owner = activeTasks(beforeScaleIn, a).contains(task) ? a : b; + final Instance leaver = owner == a ? b : a; + final Set<Integer> tasksOfOwner = activeTasks(beforeScaleIn, owner); + final Set<Integer> tasksOfLeaver = activeTasks(beforeScaleIn, leaver); + final int restoresByOwner = owner.restoreEnds(task); + + // With the leaver gone, the owner is no longer over its share and keeps the task it was to hand to c, while c + // takes over the leaver's tasks instead. So the warm-up task on c is given up rather than promoted. + leaver.streams.close(Duration.ofSeconds(30)); + instances.remove(leaver); + waitForGroup("c to give up the warm-up task " + task + " and take over the leaver's tasks", group -> isStable(group, 2) + && targetActiveTasks(group, c).equals(tasksOfLeaver) + && activeTasks(group, c).equals(tasksOfLeaver) + && activeTasks(group, owner).equals(tasksOfOwner) + && totalWarmupTasks(group) == 0); + // Observed through the standby update listener rather than the thread metadata, which a stream thread only + // refreshes once all its active tasks are restored -- and c's gate holds back the restore of its new tasks. + TestUtils.waitForCondition( + () -> c.standbySuspensionsByTask.containsKey(task), + WAIT_MS, + () -> "c never closed the warm-up task " + task + ); + assertEquals(StandbyUpdateListener.SuspendReason.MIGRATED, c.standbySuspensionsByTask.get(task), + "c should have given up the warm-up task " + task + " rather than promoted it"); + assertEquals(restoresByOwner, owner.restoreEnds(task), "the owner should have kept running " + task + " throughout"); + } + + @Test + public void shouldMoveTasksColdWhenWarmupTasksAreDisabled(final TestInfo testInfo) throws Exception { + setUp(testInfo, 2); + CLUSTER.setGroupNumWarmupReplicas(appId, 0); + final Instance a = startOwnerOfAllTasks(); + + final Instance b = startInstance("b"); + final AtomicBoolean sawWarmupTask = new AtomicBoolean(); + waitForGroup("b to take over a task", group -> { + sawWarmupTask.compareAndSet(false, totalWarmupTasks(group) > 0); + return isStable(group, 2) && activeTasks(group, a).size() == 1 && activeTasks(group, b).size() == 1; + }); + + assertFalse(sawWarmupTask.get(), "no warm-up task should be assigned with warm-up tasks disabled"); + final int movedTask = single(activeTasks(describeGroup(), b)); + awaitRestoreEnd(b, movedTask, 1); + assertEquals(changelogSizes().get(movedTask), b.totalRestoredByTask.get(movedTask), + "b should have taken over " + movedTask + " cold, restoring all of it from the changelog"); + } + + /** + * Waits until the broker no longer takes any active task to be restoring. A client reports the changelog offsets of + * an active task while it restores it, but only reports that the restore has finished with its next offset report, + * which may be up to the task offset interval later. Until then the refiner sees no caught-up state on the owner to + * protect and moves the task without warming it up. + */ + private void awaitNoActiveTaskReportedAsRestoring() throws InterruptedException { Review Comment: While we cannot fully get rid of this one, it still surfaces a real issue. We should send a new task-(end-)offset report right away when an active-restoring becomes hot, and when restoration completes, to at least reduce the race condition window to max one heartbeat.... Follow up PR: https://github.com/apache/kafka/pull/23987 -- 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]
