SimpleBalancer implements Balancer { @Override public void balance(int rebalanceWaterMark, double rebalanceTolerancePercentage, int rebalanceConcurrency, Optional<RateLimiter> rebalanceRateLimiter) { Map<SocketAddress, Set<String>> distribution1 = targetMonitor1.getStreamOwnershipDistribution(); Map<SocketAddress, Set<String>> distribution2 = targetMonitor2.getStreamOwnershipDistribution(); int proxyCount1 = distribution1.size(); int streamCount1 = countNumberStreams(distribution1); int proxyCount2 = distribution2.size(); int streamCount2 = countNumberStreams(distribution2); logger.info("'{}' has {} streams by {} proxies; while '{}' has {} streams by {} proxies.", new Object[] {target1, streamCount1, proxyCount1, target2, streamCount2, proxyCount2 }); String source, target; Map<SocketAddress, Set<String>> srcDistribution; DistributedLogClient srcClient, targetClient; MonitorServiceClient srcMonitor, targetMonitor; int srcStreamCount, targetStreamCount; if (streamCount1 > streamCount2) { source = target1; srcStreamCount = streamCount1; srcClient = targetClient1; srcMonitor = targetMonitor1; srcDistribution = distribution1; target = target2; targetStreamCount = streamCount2; targetClient = targetClient2; targetMonitor = targetMonitor2; } else { source = target2; srcStreamCount = streamCount2; srcClient = targetClient2; srcMonitor = targetMonitor2; srcDistribution = distribution2; target = target1; targetStreamCount = streamCount1; targetClient = targetClient1; targetMonitor = targetMonitor1; } Map<String, Integer> loadDistribution = new HashMap<String, Integer>(); loadDistribution.put(source, srcStreamCount); loadDistribution.put(target, targetStreamCount); int numStreamsToRebalance = BalancerUtils.calculateNumStreamsToRebalance(source, loadDistribution, rebalanceWaterMark, rebalanceTolerancePercentage); if (numStreamsToRebalance <= 0) { logger.info("No streams need to be rebalanced from '{}' to '{}'.", source, target); return; } StreamChooser streamChooser = LimitedStreamChooser.of(new CountBasedStreamChooser(srcDistribution), numStreamsToRebalance); StreamMover streamMover = new StreamMoverImpl(source, srcClient, srcMonitor, target, targetClient, targetMonitor); moveStreams(streamChooser, streamMover, rebalanceConcurrency, rebalanceRateLimiter); } SimpleBalancer(String name1,
                          DistributedLogClient client1,
                          MonitorServiceClient monitor1,
                          String name2,
                          DistributedLogClient client2,
                          MonitorServiceClient monitor2); @Override void balance(int rebalanceWaterMark,
                        double rebalanceTolerancePercentage,
                        int rebalanceConcurrency,
                        Optional<RateLimiter> rebalanceRateLimiter); @Override void balanceAll(String source,
                           int rebalanceConcurrency,
                           Optional<RateLimiter> rebalanceRateLimiter); @Override void close();  }
@Test(timeout = 60000) public void testBalanceStreams() throws Exception { String namePrefix = "simplebalancer-balance-streams-"; int numStreams = 10; for (int i = 0; i < numStreams; i++) { String name = namePrefix + i; dlClient.routingService.addHost(name, dlServer.getAddress()); targetClient.routingService.addHost(name, targetServer.getAddress()); } for (int i = 0; i < numStreams; i++) { String name = namePrefix + i; Await.result(((DistributedLogClient) dlClient.dlClient).write(name, ByteBuffer.wrap(("" + i).getBytes(UTF_8)))); } for (int i = 0; i < numStreams; i++) { String name = namePrefix + i; checkStream(name, dlClient, dlServer, 1, numStreams, numStreams, true, true); checkStream(name, targetClient, targetServer, 0, 0, 0, false, false); } Optional<RateLimiter> rateLimiter = Optional.absent(); Balancer balancer = new SimpleBalancer("source", dlClient.dlClient, dlClient.dlClient, "target", targetClient.dlClient, targetClient.dlClient); logger.info("Rebalancing streams between targets"); balancer.balance(0, 0, 10, rateLimiter); Set<String> sourceStreams = getAllStreamsFromDistribution(getStreamOwnershipDistribution(dlClient)); Set<String> targetStreams = getAllStreamsFromDistribution(getStreamOwnershipDistribution(targetClient)); assertEquals(numStreams / 2, sourceStreams.size()); assertEquals(numStreams / 2, targetStreams.size()); for (String name : sourceStreams) { checkStream(name, dlClient, dlServer, 1, numStreams / 2, numStreams / 2, true, true); checkStream(name, targetClient, targetServer, 1, numStreams / 2, numStreams / 2, false, false); } for (String name : targetStreams) { checkStream(name, targetClient, targetServer, 1, numStreams / 2, numStreams / 2, true, true); checkStream(name, dlClient, dlServer, 1, numStreams / 2, numStreams / 2, false, false); } }