SimpleBalancer implements Balancer { @Override public void balanceAll(String source, int rebalanceConcurrency, Optional<RateLimiter> rebalanceRateLimiter) { String target; DistributedLogClient sourceClient, targetClient; MonitorServiceClient sourceMonitor, targetMonitor; if (target1.equals(source)) { sourceClient = targetClient1; sourceMonitor = targetMonitor1; target = target2; targetClient = targetClient2; targetMonitor = targetMonitor2; } else if (target2.equals(source)) { sourceClient = targetClient2; sourceMonitor = targetMonitor2; target = target1; targetClient = targetClient1; targetMonitor = targetMonitor1; } else { throw new IllegalArgumentException("Unknown target " + source); } Map<SocketAddress, Set<String>> distribution = sourceMonitor.getStreamOwnershipDistribution(); if (distribution.isEmpty()) { return; } StreamChooser streamChooser = new CountBasedStreamChooser(distribution); StreamMover streamMover = new StreamMoverImpl(source, sourceClient, sourceMonitor, 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 testBalanceAll() throws Exception { String namePrefix = "simplebalancer-balance-all-"; 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 from 'unknown' target"); try { balancer.balanceAll("unknown", 10, rateLimiter); fail("Should fail on balanceAll from 'unknown' target."); } catch (IllegalArgumentException iae) { } logger.info("Rebalancing from 'target' target"); balancer.balanceAll("target", 1, rateLimiter); 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); } logger.info("Rebalancing from 'source' target"); balancer.balanceAll("source", 10, rateLimiter); for (int i = 0; i < numStreams; i++) { String name = namePrefix + i; checkStream(name, targetClient, targetServer, 1, numStreams, numStreams, true, true); checkStream(name, dlClient, dlServer, 0, 0, 0, false, false); } }
@Test(timeout = 60000) public void testBalanceAll() throws Exception { String namePrefix = "simplebalancer-balance-all-"; 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(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 from 'unknown' target"); try { balancer.balanceAll("unknown", 10, rateLimiter); fail("Should fail on balanceAll from 'unknown' target."); } catch (IllegalArgumentException iae) { } logger.info("Rebalancing from 'target' target"); balancer.balanceAll("target", 1, rateLimiter); 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); } logger.info("Rebalancing from 'source' target"); balancer.balanceAll("source", 10, rateLimiter); for (int i = 0; i < numStreams; i++) { String name = namePrefix + i; checkStream(name, targetClient, targetServer, 1, numStreams, numStreams, true, true); checkStream(name, dlClient, dlServer, 0, 0, 0, false, false); } }