@Override
    public List<Assignment> rebalance(List<Assignment> currentAssignments, Set<String> workers) {

        Map<String, List<Instance>> workerToAssignmentMap = new HashMap<>();

        workers.forEach(workerId -> workerToAssignmentMap.put(workerId, new LinkedList<>()));

        currentAssignments.forEach(
                assignment -> workerToAssignmentMap.computeIfAbsent(assignment.getWorkerId(), s -> new LinkedList<>())
                        .add(assignment.getInstance()));

        List<Assignment> newAssignments = new LinkedList<>();

        int iterations = 0;
        while (true) {
            iterations++;

            Map.Entry<String, List<Instance>> mostAssignmentsWorker =
                    findWorkerWithMostAssignments(workerToAssignmentMap);

            Map.Entry<String, List<Instance>> leastAssignmentsWorker =
                    findWorkerWithLeastAssignments(workerToAssignmentMap);

            if (mostAssignmentsWorker.getValue().size() == leastAssignmentsWorker.getValue().size()
                    || mostAssignmentsWorker.getValue().size() == leastAssignmentsWorker.getValue().size() + 1) {
                break;
            }

            String mostAssignmentsWorkerId = mostAssignmentsWorker.getKey();
            String leastAssignmentsWorkerId = leastAssignmentsWorker.getKey();

            Queue<Instance> src = (Queue) workerToAssignmentMap.get(mostAssignmentsWorkerId);
            Queue<Instance> dest = (Queue) workerToAssignmentMap.get(leastAssignmentsWorkerId);

            Instance instance = src.poll();
            Assignment newAssignment = Assignment.newBuilder()
                    .setInstance(instance)
                    .setWorkerId(leastAssignmentsWorkerId)
                    .build();
            newAssignments.add(newAssignment);

            dest.add(instance);
        }

        log.info("Rebalance - iterations: {}", iterations);

        return newAssignments;
    }
@Test
    public void testRebalance() {
        Function.FunctionMetaData function1 = Function.FunctionMetaData.newBuilder()
                .setFunctionDetails(Function.FunctionDetails.newBuilder().setName("func-1")
                        .setNamespace("namespace-1").setTenant("tenant-1").setParallelism(1)).setVersion(0)
                .build();

        List<Function.Assignment> assignments = new LinkedList<>();
        for (int i = 0; i < 10; i++) {
            Function.Assignment assignment1 = Function.Assignment.newBuilder()
                    .setWorkerId("worker-1")
                    .setInstance(Function.Instance.newBuilder()
                            .setFunctionMetaData(function1).setInstanceId(i).build())
                    .build();

            assignments.add(assignment1);
        }

        Set<String> workers = new HashSet<>();
        for (int i = 0; i < 3; i++) {
            workers.add("worker-" + i);
        }

        RoundRobinScheduler roundRobinScheduler = new RoundRobinScheduler();

        List<Function.Assignment> newAssignments = roundRobinScheduler.rebalance(assignments, workers);

        Map<String, Integer> workerAssignments = new HashMap<>();
        for (Function.Assignment assignment : newAssignments) {
            Integer count = workerAssignments.get((assignment.getWorkerId()));
            if (count == null) {
                count = 0;
            }
            count++;
            workerAssignments.put(assignment.getWorkerId(), count);
        }

        Assert.assertEquals(workerAssignments.size(), 2);
        for (Map.Entry<String, Integer> entry : workerAssignments.entrySet()) {
            Assert.assertEquals(entry.getValue().intValue(), 3);
        }
    }