Skip to content

KAFKA-13600: Kafka Streams - Fall back to most caught up client if no caught up clients exist #11760

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cadonna merged 10 commits into from
Mar 28, 2022
Merged

KAFKA-13600: Kafka Streams - Fall back to most caught up client if no caught up clients exist #11760

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
382ae5a
[KAFKA-13600] Unit test to capture not assigning to caught up task if…
tim-patterson Feb 14, 2022
6cec06c
[KAFKA-13600] Fall back to most caught up client if a caught up clien…
tim-patterson Feb 15, 2022
5284d29
[KAFKA-13600] Fix copy paste mistake
tim-patterson Feb 16, 2022
7b1e879
[KAFKA-13600] Make all cases for active assignment more explicit
tim-patterson Feb 18, 2022
26149ab
[KAFKA-13600] Use SortedSet for clients in tasksToClientByLag
tim-patterson Feb 28, 2022
67d5ba0
[KAFKA-13600] Merge remote-tracking branch 'origin/trunk' into KAFKA-…
tim-patterson Mar 11, 2022
ec7e24c
[KAFKA-13600] Fix merge issues
tim-patterson Mar 11, 2022
2dfb79d
[KAFKA-13600] break up large method
tim-patterson Mar 14, 2022
067897e
Update streams/src/main/java/org/apache/kafka/streams/processor/inter…
tim-patterson Mar 25, 2022
d3798f5
[KAFKA-13600] Add taskmovement test for moving active to most caught up
tim-patterson Mar 25, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions ...org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -67,6 +70,8 @@ public boolean assign(final Map<UUID, ClientState> clients,
configs.acceptableRecoveryLag
);

