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

        final ConstrainedPrioritySet caughtUpClientsByTaskLoad = new ConstrainedPrioritySet(
            caughtUpPredicate,
            client -> clientStates.get(client).assignedTaskLoad()
        );

        final Queue<TaskMovement> taskMovements = new PriorityQueue<>(
            Comparator.comparing(TaskMovement::numCaughtUpClients).thenComparing(TaskMovement::task)
        );

        for (final Map.Entry<ProcessId, ClientState> clientStateEntry : clientStates.entrySet()) {
            final ProcessId 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_ 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 (taskIsNotCaughtUpOnClientAndOtherMoreCaughtUpClientsExist(task, client, clientStates, tasksToCaughtUpClients, tasksToClientByLag)) {
                    taskMovements.add(new TaskMovement(task, client, tasksToCaughtUpClients.get(task)));
                }
            }
            caughtUpClientsByTaskLoad.offer(client);
        }

        final int movementsNeeded = taskMovements.size();

        while (!taskMovements.isEmpty()) {
            final TaskMovement movement = taskMovements.poll();
            // Attempt to find a caught up standby, otherwise find any caught up client, failing that use the most
            // caught up client.
            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 as scheduled before but none exist");
            }
        }

        return movementsNeeded;
    }
@Test
    public void shouldMoveTasksToCaughtUpClientsAndAssignWarmupReplicasInTheirPlace() {
        final int maxWarmupReplicas = Integer.MAX_VALUE;
        final Set<TaskId> 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<ProcessId, ClientState> clientStates = getClientStatesMap(client1, client2, client3);

        final Map<TaskId, SortedSet<ProcessId>> tasksToCaughtUpClients = mkMap(
            mkEntry(TASK_0_0, mkSortedSet(PID_1)),
            mkEntry(TASK_0_1, mkSortedSet(PID_3)),
            mkEntry(TASK_0_2, mkSortedSet(PID_2))
        );
        final Map<TaskId, SortedSet<ProcessId>> tasksToClientByLag = mkMap(
            mkEntry(TASK_0_0, mkOrderedSet(PID_1, PID_2, PID_3)),
            mkEntry(TASK_0_1, mkOrderedSet(PID_3, PID_1, PID_2)),
            mkEntry(TASK_0_2, mkOrderedSet(PID_2, PID_1, PID_3))
        );

        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 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)));
    }