private void rebalance() {
        int activeNodes = (int) clusterService.getNodes()
                .stream()
                .filter(node -> clusterService.getState(node.id()).isActive())
                .count();

        int myShare = (int) Math.ceil((double) NUM_PARTITIONS / activeNodes);

        // First make sure this node is a candidate for all partitions.
        IntStream.range(0, NUM_PARTITIONS)
                 .mapToObj(this::getPartitionPath)
                 .map(leadershipService::getLeadership)
                 .filter(leadership -> !leadership.candidates().contains(localNodeId))
                 .map(Leadership::topic)
                 .forEach(leadershipService::runForLeadership);

        List<String> myPartitions = IntStream.range(0, NUM_PARTITIONS)
                                             .mapToObj(this::getPartitionPath)
                                             .map(leadershipService::getLeadership)
                                             .filter(Objects::nonNull)
                                             .filter(leadership -> localNodeId.equals(leadership.leaderNodeId()))
                                             .map(Leadership::topic)
                                             .collect(Collectors.toList());

        int relinquish = myPartitions.size() - myShare;


        for (int i = 0; i < relinquish; i++) {
            String topic = myPartitions.get(i);
            // Wait till all active nodes are in contention for partition ownership.
            // This avoids too many relinquish/reclaim cycles.
            if (leadershipService.getCandidates(topic).size() == activeNodes) {
                leadershipService.withdraw(topic);
                executor.schedule(() -> recontest(topic), BACKOFF_TIME, TimeUnit.SECONDS);
            }
        }
    }
@Test
    public void testRebalance() {
        // We have all the partitions so we'll need to relinquish some
        setUpLeadershipService(WorkPartitionManager.NUM_PARTITIONS);

        leadershipService.withdraw(anyString());
        expectLastCall().times(7);

        replay(leadershipService);

        partitionManager.activate();

        // trigger rebalance
        partitionManager.doRebalance();

        verify(leadershipService);
    }