final Map<TaskId, List<UUID>> 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
Expand All @@ -76,13 +81,15 @@ public boolean assign(final Map<UUID, ClientState> clients,

final int neededActiveTaskMovements = assignActiveTaskMovements(
tasksToCaughtUpClients,
tasksToClientByLag,
clientStates,
warmups,
remainingWarmupReplicas
);

final int neededStandbyTaskMovements = assignStandbyTaskMovements(
tasksToCaughtUpClients,
tasksToClientByLag,
clientStates,
remainingWarmupReplicas,
warmups
Expand Down Expand Up @@ -252,6 +259,18 @@ private static Map<TaskId, SortedSet<UUID>> tasksToCaughtUpClients(final Set<Tas
return taskToCaughtUpClients;
}

private static Map<TaskId, List<UUID>> tasksToClientByLag(final Set<TaskId> statefulTasks,
final Map<UUID, ClientState> clientStates) {
final Map<TaskId, List<UUID>> tasksToClientByLag = new HashMap<>();
for (final TaskId task : statefulTasks) {
final List<Map.Entry<UUID, ClientState>> clientLag = new ArrayList<>(clientStates.entrySet());
clientLag.sort(Comparator.<Map.Entry<UUID, ClientState>>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;
}
Expand Down
101 changes: 72 additions & 29 deletions ...s/src/main/java/org/apache/kafka/streams/processor/internals/assignment/TaskMovement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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() {
Expand All @@ -56,25 +54,34 @@ private int numCaughtUpClients() {
return caughtUpClients.size();
}

private static boolean taskIsNotCaughtUpOnClientAndOtherCaughtUpClientsExist(final TaskId task,
final UUID client,
final Map<TaskId, SortedSet<UUID>> tasksToCaughtUpClients) {
return !taskIsCaughtUpOnClientOrNoCaughtUpClientsExist(task, client, tasksToCaughtUpClients);
private static boolean taskIsNotCaughtUpOnClientAndOtherMoreCaughtUpClientsExist(final TaskId task,
final UUID client,
final Map<UUID, ClientState> clientStates,
final Map<TaskId, SortedSet<UUID>> tasksToCaughtUpClients,
final Map<TaskId, List<UUID>> tasksToClientByLag) {

@cadonna cadonna Feb 21, 2022 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be possible to use a SortedSet instead of a List here? Would make it clearer that client IDs should only appear once in this list.

@tim-patterson tim-patterson Feb 28, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure I've pushed up a commit that does this.
My only concern is that these SortedSet's now end up holding a reference to clientStates rather than just being a plain old list/set etc.

final List<UUID> taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized map");

@cadonna cadonna Feb 16, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
final List<UUID> taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized map");
final List<UUID> taskClients = requireNonNull(tasksToClientByLag.get(task), "uninitialized list");

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<TaskId, SortedSet<UUID>> tasksToCaughtUpClients) {
private static boolean taskIsCaughtUpOnClient(final TaskId task,
final UUID client,
final Map<TaskId, SortedSet<UUID>> tasksToCaughtUpClients) {
final Set<UUID> caughtUpClients = requireNonNull(tasksToCaughtUpClients.get(task), "uninitialized set");
return caughtUpClients.isEmpty() || caughtUpClients.contains(client);
return caughtUpClients.contains(client);
}

static int assignActiveTaskMovements(final Map<TaskId, SortedSet<UUID>> tasksToCaughtUpClients,
final Map<TaskId, List<UUID>> tasksToClientByLag,
final Map<UUID, ClientState> clientStates,
final Map<UUID, Set<TaskId>> warmups,
final AtomicInteger remainingWarmupReplicas) {
final BiFunction<UUID, TaskId, Boolean> caughtUpPredicate =
(client, task) -> taskIsCaughtUpOnClientOrNoCaughtUpClientsExist(task, client, tasksToCaughtUpClients);
(client, task) -> taskIsCaughtUpOnClient(task, client, tasksToCaughtUpClients);

final ConstrainedPrioritySet caughtUpClientsByTaskLoad = new ConstrainedPrioritySet(
caughtUpPredicate,
Expand All @@ -89,10 +96,10 @@ static int assignActiveTaskMovements(final Map<TaskId, SortedSet<UUID>> 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)));
}
}
Expand All @@ -102,17 +109,26 @@ static int assignActiveTaskMovements(final Map<TaskId, SortedSet<UUID>> 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

@cadonna cadonna Feb 16, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This part of the code is a bit hard to read. After it is clear that there is a client with a standby after line 117 (if sourceClient != null), you could execute the content of the else branch on line 140. If there is no such client, then you can check for the caught up and most caught up client. Something like:

if (sourceClient != null) {
    swapStandbyAndActive(...)
}  else {
    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"
        );
    }
    moveActiveAndTryToWarmUp(...)
}
caughtUpClientsByTaskLoad.offerAll(asList(sourceClient, movement.destination));

@tim-patterson tim-patterson Feb 16, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I sort of originally did something like that , but a unit test caused a runtime assertion.

Basically there's ~5 cases.

  1. Caught up client with standby
  2. Caught up client without standby
  3. Not Caught up client with standby
  4. Not Caught up client without standby
  5. No client found (runtime assertion).

So the sourceClient from mostCaughtUpEligibleClient(...) may actually need to call swapStandbyAndActive(...) depending on if it's actually got a standby assigned to it.

I guess can do something like

if (sourceClient != null) {
    swapStandbyAndActive(...)
}  else {
    sourceClient = caughtUpClientsByTaskLoad.poll(movement.task);
    if (sourceClient != null) {
        moveActiveAndTryToWarmUp(...)
    } else {
        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)) {
            swapStandbyAndActive(...)
        } else {
            moveActiveAndTryToWarmUp(...)
        }
    }
}
caughtUpClientsByTaskLoad.offerAll(asList(sourceClient, movement.destination));

@cadonna cadonna Feb 17, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, I totally missed case 3.

What about the following:

if (sourceClient != null) {
    swapStandbyAndActive(...);
}  else {
    sourceClient = caughtUpClientsByTaskLoad.poll(movement.task);
    if (sourceClient != null) {
        moveActiveAndTryToWarmUp(...);
    } else {
        sourceClient = mostCaughtUpEligibleClient(tasksToClientByLag, movement.task, movement.destination);
        if (sourceClient != null) {
            if (clientStates.get(sourceClient).hasStandbyTask(movement.task)) {
                swapStandbyAndActive(...);
            } else {
                moveActiveAndTryToWarmUp(...);
            }
        } else {
            throw new IllegalStateException("Tried to move task to more caught-up client but none exist");
        }
    }
}
caughtUpClientsByTaskLoad.offerAll(asList(sourceClient, movement.destination));

I think it makes the cases more explicit in the code.

@cadonna cadonna Feb 17, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could even make it more explicit by introducing methods like:

private boolean tryToSwapStandbyAndActiveOnCaughtUpClient(...);
private boolean tryToMoveToCaughtUpClientAndTryToWarmUp(...);
private boolean tryToSwapStandbyAndActiveOnMostCaughtUpClient(...);
private boolean tryToMoveToMostCaughtUpClientAndTryToWarmUp(...);

But that is optional and up to you.

@tim-patterson tim-patterson Feb 21, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've done it with the nested if/else's but stopped reusing the same sourceClient var to try make it a little clearer.
Let me know what you think

@cadonna cadonna Feb 21, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is fine as you did. However, the method is already over 100 lines long. Maybe the code would benefit from a bit of refactoring. The method was already quite long before you added your part. WDYT?

@tim-patterson tim-patterson Feb 21, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure let me have a bit of a play around.
All these methods being static means that there's a lot of state/method arguments to pass to each method call.
So I'm unsure about how much of a win we're going to get extracting smaller bits of code out into separate methods.
Maybe some closures/local functions might help....
I'll have a bit of a play and get back to you :)

@cadonna cadonna Feb 21, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That sounds great! Thanks!

moveActiveAndTryToWarmUp(
remainingWarmupReplicas,
movement.task,
Expand All @@ -125,22 +141,23 @@ static int assignActiveTaskMovements(final Map<TaskId, SortedSet<UUID>> 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));
}
}

