@Override
    public boolean assign(final Map<ProcessId, ClientState> clients,
                          final Set<TaskId> allTaskIds,
                          final Set<TaskId> statefulTaskIds,
                          final AssignmentConfigs configs) {
        final int numStandbyReplicas = configs.numStandbyReplicas();
        final Set<String> rackAwareAssignmentTags = new HashSet<>(tagsFunction.apply(configs));

        final Map<TaskId, Integer> tasksToRemainingStandbys = computeTasksToRemainingStandbys(
            numStandbyReplicas,
            statefulTaskIds
        );

        final Map<String, Set<String>> tagKeyToValues = new HashMap<>();
        final Map<TagEntry, Set<ProcessId>> tagEntryToClients = new HashMap<>();

        fillClientsTagStatistics(clients, tagEntryToClients, tagKeyToValues);

        final ConstrainedPrioritySet standbyTaskClientsByTaskLoad = createLeastLoadedPrioritySetConstrainedByAssignedTask(clients);

        final Map<TaskId, ProcessId> pendingStandbyTasksToClientId = new HashMap<>();

        for (final TaskId statefulTaskId : statefulTaskIds) {
            for (final Map.Entry<ProcessId, ClientState> entry : clients.entrySet()) {
                final ProcessId clientId = entry.getKey();
                final ClientState clientState = entry.getValue();

                if (clientState.activeTasks().contains(statefulTaskId)) {
                    assignStandbyTasksToClientsWithDifferentTags(
                        numStandbyReplicas,
                        standbyTaskClientsByTaskLoad,
                        statefulTaskId,
                        clientId,
                        rackAwareAssignmentTags,
                        clients,
                        tasksToRemainingStandbys,
                        tagKeyToValues,
                        tagEntryToClients,
                        pendingStandbyTasksToClientId
                    );
                }
            }
        }

        if (!tasksToRemainingStandbys.isEmpty()) {
            assignPendingStandbyTasksToLeastLoadedClients(clients,
                                                          numStandbyReplicas,
                                                          standbyTaskClientsByTaskLoad,
                                                          tasksToRemainingStandbys);
        }

        // returning false, because standby task assignment will never require a follow-up probing rebalance.
        return false;
    }
@Test
    public void shouldDistributeStandbyTasksWhenActiveTasksAreLocatedOnSameZone() {
        final Map<ProcessId, ClientState> clientStates = mkMap(
            mkEntry(PID_1, createClientStateWithCapacity(PID_1, 2, mkMap(mkEntry(ZONE_TAG, ZONE_1), mkEntry(CLUSTER_TAG, CLUSTER_1)), TASK_0_0, TASK_1_0)),
            mkEntry(PID_2, createClientStateWithCapacity(PID_2, 2, mkMap(mkEntry(ZONE_TAG, ZONE_2), mkEntry(CLUSTER_TAG, CLUSTER_1)))),
            mkEntry(PID_3, createClientStateWithCapacity(PID_3, 2, mkMap(mkEntry(ZONE_TAG, ZONE_3), mkEntry(CLUSTER_TAG, CLUSTER_1)))),

            mkEntry(PID_4, createClientStateWithCapacity(PID_4, 2, mkMap(mkEntry(ZONE_TAG, ZONE_1), mkEntry(CLUSTER_TAG, CLUSTER_2)), TASK_0_1, TASK_1_1)),
            mkEntry(PID_5, createClientStateWithCapacity(PID_5, 2, mkMap(mkEntry(ZONE_TAG, ZONE_2), mkEntry(CLUSTER_TAG, CLUSTER_2)))),
            mkEntry(PID_6, createClientStateWithCapacity(PID_6, 2, mkMap(mkEntry(ZONE_TAG, ZONE_3), mkEntry(CLUSTER_TAG, CLUSTER_2)))),

            mkEntry(PID_7, createClientStateWithCapacity(PID_7, 2, mkMap(mkEntry(ZONE_TAG, ZONE_1), mkEntry(CLUSTER_TAG, CLUSTER_3)), TASK_0_2, TASK_1_2)),
            mkEntry(PID_8, createClientStateWithCapacity(PID_8, 2, mkMap(mkEntry(ZONE_TAG, ZONE_2), mkEntry(CLUSTER_TAG, CLUSTER_3)))),
            mkEntry(PID_9, createClientStateWithCapacity(PID_9, 2, mkMap(mkEntry(ZONE_TAG, ZONE_3), mkEntry(CLUSTER_TAG, CLUSTER_3))))
        );

        final Set<TaskId> allActiveTasks = findAllActiveTasks(clientStates);
        final AssignmentConfigs assignmentConfigs = newAssignmentConfigs(2, ZONE_TAG, CLUSTER_TAG);

        standbyTaskAssignor.assign(clientStates, allActiveTasks, allActiveTasks, assignmentConfigs);

        assertTrue(clientStates.values().stream().allMatch(ClientState::reachedCapacity));

        Stream.of(PID_1, PID_4, PID_7).forEach(client -> assertStandbyTaskCountForClientEqualsTo(clientStates, client, 0));
        Stream.of(PID_2, PID_3, PID_5, PID_6, PID_8, PID_9).forEach(client -> assertStandbyTaskCountForClientEqualsTo(clientStates, client, 2));
        assertTotalNumberOfStandbyTasksEqualsTo(clientStates, 12);

        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_0_0,
                clientStates,
                asList(
                    mkSet(PID_9, PID_5), mkSet(PID_6, PID_8)
                )
            )
        );
        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_1_0,
                clientStates,
                asList(
                    mkSet(PID_9, PID_5), mkSet(PID_6, PID_8)
                )
            )
        );

        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_0_1,
                clientStates,
                asList(
                    mkSet(PID_2, PID_9), mkSet(PID_3, PID_8)
                )
            )
        );
        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_1_1,
                clientStates,
                asList(
                    mkSet(PID_2, PID_9), mkSet(PID_3, PID_8)
                )
            )
        );

        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_0_2,
                clientStates,
                asList(
                    mkSet(PID_5, PID_3), mkSet(PID_2, PID_6)
                )
            )
        );
        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_1_2,
                clientStates,
                asList(
                    mkSet(PID_5, PID_3), mkSet(PID_2, PID_6)
                )
            )
        );
    }