@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 shouldDoThePartialRackAwareness() {
        final Map<ProcessId, ClientState> clientStates = mkMap(
            mkEntry(PID_1, createClientStateWithCapacity(PID_1, 1, mkMap(mkEntry(CLUSTER_TAG, CLUSTER_1), mkEntry(ZONE_TAG, ZONE_1)), TASK_0_0)),
            mkEntry(PID_2, createClientStateWithCapacity(PID_2, 1, mkMap(mkEntry(CLUSTER_TAG, CLUSTER_1), mkEntry(ZONE_TAG, ZONE_2)))),
            mkEntry(PID_3, createClientStateWithCapacity(PID_3, 1, mkMap(mkEntry(CLUSTER_TAG, CLUSTER_1), mkEntry(ZONE_TAG, ZONE_3)))),

            mkEntry(PID_4, createClientStateWithCapacity(PID_4, 1, mkMap(mkEntry(CLUSTER_TAG, CLUSTER_2), mkEntry(ZONE_TAG, ZONE_1)))),
            mkEntry(PID_5, createClientStateWithCapacity(PID_5, 1, mkMap(mkEntry(CLUSTER_TAG, CLUSTER_2), mkEntry(ZONE_TAG, ZONE_2)))),
            mkEntry(PID_6, createClientStateWithCapacity(PID_6, 1, mkMap(mkEntry(CLUSTER_TAG, CLUSTER_2), mkEntry(ZONE_TAG, ZONE_3)), TASK_1_0))
        );

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

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

        // We need to distribute 2 standby tasks (+1 active task).
        // Since we have only two unique `cluster` tag values,
        // we can only achieve "ideal" distribution on the 1st standby task assignment.
        // We can't consider the `cluster` tag for the 2nd standby task assignment because the 1st standby
        // task would already be assigned on different clusters compared to the active one, which means
        // we have already used all the available cluster tag values. Taking the `cluster` tag into consideration
        // for the  2nd standby task assignment would affectively mean excluding all the clients.
        // Instead, for the 2nd standby task, we can only achieve partial rack awareness based on the `zone` tag.
        // As we don't consider the `cluster` tag for the 2nd standby task assignment, partial rack awareness
        // can be satisfied by placing the 2nd standby client on a different `zone` tag compared to active and corresponding standby tasks.
        // The `zone` on either `cluster` tags are valid candidates for the partial rack awareness, as our goal is to distribute clients on the different `zone` tags.

        Stream.of(PID_2, PID_5).forEach(client -> assertStandbyTaskCountForClientEqualsTo(clientStates, client, 1));
        // There's no strong guarantee where 2nd standby task will end up.
        Stream.of(PID_1, PID_3, PID_4, PID_6).forEach(client -> assertStandbyTaskCountForClientEqualsTo(clientStates, client, 0, 1));
        assertTotalNumberOfStandbyTasksEqualsTo(clientStates, 4);

        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_0_0,
                clientStates,
                asList(
                    // Since it's located on a different `cluster` and `zone` tag dimensions,
                    // `PID_5` is the "ideal" distribution for the 1st standby task assignment.
                    // For the 2nd standby, either `PID_3` or `PID_6` are valid destinations as
                    // we need to distribute the clients on different `zone`
                    // tags without considering the `cluster` tag value.
                    mkSet(PID_5, PID_3),
                    mkSet(PID_5, PID_6)
                )
            )
        );
        assertTrue(
            standbyClientsHonorRackAwareness(
                TASK_1_0,
                clientStates,
                asList(
                    // The same comment as above applies here too.
                    // `PID_2` is the ideal distribution on different `cluster`
                    // and `zone` tag dimensions. In contrast, `PID_4` and `PID_1`
                    // satisfy only the partial rack awareness as they are located on a different `zone` tag dimension.
                    mkSet(PID_2, PID_4),
                    mkSet(PID_2, PID_1)
                )
            )
        );
    }