From 382ae5a40c43e7c087e1a689f4318fe5f8f85f92 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Tue, 15 Feb 2022 08:27:36 +1300 Subject: [PATCH 1/9] [KAFKA-13600] Unit test to capture not assigning to caught up task if all other replicas are behind --- .../HighAvailabilityTaskAssignorTest.java | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java index a2d4716b45870..e878b13631658 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java @@ -415,6 +415,31 @@ public void shouldComputeNewAssignmentIfActiveTasksWasNotOnCaughtUpClient() { assertBalancedTasks(clientStates); } + @Test + public void shouldAssignToMostCaughtUpIfActiveTasksWasNotOnCaughtUpClient() { + final Set allTasks = mkSet(TASK_0_0); + final Set statefulTasks = mkSet(TASK_0_0); + final ClientState client1 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, Long.MAX_VALUE), 1); + final ClientState client2 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 500L), 1); + final Map clientStates = mkMap( + mkEntry(UUID_1, client1), + mkEntry(UUID_2, client2) + ); + + final boolean probingRebalanceNeeded = + new HighAvailabilityTaskAssignor().assign(clientStates, allTasks, statefulTasks, configWithStandbys); + + assertThat(clientStates.get(UUID_1).activeTasks(), is(emptySet())); + assertThat(clientStates.get(UUID_2).activeTasks(), is(singleton(TASK_0_0))); + // we'll warm up task 0_0 on client1 because it's first in sorted order, + // although this isn't an optimal convergence + assertThat(probingRebalanceNeeded, is(true)); + assertValidAssignment(0, 1, allTasks, emptySet(), clientStates, new StringBuilder()); + assertBalancedActiveAssignment(clientStates, new StringBuilder()); + assertBalancedStatefulAssignment(allTasks, clientStates, new StringBuilder()); + assertBalancedTasks(clientStates); + } + @Test public void shouldAssignStandbysForStatefulTasks() { final Set allTasks = mkSet(TASK_0_0, TASK_0_1); From 6cec06cf14a67e1d880e7d25c388b79e9007ce82 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Tue, 15 Feb 2022 22:32:37 +1300 Subject: [PATCH 2/9] [KAFKA-13600] Fall back to most caught up client if a caught up client can't be found --- .../HighAvailabilityTaskAssignor.java | 19 ++++ .../internals/assignment/TaskMovement.java | 101 +++++++++++++----- .../HighAvailabilityTaskAssignorTest.java | 18 ++-- .../assignment/TaskMovementTest.java | 79 ++++++++++---- 4 files changed, 163 insertions(+), 54 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java index f6464f8ac3117..25a3e8dd4a489 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java @@ -22,8 +22,11 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.ArrayList; +import java.util.Comparator; import java.util.HashMap; import java.util.Iterator; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.SortedMap; @@ -67,6 +70,8 @@ public boolean assign(final Map clients, configs.acceptableRecoveryLag ); + final Map> tasksToClientByLag = tasksToClientByLag(statefulTasks, clientStates); + // We temporarily need to know which standby tasks were intended as warmups // for active tasks, so that we don't move them (again) when we plan standby // task movements. We can then immediately treat warmups exactly the same as @@ -76,6 +81,7 @@ public boolean assign(final Map clients, final int neededActiveTaskMovements = assignActiveTaskMovements( tasksToCaughtUpClients, + tasksToClientByLag, clientStates, warmups, remainingWarmupReplicas @@ -83,6 +89,7 @@ public boolean assign(final Map clients, final int neededStandbyTaskMovements = assignStandbyTaskMovements( tasksToCaughtUpClients, + tasksToClientByLag, clientStates, remainingWarmupReplicas, warmups @@ -252,6 +259,18 @@ private static Map> tasksToCaughtUpClients(final Set> tasksToClientByLag(final Set statefulTasks, + final Map clientStates) { + final Map> tasksToClientByLag = new HashMap<>(); + for (final TaskId task : statefulTasks) { + final List> clientLag = new ArrayList<>(clientStates.entrySet()); + clientLag.sort(Comparator.>comparingLong( + a -> a.getValue().lagFor(task)).thenComparing(Map.Entry::getKey)); + tasksToClientByLag.put(task, clientLag.stream().map(Map.Entry::getKey).collect(Collectors.toList())); + } + return tasksToClientByLag; + } + private static boolean unbounded(final long acceptableRecoveryLag) { return acceptableRecoveryLag == Long.MAX_VALUE; } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java index cbfa3daaedf90..c6cc79e306748 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java @@ -20,6 +20,7 @@ import java.util.Collections; import java.util.Comparator; +import java.util.List; import java.util.Map; import java.util.PriorityQueue; import java.util.Queue; @@ -29,6 +30,7 @@ import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiFunction; +import java.util.function.Function; import static java.util.Arrays.asList; import static java.util.Objects.requireNonNull; @@ -42,10 +44,6 @@ private TaskMovement(final TaskId task, final UUID destination, final SortedSet< this.task = task; this.destination = destination; this.caughtUpClients = caughtUpClients; - - if (caughtUpClients == null || caughtUpClients.isEmpty()) { - throw new IllegalStateException("Should not attempt to move a task if no caught up clients exist"); - } } private TaskId task() { @@ -56,25 +54,34 @@ private int numCaughtUpClients() { return caughtUpClients.size(); } - private static boolean taskIsNotCaughtUpOnClientAndOtherCaughtUpClientsExist(final TaskId task, - final UUID client, - final Map> tasksToCaughtUpClients) { - return !taskIsCaughtUpOnClientOrNoCaughtUpClientsExist(task, client, tasksToCaughtUpClients); + private static boolean taskIsNotCaughtUpOnClientAndOtherMoreCaughtUpClientsExist(final TaskId task, + final UUID client, + final Map clientStates, + final Map> tasksToCaughtUpClients, + final Map> tasksToClientByLag) { + final List taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized map"); + if (taskIsCaughtUpOnClient(task, client, tasksToCaughtUpClients)) { + return false; + } + final long mostCaughtUpLag = clientStates.get(taskClients.get(0)).lagFor(task); + final long clientLag = clientStates.get(client).lagFor(task); + return mostCaughtUpLag < clientLag; } - private static boolean taskIsCaughtUpOnClientOrNoCaughtUpClientsExist(final TaskId task, - final UUID client, - final Map> tasksToCaughtUpClients) { + private static boolean taskIsCaughtUpOnClient(final TaskId task, + final UUID client, + final Map> tasksToCaughtUpClients) { final Set caughtUpClients = requireNonNull(tasksToCaughtUpClients.get(task), "uninitialized set"); - return caughtUpClients.isEmpty() || caughtUpClients.contains(client); + return caughtUpClients.contains(client); } static int assignActiveTaskMovements(final Map> tasksToCaughtUpClients, + final Map> tasksToClientByLag, final Map clientStates, final Map> warmups, final AtomicInteger remainingWarmupReplicas) { final BiFunction caughtUpPredicate = - (client, task) -> taskIsCaughtUpOnClientOrNoCaughtUpClientsExist(task, client, tasksToCaughtUpClients); + (client, task) -> taskIsCaughtUpOnClient(task, client, tasksToCaughtUpClients); final ConstrainedPrioritySet caughtUpClientsByTaskLoad = new ConstrainedPrioritySet( caughtUpPredicate, @@ -89,10 +96,10 @@ static int assignActiveTaskMovements(final Map> tasksToC final UUID client = clientStateEntry.getKey(); final ClientState state = clientStateEntry.getValue(); for (final TaskId task : state.activeTasks()) { - // if the desired client is not caught up, and there is another client that _is_ caught up, then - // we schedule a movement, so we can move the active task to the caught-up client. We'll try to + // if the desired client is not caught up, and there is another client that _is_ more caught up, then + // we schedule a movement, so we can move the active task to a more caught-up client. We'll try to // assign a warm-up to the desired client so that we can move it later on. - if (taskIsNotCaughtUpOnClientAndOtherCaughtUpClientsExist(task, client, tasksToCaughtUpClients)) { + if (taskIsNotCaughtUpOnClientAndOtherMoreCaughtUpClientsExist(task, client, clientStates, tasksToCaughtUpClients, tasksToClientByLag)) { taskMovements.add(new TaskMovement(task, client, tasksToCaughtUpClients.get(task))); } } @@ -102,17 +109,26 @@ static int assignActiveTaskMovements(final Map> tasksToC final int movementsNeeded = taskMovements.size(); for (final TaskMovement movement : taskMovements) { - final UUID standbySourceClient = caughtUpClientsByTaskLoad.poll( + // Attempt to find a caught up standby, otherwise find any caught up client, failing that use the most + // caught up client. + UUID sourceClient = caughtUpClientsByTaskLoad.poll( movement.task, c -> clientStates.get(c).hasStandbyTask(movement.task) ); - if (standbySourceClient == null) { - // there's not a caught-up standby available to take over the task, so we'll schedule a warmup instead - final UUID sourceClient = requireNonNull( - caughtUpClientsByTaskLoad.poll(movement.task), - "Tried to move task to caught-up client but none exist" + + if (sourceClient == null) { + sourceClient = caughtUpClientsByTaskLoad.poll(movement.task); + } + + if (sourceClient == null) { + sourceClient = requireNonNull( + mostCaughtUpEligibleClient(tasksToClientByLag, movement.task, movement.destination), + "Tried to move task to more caught-up client but none exist" ); + } + if (!clientStates.get(sourceClient).hasStandbyTask(movement.task)) { + // there's not a standby available to take over the task, so we'll schedule a warmup instead moveActiveAndTryToWarmUp( remainingWarmupReplicas, movement.task, @@ -125,10 +141,10 @@ static int assignActiveTaskMovements(final Map> tasksToC // we found a candidate to trade standby/active state with our destination, so we don't need a warmup swapStandbyAndActive( movement.task, - clientStates.get(standbySourceClient), + clientStates.get(sourceClient), clientStates.get(movement.destination) ); - caughtUpClientsByTaskLoad.offerAll(asList(standbySourceClient, movement.destination)); + caughtUpClientsByTaskLoad.offerAll(asList(sourceClient, movement.destination)); } } @@ -136,11 +152,12 @@ static int assignActiveTaskMovements(final Map> tasksToC } static int assignStandbyTaskMovements(final Map> tasksToCaughtUpClients, + final Map> tasksToClientByLag, final Map clientStates, final AtomicInteger remainingWarmupReplicas, final Map> warmups) { final BiFunction caughtUpPredicate = - (client, task) -> taskIsCaughtUpOnClientOrNoCaughtUpClientsExist(task, client, tasksToCaughtUpClients); + (client, task) -> taskIsCaughtUpOnClient(task, client, tasksToCaughtUpClients); final ConstrainedPrioritySet caughtUpClientsByTaskLoad = new ConstrainedPrioritySet( caughtUpPredicate, @@ -157,8 +174,8 @@ static int assignStandbyTaskMovements(final Map> tasksTo for (final TaskId task : state.standbyTasks()) { if (warmups.getOrDefault(destination, Collections.emptySet()).contains(task)) { // this is a warmup, so we won't move it. - } else if (taskIsNotCaughtUpOnClientAndOtherCaughtUpClientsExist(task, destination, tasksToCaughtUpClients)) { - // if the desired client is not caught up, and there is another client that _is_ caught up, then + } else if (taskIsNotCaughtUpOnClientAndOtherMoreCaughtUpClientsExist(task, destination, clientStates, tasksToCaughtUpClients, tasksToClientByLag)) { + // if the desired client is not caught up, and there is another client that _is_ more caught up, then // we schedule a movement, so we can move the active task to the caught-up client. We'll try to // assign a warm-up to the desired client so that we can move it later on. taskMovements.add(new TaskMovement(task, destination, tasksToCaughtUpClients.get(task))); @@ -170,11 +187,17 @@ static int assignStandbyTaskMovements(final Map> tasksTo int movementsNeeded = 0; for (final TaskMovement movement : taskMovements) { - final UUID sourceClient = caughtUpClientsByTaskLoad.poll( + final Function eligibleClientPredicate = + clientId -> !clientStates.get(clientId).hasAssignedTask(movement.task); + UUID sourceClient = caughtUpClientsByTaskLoad.poll( movement.task, - clientId -> !clientStates.get(clientId).hasAssignedTask(movement.task) + eligibleClientPredicate ); + if (sourceClient == null) { + sourceClient = mostCaughtUpEligibleClient(tasksToClientByLag, eligibleClientPredicate, movement.task, movement.destination); + } + if (sourceClient == null) { // then there's no caught-up client that doesn't already have a copy of this task, so there's // nowhere to move it. @@ -235,4 +258,24 @@ private static void swapStandbyAndActive(final TaskId task, destinationClientState.assignStandby(task); } + private static UUID mostCaughtUpEligibleClient(final Map> tasksToClientByLag, + final TaskId task, + final UUID destinationClient) { + return mostCaughtUpEligibleClient(tasksToClientByLag, client -> true, task, destinationClient); + } + + private static UUID mostCaughtUpEligibleClient(final Map> tasksToClientByLag, + final Function constraint, + final TaskId task, + final UUID destinationClient) { + for (final UUID client : tasksToClientByLag.get(task)) { + if (destinationClient.equals(client)) { + break; + } else if (constraint.apply(client)) { + return client; + } + } + return null; + } + } diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java index e878b13631658..f7a3d2e05ff30 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java @@ -420,21 +420,27 @@ public void shouldAssignToMostCaughtUpIfActiveTasksWasNotOnCaughtUpClient() { final Set allTasks = mkSet(TASK_0_0); final Set statefulTasks = mkSet(TASK_0_0); final ClientState client1 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, Long.MAX_VALUE), 1); - final ClientState client2 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 500L), 1); + final ClientState client2 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 1000L), 1); + final ClientState client3 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 500L), 1); final Map clientStates = mkMap( mkEntry(UUID_1, client1), - mkEntry(UUID_2, client2) + mkEntry(UUID_2, client2), + mkEntry(UUID_3, client3) ); final boolean probingRebalanceNeeded = new HighAvailabilityTaskAssignor().assign(clientStates, allTasks, statefulTasks, configWithStandbys); assertThat(clientStates.get(UUID_1).activeTasks(), is(emptySet())); - assertThat(clientStates.get(UUID_2).activeTasks(), is(singleton(TASK_0_0))); - // we'll warm up task 0_0 on client1 because it's first in sorted order, - // although this isn't an optimal convergence + assertThat(clientStates.get(UUID_2).activeTasks(), is(emptySet())); + assertThat(clientStates.get(UUID_3).activeTasks(), is(singleton(TASK_0_0))); + + assertThat(clientStates.get(UUID_1).standbyTasks(), is(singleton(TASK_0_0))); // warm up + assertThat(clientStates.get(UUID_2).standbyTasks(), is(singleton(TASK_0_0))); // standby + assertThat(clientStates.get(UUID_3).standbyTasks(), is(emptySet())); + assertThat(probingRebalanceNeeded, is(true)); - assertValidAssignment(0, 1, allTasks, emptySet(), clientStates, new StringBuilder()); + assertValidAssignment(1, 1, allTasks, emptySet(), clientStates, new StringBuilder()); assertBalancedActiveAssignment(clientStates, new StringBuilder()); assertBalancedStatefulAssignment(allTasks, clientStates, new StringBuilder()); assertBalancedTasks(clientStates); diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java index 9b58d18c88f60..e0fe05efcd863 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java @@ -19,8 +19,8 @@ import org.apache.kafka.streams.processor.TaskId; import org.junit.Test; -import java.util.Collection; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.SortedSet; @@ -29,9 +29,8 @@ import java.util.concurrent.atomic.AtomicInteger; import static java.util.Arrays.asList; -import static java.util.Collections.emptyList; +import static java.util.Collections.emptySet; import static java.util.Collections.emptySortedSet; -import static java.util.Collections.singletonList; 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.mkSet; @@ -58,17 +57,20 @@ public void shouldAssignTasksToClientsAndReturnFalseWhenAllClientsCaughtUp() { final Set allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2, TASK_1_0, TASK_1_1, TASK_1_2); final Map> tasksToCaughtUpClients = new HashMap<>(); + final Map> tasksToClientByLag = new HashMap<>(); for (final TaskId task : allTasks) { tasksToCaughtUpClients.put(task, mkSortedSet(UUID_1, UUID_2, UUID_3)); + tasksToClientByLag.put(task, asList(UUID_1, UUID_2, UUID_3)); } - final ClientState client1 = getClientStateWithActiveAssignment(asList(TASK_0_0, TASK_1_0)); - final ClientState client2 = getClientStateWithActiveAssignment(asList(TASK_0_1, TASK_1_1)); - final ClientState client3 = getClientStateWithActiveAssignment(asList(TASK_0_2, TASK_1_2)); + final ClientState client1 = getClientStateWithActiveAssignment(mkSet(TASK_0_0, TASK_1_0), allTasks, allTasks); + final ClientState client2 = getClientStateWithActiveAssignment(mkSet(TASK_0_1, TASK_1_1), allTasks, allTasks); + final ClientState client3 = getClientStateWithActiveAssignment(mkSet(TASK_0_2, TASK_1_2), allTasks, allTasks); assertThat( assignActiveTaskMovements( tasksToCaughtUpClients, + tasksToClientByLag, getClientStatesMap(client1, client2, client3), new TreeMap<>(), new AtomicInteger(maxWarmupReplicas) @@ -80,10 +82,11 @@ public void shouldAssignTasksToClientsAndReturnFalseWhenAllClientsCaughtUp() { @Test public void shouldAssignAllTasksToClientsAndReturnFalseIfNoClientsAreCaughtUp() { final int maxWarmupReplicas = Integer.MAX_VALUE; + final Set allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2, TASK_1_0, TASK_1_1, TASK_1_2); - final ClientState client1 = getClientStateWithActiveAssignment(asList(TASK_0_0, TASK_1_0)); - final ClientState client2 = getClientStateWithActiveAssignment(asList(TASK_0_1, TASK_1_1)); - final ClientState client3 = getClientStateWithActiveAssignment(asList(TASK_0_2, TASK_1_2)); + final ClientState client1 = getClientStateWithActiveAssignment(mkSet(TASK_0_0, TASK_1_0), mkSet(), allTasks); + final ClientState client2 = getClientStateWithActiveAssignment(mkSet(TASK_0_1, TASK_1_1), mkSet(), allTasks); + final ClientState client3 = getClientStateWithActiveAssignment(mkSet(TASK_0_2, TASK_1_2), mkSet(), allTasks); final Map> tasksToCaughtUpClients = mkMap( mkEntry(TASK_0_0, emptySortedSet()), @@ -93,9 +96,18 @@ public void shouldAssignAllTasksToClientsAndReturnFalseIfNoClientsAreCaughtUp() mkEntry(TASK_1_1, emptySortedSet()), mkEntry(TASK_1_2, emptySortedSet()) ); + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, asList(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_1, asList(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_2, asList(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_1_0, asList(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_1_1, asList(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_1_2, asList(UUID_1, UUID_2, UUID_3)) + ); assertThat( assignActiveTaskMovements( tasksToCaughtUpClients, + tasksToClientByLag, getClientStatesMap(client1, client2, client3), new TreeMap<>(), new AtomicInteger(maxWarmupReplicas) @@ -107,9 +119,10 @@ public void shouldAssignAllTasksToClientsAndReturnFalseIfNoClientsAreCaughtUp() @Test public void shouldMoveTasksToCaughtUpClientsAndAssignWarmupReplicasInTheirPlace() { final int maxWarmupReplicas = Integer.MAX_VALUE; - final ClientState client1 = getClientStateWithActiveAssignment(singletonList(TASK_0_0)); - final ClientState client2 = getClientStateWithActiveAssignment(singletonList(TASK_0_1)); - final ClientState client3 = getClientStateWithActiveAssignment(singletonList(TASK_0_2)); + final Set allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2); + final ClientState client1 = getClientStateWithActiveAssignment(mkSet(TASK_0_0), mkSet(TASK_0_0), allTasks); + final ClientState client2 = getClientStateWithActiveAssignment(mkSet(TASK_0_1), mkSet(TASK_0_2), allTasks); + final ClientState client3 = getClientStateWithActiveAssignment(mkSet(TASK_0_2), mkSet(TASK_0_1), allTasks); final Map clientStates = getClientStatesMap(client1, client2, client3); final Map> tasksToCaughtUpClients = mkMap( @@ -117,11 +130,17 @@ public void shouldMoveTasksToCaughtUpClientsAndAssignWarmupReplicasInTheirPlace( mkEntry(TASK_0_1, mkSortedSet(UUID_3)), mkEntry(TASK_0_2, mkSortedSet(UUID_2)) ); + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, asList(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_1, asList(UUID_3, UUID_1, UUID_2)), + mkEntry(TASK_0_2, asList(UUID_2, UUID_1, UUID_3)) + ); assertThat( "should have assigned movements", assignActiveTaskMovements( tasksToCaughtUpClients, + tasksToClientByLag, clientStates, new TreeMap<>(), new AtomicInteger(maxWarmupReplicas) @@ -142,9 +161,10 @@ public void shouldMoveTasksToCaughtUpClientsAndAssignWarmupReplicasInTheirPlace( @Test public void shouldOnlyGetUpToMaxWarmupReplicasAndReturnTrue() { final int maxWarmupReplicas = 1; - final ClientState client1 = getClientStateWithActiveAssignment(singletonList(TASK_0_0)); - final ClientState client2 = getClientStateWithActiveAssignment(singletonList(TASK_0_1)); - final ClientState client3 = getClientStateWithActiveAssignment(singletonList(TASK_0_2)); + final Set allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2); + final ClientState client1 = getClientStateWithActiveAssignment(mkSet(TASK_0_0), mkSet(TASK_0_0), allTasks); + final ClientState client2 = getClientStateWithActiveAssignment(mkSet(TASK_0_1), mkSet(TASK_0_2), allTasks); + final ClientState client3 = getClientStateWithActiveAssignment(mkSet(TASK_0_2), mkSet(TASK_0_1), allTasks); final Map clientStates = getClientStatesMap(client1, client2, client3); final Map> tasksToCaughtUpClients = mkMap( @@ -152,11 +172,17 @@ public void shouldOnlyGetUpToMaxWarmupReplicasAndReturnTrue() { mkEntry(TASK_0_1, mkSortedSet(UUID_3)), mkEntry(TASK_0_2, mkSortedSet(UUID_2)) ); + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, asList(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_1, asList(UUID_3, UUID_1, UUID_2)), + mkEntry(TASK_0_2, asList(UUID_2, UUID_1, UUID_3)) + ); assertThat( "should have assigned movements", assignActiveTaskMovements( tasksToCaughtUpClients, + tasksToClientByLag, clientStates, new TreeMap<>(), new AtomicInteger(maxWarmupReplicas) @@ -182,19 +208,24 @@ public void shouldOnlyGetUpToMaxWarmupReplicasAndReturnTrue() { @Test public void shouldNotCountPreviousStandbyTasksTowardsMaxWarmupReplicas() { final int maxWarmupReplicas = 0; - final ClientState client1 = getClientStateWithActiveAssignment(emptyList()); + final Set allTasks = mkSet(TASK_0_0); + final ClientState client1 = getClientStateWithActiveAssignment(mkSet(), mkSet(TASK_0_0), allTasks); client1.assignStandby(TASK_0_0); - final ClientState client2 = getClientStateWithActiveAssignment(singletonList(TASK_0_0)); + final ClientState client2 = getClientStateWithActiveAssignment(mkSet(TASK_0_0), mkSet(), allTasks); final Map clientStates = getClientStatesMap(client1, client2); final Map> tasksToCaughtUpClients = mkMap( mkEntry(TASK_0_0, mkSortedSet(UUID_1)) ); + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, asList(UUID_1, UUID_2)) + ); assertThat( "should have assigned movements", assignActiveTaskMovements( tasksToCaughtUpClients, + tasksToClientByLag, clientStates, new TreeMap<>(), new AtomicInteger(maxWarmupReplicas) @@ -215,8 +246,18 @@ public void shouldNotCountPreviousStandbyTasksTowardsMaxWarmupReplicas() { } - private static ClientState getClientStateWithActiveAssignment(final Collection activeTasks) { - final ClientState client1 = new ClientState(1); + private static ClientState getClientStateWithActiveAssignment(final Set activeTasks, + final Set caughtUpTasks, + final Set allTasks) { + final Map lags = new HashMap<>(); + for (final TaskId task : allTasks) { + if (caughtUpTasks.contains(task)) { + lags.put(task, 0L); + } else { + lags.put(task, 10000L); + } + } + final ClientState client1 = new ClientState(activeTasks, emptySet(), lags, 1); client1.assignActiveTasks(activeTasks); return client1; } From 5284d299cc4e0dad6402cc9e22ebb4f73cf27797 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Thu, 17 Feb 2022 09:18:09 +1300 Subject: [PATCH 3/9] [KAFKA-13600] Fix copy paste mistake --- .../streams/processor/internals/assignment/TaskMovement.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java index c6cc79e306748..472b38725bab8 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java @@ -59,7 +59,7 @@ private static boolean taskIsNotCaughtUpOnClientAndOtherMoreCaughtUpClientsExist final Map clientStates, final Map> tasksToCaughtUpClients, final Map> tasksToClientByLag) { - final List taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized map"); + final List taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized list"); if (taskIsCaughtUpOnClient(task, client, tasksToCaughtUpClients)) { return false; } From 7b1e8797de57637d75191465e274b4f253b77f97 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Fri, 18 Feb 2022 14:03:18 +1300 Subject: [PATCH 4/9] [KAFKA-13600] Make all cases for active assignment more explicit --- .../internals/assignment/TaskMovement.java | 69 +++++++++++-------- 1 file changed, 40 insertions(+), 29 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java index 472b38725bab8..2d829130cf695 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java @@ -111,40 +111,51 @@ static int assignActiveTaskMovements(final Map> tasksToC for (final TaskMovement movement : taskMovements) { // Attempt to find a caught up standby, otherwise find any caught up client, failing that use the most // caught up client. - UUID sourceClient = caughtUpClientsByTaskLoad.poll( + final UUID caughtUpStandbySourceClient = caughtUpClientsByTaskLoad.poll( movement.task, c -> clientStates.get(c).hasStandbyTask(movement.task) ); - - if (sourceClient == null) { - sourceClient = caughtUpClientsByTaskLoad.poll(movement.task); - } - - if (sourceClient == null) { - sourceClient = requireNonNull( - mostCaughtUpEligibleClient(tasksToClientByLag, movement.task, movement.destination), - "Tried to move task to more caught-up client but none exist" - ); - } - - if (!clientStates.get(sourceClient).hasStandbyTask(movement.task)) { - // there's not a standby available to take over the task, so we'll schedule a warmup instead - moveActiveAndTryToWarmUp( - remainingWarmupReplicas, - movement.task, - clientStates.get(sourceClient), - clientStates.get(movement.destination), - warmups.computeIfAbsent(movement.destination, x -> new TreeSet<>()) - ); - caughtUpClientsByTaskLoad.offerAll(asList(sourceClient, movement.destination)); - } else { - // we found a candidate to trade standby/active state with our destination, so we don't need a warmup + if (caughtUpStandbySourceClient != null) { swapStandbyAndActive( - movement.task, - clientStates.get(sourceClient), - clientStates.get(movement.destination) + movement.task, + clientStates.get(caughtUpStandbySourceClient), + clientStates.get(movement.destination) ); - caughtUpClientsByTaskLoad.offerAll(asList(sourceClient, movement.destination)); + caughtUpClientsByTaskLoad.offerAll(asList(caughtUpStandbySourceClient, movement.destination)); + } else { + final UUID caughtUpSourceClient = caughtUpClientsByTaskLoad.poll(movement.task); + if (caughtUpSourceClient != null) { + moveActiveAndTryToWarmUp( + remainingWarmupReplicas, + movement.task, + clientStates.get(caughtUpSourceClient), + clientStates.get(movement.destination), + warmups.computeIfAbsent(movement.destination, x -> new TreeSet<>()) + ); + caughtUpClientsByTaskLoad.offerAll(asList(caughtUpSourceClient, movement.destination)); + } else { + final UUID mostCaughtUpSourceClient = mostCaughtUpEligibleClient(tasksToClientByLag, movement.task, movement.destination); + if (mostCaughtUpSourceClient != null) { + if (clientStates.get(mostCaughtUpSourceClient).hasStandbyTask(movement.task)) { + swapStandbyAndActive( + movement.task, + clientStates.get(mostCaughtUpSourceClient), + clientStates.get(movement.destination) + ); + } else { + moveActiveAndTryToWarmUp( + remainingWarmupReplicas, + movement.task, + clientStates.get(mostCaughtUpSourceClient), + clientStates.get(movement.destination), + warmups.computeIfAbsent(movement.destination, x -> new TreeSet<>()) + ); + } + caughtUpClientsByTaskLoad.offerAll(asList(mostCaughtUpSourceClient, movement.destination)); + } else { + throw new IllegalStateException("Tried to move task to more caught-up client but none exist"); + } + } } } From 26149abd476ea427f2a7d49f1c5dc8dcdd40e371 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Mon, 28 Feb 2022 13:51:46 +1300 Subject: [PATCH 5/9] [KAFKA-13600] Use SortedSet for clients in tasksToClientByLag --- .../HighAvailabilityTaskAssignor.java | 16 +++--- .../internals/assignment/TaskMovement.java | 15 +++--- .../assignment/TaskMovementTest.java | 50 ++++++++++++------- 3 files changed, 45 insertions(+), 36 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java index 25a3e8dd4a489..d5ed53264db50 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java @@ -22,11 +22,9 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.ArrayList; import java.util.Comparator; import java.util.HashMap; import java.util.Iterator; -import java.util.List; import java.util.Map; import java.util.Set; import java.util.SortedMap; @@ -70,7 +68,7 @@ public boolean assign(final Map clients, configs.acceptableRecoveryLag ); - final Map> tasksToClientByLag = tasksToClientByLag(statefulTasks, clientStates); + final Map> tasksToClientByLag = tasksToClientByLag(statefulTasks, clientStates); // We temporarily need to know which standby tasks were intended as warmups // for active tasks, so that we don't move them (again) when we plan standby @@ -259,14 +257,14 @@ private static Map> tasksToCaughtUpClients(final Set> tasksToClientByLag(final Set statefulTasks, + private static Map> tasksToClientByLag(final Set statefulTasks, final Map clientStates) { - final Map> tasksToClientByLag = new HashMap<>(); + final Map> tasksToClientByLag = new HashMap<>(); for (final TaskId task : statefulTasks) { - final List> clientLag = new ArrayList<>(clientStates.entrySet()); - clientLag.sort(Comparator.>comparingLong( - a -> a.getValue().lagFor(task)).thenComparing(Map.Entry::getKey)); - tasksToClientByLag.put(task, clientLag.stream().map(Map.Entry::getKey).collect(Collectors.toList())); + final SortedSet clientLag = new TreeSet<>(Comparator.comparingLong(a -> + clientStates.get(a).lagFor(task)).thenComparing(a -> a)); + clientLag.addAll(clientStates.keySet()); + tasksToClientByLag.put(task, clientLag); } return tasksToClientByLag; } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java index 2d829130cf695..b260fe9ec41aa 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java @@ -20,7 +20,6 @@ import java.util.Collections; import java.util.Comparator; -import java.util.List; import java.util.Map; import java.util.PriorityQueue; import java.util.Queue; @@ -58,12 +57,12 @@ private static boolean taskIsNotCaughtUpOnClientAndOtherMoreCaughtUpClientsExist final UUID client, final Map clientStates, final Map> tasksToCaughtUpClients, - final Map> tasksToClientByLag) { - final List taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized list"); + final Map> tasksToClientByLag) { + final SortedSet taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized set"); if (taskIsCaughtUpOnClient(task, client, tasksToCaughtUpClients)) { return false; } - final long mostCaughtUpLag = clientStates.get(taskClients.get(0)).lagFor(task); + final long mostCaughtUpLag = clientStates.get(taskClients.first()).lagFor(task); final long clientLag = clientStates.get(client).lagFor(task); return mostCaughtUpLag < clientLag; } @@ -76,7 +75,7 @@ private static boolean taskIsCaughtUpOnClient(final TaskId task, } static int assignActiveTaskMovements(final Map> tasksToCaughtUpClients, - final Map> tasksToClientByLag, + final Map> tasksToClientByLag, final Map clientStates, final Map> warmups, final AtomicInteger remainingWarmupReplicas) { @@ -163,7 +162,7 @@ static int assignActiveTaskMovements(final Map> tasksToC } static int assignStandbyTaskMovements(final Map> tasksToCaughtUpClients, - final Map> tasksToClientByLag, + final Map> tasksToClientByLag, final Map clientStates, final AtomicInteger remainingWarmupReplicas, final Map> warmups) { @@ -269,13 +268,13 @@ private static void swapStandbyAndActive(final TaskId task, destinationClientState.assignStandby(task); } - private static UUID mostCaughtUpEligibleClient(final Map> tasksToClientByLag, + private static UUID mostCaughtUpEligibleClient(final Map> tasksToClientByLag, final TaskId task, final UUID destinationClient) { return mostCaughtUpEligibleClient(tasksToClientByLag, client -> true, task, destinationClient); } - private static UUID mostCaughtUpEligibleClient(final Map> tasksToClientByLag, + private static UUID mostCaughtUpEligibleClient(final Map> tasksToClientByLag, final Function constraint, final TaskId task, final UUID destinationClient) { diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java index e0fe05efcd863..b734688e6d319 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java @@ -19,12 +19,14 @@ import org.apache.kafka.streams.processor.TaskId; import org.junit.Test; +import java.util.Comparator; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; import java.util.SortedSet; import java.util.TreeMap; +import java.util.TreeSet; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; @@ -57,10 +59,10 @@ public void shouldAssignTasksToClientsAndReturnFalseWhenAllClientsCaughtUp() { final Set allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2, TASK_1_0, TASK_1_1, TASK_1_2); final Map> tasksToCaughtUpClients = new HashMap<>(); - final Map> tasksToClientByLag = new HashMap<>(); + final Map> tasksToClientByLag = new HashMap<>(); for (final TaskId task : allTasks) { tasksToCaughtUpClients.put(task, mkSortedSet(UUID_1, UUID_2, UUID_3)); - tasksToClientByLag.put(task, asList(UUID_1, UUID_2, UUID_3)); + tasksToClientByLag.put(task, mkOrderedSet(UUID_1, UUID_2, UUID_3)); } final ClientState client1 = getClientStateWithActiveAssignment(mkSet(TASK_0_0, TASK_1_0), allTasks, allTasks); @@ -96,13 +98,13 @@ public void shouldAssignAllTasksToClientsAndReturnFalseIfNoClientsAreCaughtUp() mkEntry(TASK_1_1, emptySortedSet()), mkEntry(TASK_1_2, emptySortedSet()) ); - final Map> tasksToClientByLag = mkMap( - mkEntry(TASK_0_0, asList(UUID_1, UUID_2, UUID_3)), - mkEntry(TASK_0_1, asList(UUID_1, UUID_2, UUID_3)), - mkEntry(TASK_0_2, asList(UUID_1, UUID_2, UUID_3)), - mkEntry(TASK_1_0, asList(UUID_1, UUID_2, UUID_3)), - mkEntry(TASK_1_1, asList(UUID_1, UUID_2, UUID_3)), - mkEntry(TASK_1_2, asList(UUID_1, UUID_2, UUID_3)) + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_1, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_2, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_1_0, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_1_1, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_1_2, mkOrderedSet(UUID_1, UUID_2, UUID_3)) ); assertThat( assignActiveTaskMovements( @@ -130,10 +132,10 @@ public void shouldMoveTasksToCaughtUpClientsAndAssignWarmupReplicasInTheirPlace( mkEntry(TASK_0_1, mkSortedSet(UUID_3)), mkEntry(TASK_0_2, mkSortedSet(UUID_2)) ); - final Map> tasksToClientByLag = mkMap( - mkEntry(TASK_0_0, asList(UUID_1, UUID_2, UUID_3)), - mkEntry(TASK_0_1, asList(UUID_3, UUID_1, UUID_2)), - mkEntry(TASK_0_2, asList(UUID_2, UUID_1, UUID_3)) + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_1, mkOrderedSet(UUID_3, UUID_1, UUID_2)), + mkEntry(TASK_0_2, mkOrderedSet(UUID_2, UUID_1, UUID_3)) ); assertThat( @@ -172,10 +174,10 @@ public void shouldOnlyGetUpToMaxWarmupReplicasAndReturnTrue() { mkEntry(TASK_0_1, mkSortedSet(UUID_3)), mkEntry(TASK_0_2, mkSortedSet(UUID_2)) ); - final Map> tasksToClientByLag = mkMap( - mkEntry(TASK_0_0, asList(UUID_1, UUID_2, UUID_3)), - mkEntry(TASK_0_1, asList(UUID_3, UUID_1, UUID_2)), - mkEntry(TASK_0_2, asList(UUID_2, UUID_1, UUID_3)) + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_1, mkOrderedSet(UUID_3, UUID_1, UUID_2)), + mkEntry(TASK_0_2, mkOrderedSet(UUID_2, UUID_1, UUID_3)) ); assertThat( @@ -217,8 +219,8 @@ public void shouldNotCountPreviousStandbyTasksTowardsMaxWarmupReplicas() { final Map> tasksToCaughtUpClients = mkMap( mkEntry(TASK_0_0, mkSortedSet(UUID_1)) ); - final Map> tasksToClientByLag = mkMap( - mkEntry(TASK_0_0, asList(UUID_1, UUID_2)) + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, mkOrderedSet(UUID_1, UUID_2)) ); assertThat( @@ -262,4 +264,14 @@ private static ClientState getClientStateWithActiveAssignment(final Set return client1; } + /** + * Creates a SortedSet with the sort order being the order of elements in the parameter list + */ + private static SortedSet mkOrderedSet(final UUID... clients) { + final List clientList = asList(clients); + final SortedSet set = new TreeSet<>(Comparator.comparing(clientList::indexOf)); + set.addAll(clientList); + return set; + } + } From ec7e24c3c83df95ee3d87e7cdef114b1e1bcbfc9 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Sat, 12 Mar 2022 09:00:55 +1300 Subject: [PATCH 6/9] [KAFKA-13600] Fix merge issues --- .../assignment/HighAvailabilityTaskAssignorTest.java | 6 +++--- .../processor/internals/assignment/TaskMovementTest.java | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java index 6eda09dda5150..a96dad43cb16f 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java @@ -427,9 +427,9 @@ public void shouldComputeNewAssignmentIfActiveTasksWasNotOnCaughtUpClient() { public void shouldAssignToMostCaughtUpIfActiveTasksWasNotOnCaughtUpClient() { final Set allTasks = mkSet(TASK_0_0); final Set statefulTasks = mkSet(TASK_0_0); - final ClientState client1 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, Long.MAX_VALUE), 1); - final ClientState client2 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 1000L), 1); - final ClientState client3 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 500L), 1); + final ClientState client1 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, Long.MAX_VALUE), EMPTY_CLIENT_TAGS, 1); + final ClientState client2 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 1000L), EMPTY_CLIENT_TAGS, 1); + final ClientState client3 = new ClientState(emptySet(), emptySet(), singletonMap(TASK_0_0, 500L), EMPTY_CLIENT_TAGS, 1); final Map clientStates = mkMap( mkEntry(UUID_1, client1), mkEntry(UUID_2, client2), diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java index b734688e6d319..f9d452b137fbe 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java @@ -31,6 +31,7 @@ import java.util.concurrent.atomic.AtomicInteger; import static java.util.Arrays.asList; +import static java.util.Collections.emptyMap; import static java.util.Collections.emptySet; import static java.util.Collections.emptySortedSet; import static org.apache.kafka.common.utils.Utils.mkEntry; @@ -259,7 +260,7 @@ private static ClientState getClientStateWithActiveAssignment(final Set lags.put(task, 10000L); } } - final ClientState client1 = new ClientState(activeTasks, emptySet(), lags, 1); + final ClientState client1 = new ClientState(activeTasks, emptySet(), lags, emptyMap(), 1); client1.assignActiveTasks(activeTasks); return client1; } From 2dfb79d66c207f99a4c6cd73b8c57a134a4e9132 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Mon, 14 Mar 2022 16:13:45 +1300 Subject: [PATCH 7/9] [KAFKA-13600] break up large method --- .../internals/assignment/TaskMovement.java | 119 +++++++++++------- 1 file changed, 74 insertions(+), 45 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java index b260fe9ec41aa..0a050ed03ff9b 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java @@ -110,51 +110,12 @@ static int assignActiveTaskMovements(final Map> tasksToC for (final TaskMovement movement : taskMovements) { // Attempt to find a caught up standby, otherwise find any caught up client, failing that use the most // caught up client. - final UUID caughtUpStandbySourceClient = caughtUpClientsByTaskLoad.poll( - movement.task, - c -> clientStates.get(c).hasStandbyTask(movement.task) - ); - if (caughtUpStandbySourceClient != null) { - swapStandbyAndActive( - movement.task, - clientStates.get(caughtUpStandbySourceClient), - clientStates.get(movement.destination) - ); - caughtUpClientsByTaskLoad.offerAll(asList(caughtUpStandbySourceClient, movement.destination)); - } else { - final UUID caughtUpSourceClient = caughtUpClientsByTaskLoad.poll(movement.task); - if (caughtUpSourceClient != null) { - moveActiveAndTryToWarmUp( - remainingWarmupReplicas, - movement.task, - clientStates.get(caughtUpSourceClient), - clientStates.get(movement.destination), - warmups.computeIfAbsent(movement.destination, x -> new TreeSet<>()) - ); - caughtUpClientsByTaskLoad.offerAll(asList(caughtUpSourceClient, movement.destination)); - } else { - final UUID mostCaughtUpSourceClient = mostCaughtUpEligibleClient(tasksToClientByLag, movement.task, movement.destination); - if (mostCaughtUpSourceClient != null) { - if (clientStates.get(mostCaughtUpSourceClient).hasStandbyTask(movement.task)) { - swapStandbyAndActive( - movement.task, - clientStates.get(mostCaughtUpSourceClient), - clientStates.get(movement.destination) - ); - } else { - moveActiveAndTryToWarmUp( - remainingWarmupReplicas, - movement.task, - clientStates.get(mostCaughtUpSourceClient), - clientStates.get(movement.destination), - warmups.computeIfAbsent(movement.destination, x -> new TreeSet<>()) - ); - } - caughtUpClientsByTaskLoad.offerAll(asList(mostCaughtUpSourceClient, movement.destination)); - } else { - throw new IllegalStateException("Tried to move task to more caught-up client but none exist"); - } - } + final boolean moved = tryToSwapStandbyAndActiveOnCaughtUpClient(clientStates, caughtUpClientsByTaskLoad, movement) || + tryToMoveActiveToCaughtUpClientAndTryToWarmUp(clientStates, warmups, remainingWarmupReplicas, caughtUpClientsByTaskLoad, movement) || + tryToMoveActiveToMostCaughtUpClient(tasksToClientByLag, clientStates, warmups, remainingWarmupReplicas, caughtUpClientsByTaskLoad, movement); + + if (!moved) { + throw new IllegalStateException("Tried to move task to more caught-up client but none exist"); } } @@ -226,6 +187,74 @@ static int assignStandbyTaskMovements(final Map> tasksTo return movementsNeeded; } + private static boolean tryToSwapStandbyAndActiveOnCaughtUpClient(final Map clientStates, + final ConstrainedPrioritySet caughtUpClientsByTaskLoad, + final TaskMovement movement) { + final UUID caughtUpStandbySourceClient = caughtUpClientsByTaskLoad.poll( + movement.task, + c -> clientStates.get(c).hasStandbyTask(movement.task) + ); + if (caughtUpStandbySourceClient != null) { + swapStandbyAndActive( + movement.task, + clientStates.get(caughtUpStandbySourceClient), + clientStates.get(movement.destination) + ); + caughtUpClientsByTaskLoad.offerAll(asList(caughtUpStandbySourceClient, movement.destination)); + return true; + } + return false; + } + + private static boolean tryToMoveActiveToCaughtUpClientAndTryToWarmUp(final Map clientStates, + final Map> warmups, + final AtomicInteger remainingWarmupReplicas, + final ConstrainedPrioritySet caughtUpClientsByTaskLoad, + final TaskMovement movement) { + final UUID caughtUpSourceClient = caughtUpClientsByTaskLoad.poll(movement.task); + if (caughtUpSourceClient != null) { + moveActiveAndTryToWarmUp( + remainingWarmupReplicas, + movement.task, + clientStates.get(caughtUpSourceClient), + clientStates.get(movement.destination), + warmups.computeIfAbsent(movement.destination, x -> new TreeSet<>()) + ); + caughtUpClientsByTaskLoad.offerAll(asList(caughtUpSourceClient, movement.destination)); + return true; + } + return false; + } + + private static boolean tryToMoveActiveToMostCaughtUpClient(final Map> tasksToClientByLag, + final Map clientStates, + final Map> warmups, + final AtomicInteger remainingWarmupReplicas, + final ConstrainedPrioritySet caughtUpClientsByTaskLoad, + final TaskMovement movement) { + final UUID mostCaughtUpSourceClient = mostCaughtUpEligibleClient(tasksToClientByLag, movement.task, movement.destination); + if (mostCaughtUpSourceClient != null) { + if (clientStates.get(mostCaughtUpSourceClient).hasStandbyTask(movement.task)) { + swapStandbyAndActive( + movement.task, + clientStates.get(mostCaughtUpSourceClient), + clientStates.get(movement.destination) + ); + } else { + moveActiveAndTryToWarmUp( + remainingWarmupReplicas, + movement.task, + clientStates.get(mostCaughtUpSourceClient), + clientStates.get(movement.destination), + warmups.computeIfAbsent(movement.destination, x -> new TreeSet<>()) + ); + } + caughtUpClientsByTaskLoad.offerAll(asList(mostCaughtUpSourceClient, movement.destination)); + return true; + } + return false; + } + private static void moveActiveAndTryToWarmUp(final AtomicInteger remainingWarmupReplicas, final TaskId task, final ClientState sourceClientState, From 067897e3e5c22b65eae999312f4de48313dc2048 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Fri, 25 Mar 2022 14:37:13 +1300 Subject: [PATCH 8/9] Update streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java Co-authored-by: Bruno Cadonna --- .../streams/processor/internals/assignment/TaskMovement.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java index 0a050ed03ff9b..38e64276ba0a8 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java @@ -115,7 +115,7 @@ static int assignActiveTaskMovements(final Map> tasksToC tryToMoveActiveToMostCaughtUpClient(tasksToClientByLag, clientStates, warmups, remainingWarmupReplicas, caughtUpClientsByTaskLoad, movement); if (!moved) { - throw new IllegalStateException("Tried to move task to more caught-up client but none exist"); + throw new IllegalStateException("Tried to move task to more caught-up client as scheduled before but none exist"); } } From d3798f53a3f0cd21b706caedb7271443b3e254e2 Mon Sep 17 00:00:00 2001 From: Tim Patterson Date: Fri, 25 Mar 2022 15:58:22 +1300 Subject: [PATCH 9/9] [KAFKA-13600] Add taskmovement test for moving active to most caught up --- .../assignment/TaskMovementTest.java | 54 ++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java index f9d452b137fbe..baf6d18496c1e 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovementTest.java @@ -161,6 +161,53 @@ public void shouldMoveTasksToCaughtUpClientsAndAssignWarmupReplicasInTheirPlace( assertThat(client3, hasProperty("standbyTasks", ClientState::standbyTasks, mkSet(TASK_0_2))); } + @Test + public void shouldMoveTasksToMostCaughtUpClientsAndAssignWarmupReplicasInTheirPlace() { + final int maxWarmupReplicas = Integer.MAX_VALUE; + final Map client1Lags = mkMap(mkEntry(TASK_0_0, 10000L), mkEntry(TASK_0_1, 20000L), mkEntry(TASK_0_2, 30000L)); + final Map client2Lags = mkMap(mkEntry(TASK_0_2, 10000L), mkEntry(TASK_0_0, 20000L), mkEntry(TASK_0_1, 30000L)); + final Map client3Lags = mkMap(mkEntry(TASK_0_1, 10000L), mkEntry(TASK_0_2, 20000L), mkEntry(TASK_0_0, 30000L)); + + final ClientState client1 = getClientStateWithLags(mkSet(TASK_0_0), client1Lags); + final ClientState client2 = getClientStateWithLags(mkSet(TASK_0_1), client2Lags); + final ClientState client3 = getClientStateWithLags(mkSet(TASK_0_2), client3Lags); + // To test when the task is already a standby on the most caught up node + client3.assignStandby(TASK_0_1); + final Map clientStates = getClientStatesMap(client1, client2, client3); + + final Map> tasksToCaughtUpClients = mkMap( + mkEntry(TASK_0_0, mkSortedSet()), + mkEntry(TASK_0_1, mkSortedSet()), + mkEntry(TASK_0_2, mkSortedSet()) + ); + final Map> tasksToClientByLag = mkMap( + mkEntry(TASK_0_0, mkOrderedSet(UUID_1, UUID_2, UUID_3)), + mkEntry(TASK_0_1, mkOrderedSet(UUID_3, UUID_1, UUID_2)), + mkEntry(TASK_0_2, mkOrderedSet(UUID_2, UUID_3, UUID_1)) + ); + + assertThat( + "should have assigned movements", + assignActiveTaskMovements( + tasksToCaughtUpClients, + tasksToClientByLag, + clientStates, + new TreeMap<>(), + new AtomicInteger(maxWarmupReplicas) + ), + is(2) + ); + // The active tasks have changed to the ones that each client is most caught up on + assertThat(client1, hasProperty("activeTasks", ClientState::activeTasks, mkSet(TASK_0_0))); + assertThat(client2, hasProperty("activeTasks", ClientState::activeTasks, mkSet(TASK_0_2))); + assertThat(client3, hasProperty("activeTasks", ClientState::activeTasks, mkSet(TASK_0_1))); + + // we assigned warmups to migrate to the input active assignment + assertThat(client1, hasProperty("standbyTasks", ClientState::standbyTasks, mkSet())); + assertThat(client2, hasProperty("standbyTasks", ClientState::standbyTasks, mkSet(TASK_0_1))); + assertThat(client3, hasProperty("standbyTasks", ClientState::standbyTasks, mkSet(TASK_0_2))); + } + @Test public void shouldOnlyGetUpToMaxWarmupReplicasAndReturnTrue() { final int maxWarmupReplicas = 1; @@ -260,7 +307,12 @@ private static ClientState getClientStateWithActiveAssignment(final Set lags.put(task, 10000L); } } - final ClientState client1 = new ClientState(activeTasks, emptySet(), lags, emptyMap(), 1); + return getClientStateWithLags(activeTasks, lags); + } + + private static ClientState getClientStateWithLags(final Set activeTasks, + final Map taskLags) { + final ClientState client1 = new ClientState(activeTasks, emptySet(), taskLags, emptyMap(), 1); client1.assignActiveTasks(activeTasks); return client1; }