return movementsNeeded;
}

static int assignStandbyTaskMovements(final Map<TaskId, SortedSet<UUID>> tasksToCaughtUpClients,
final Map<TaskId, List<UUID>> tasksToClientByLag,
final Map<UUID, ClientState> clientStates,
final AtomicInteger remainingWarmupReplicas,
final Map<UUID, Set<TaskId>> warmups) {
final BiFunction<UUID, TaskId, Boolean> caughtUpPredicate =
(client, task) -> taskIsCaughtUpOnClientOrNoCaughtUpClientsExist(task, client, tasksToCaughtUpClients);
(client, task) -> taskIsCaughtUpOnClient(task, client, tasksToCaughtUpClients);

final ConstrainedPrioritySet caughtUpClientsByTaskLoad = new ConstrainedPrioritySet(
caughtUpPredicate,
Expand All @@ -157,8 +174,8 @@ static int assignStandbyTaskMovements(final Map<TaskId, SortedSet<UUID>> 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)));
Expand All @@ -170,11 +187,17 @@ static int assignStandbyTaskMovements(final Map<TaskId, SortedSet<UUID>> tasksTo
int movementsNeeded = 0;

for (final TaskMovement movement : taskMovements) {
final UUID sourceClient = caughtUpClientsByTaskLoad.poll(
final Function<UUID, Boolean> 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.
Expand Down Expand Up @@ -235,4 +258,24 @@ private static void swapStandbyAndActive(final TaskId task,
destinationClientState.assignStandby(task);
}

private static UUID mostCaughtUpEligibleClient(final Map<TaskId, List<UUID>> tasksToClientByLag,
final TaskId task,
final UUID destinationClient) {
return mostCaughtUpEligibleClient(tasksToClientByLag, client -> true, task, destinationClient);
}

private static UUID mostCaughtUpEligibleClient(final Map<TaskId, List<UUID>> tasksToClientByLag,
final Function<UUID, Boolean> 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;
}

}
31 changes: 31 additions & 0 deletions ...apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,37 @@ public void shouldComputeNewAssignmentIfActiveTasksWasNotOnCaughtUpClient() {
assertBalancedTasks(clientStates);
}

@Test
public void shouldAssignToMostCaughtUpIfActiveTasksWasNotOnCaughtUpClient() {
final Set<TaskId> allTasks = mkSet(TASK_0_0);
final Set<TaskId> 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 Map<UUID, ClientState> clientStates = mkMap(
mkEntry(UUID_1, client1),
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(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(1, 1, allTasks, emptySet(), clientStates, new StringBuilder());
assertBalancedActiveAssignment(clientStates, new StringBuilder());
assertBalancedStatefulAssignment(allTasks, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
}

@Test
public void shouldAssignStandbysForStatefulTasks() {
final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1);
Expand Down
Loading