MockStore<ServiceProperties> serviceRegistry = new MockStore<ServiceProperties>();
MockStore<ClusterProperties> clusterRegistry = new MockStore<ClusterProperties>();
MockStore<UriProperties> uriRegistry = new MockStore<UriProperties>();
ScheduledExecutorService executorService = new SynchronousExecutorService();
//loadBalancerStrategyFactories.put("rr", new RandomLoadBalancerStrategyFactory());
loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
// PrpcClientFactory();
clientFactories.put("http", new DoNothingClientFactory()); // new
// HttpClientFactory();
SimpleLoadBalancerState state =
new SimpleLoadBalancerState(executorService,
uriRegistry,
clusterRegistry,
serviceRegistry,
clientFactories,
loadBalancerStrategyFactories);
SimpleLoadBalancer loadBalancer =
new SimpleLoadBalancer(state, 5, TimeUnit.SECONDS);
FutureCallback<None> balancerCallback = new FutureCallback<None>();
loadBalancer.start(balancerCallback);
balancerCallback.get();
URI uri1 = URI.create("http://test.qa1.com:1234");
URI uri2 = URI.create("http://test.qa2.com:2345");
URI uri3 = URI.create("http://test.qa3.com:6789");
Map<Integer, PartitionData> partitionData = new HashMap<Integer, PartitionData>(1);
partitionData.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1d));
Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<URI, Map<Integer, PartitionData>>(3);
uriData.put(uri1, partitionData);
uriData.put(uri2, partitionData);
uriData.put(uri3, partitionData);
prioritizedSchemes.add("http");
clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1"));
serviceRegistry.put("foo", new ServiceProperties("foo",
"cluster-1",
"/foo",
"degrader",
Collections.<String>emptyList(),
Collections.<String,Object>emptyMap(),
null,
null,
prioritizedSchemes,
null));
uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
URI expectedUri1 = URI.create("http://test.qa1.com:1234/foo");
URI expectedUri2 = URI.create("http://test.qa2.com:2345/foo");
URI expectedUri3 = URI.create("http://test.qa3.com:6789/foo");
Set<URI> expectedUris = new HashSet<URI>();
expectedUris.add(expectedUri1);
expectedUris.add(expectedUri2);
expectedUris.add(expectedUri3);
for (int i = 0; i < 100; ++i)
{
RewriteClient client =
(RewriteClient) loadBalancer.getClient(new URIRequest("d2://foo/52"),
new RequestContext());
assertTrue(expectedUris.contains(client.getUri()));
assertEquals(client.getUri().getScheme(), "http");
}
final CountDownLatch latch = new CountDownLatch(1);
PropertyEventShutdownCallback callback = new PropertyEventShutdownCallback()
{
@Override
public void done()
{
latch.countDown();
}
};
state.shutdown(callback);
if (!latch.await(60, TimeUnit.SECONDS))
{
fail("unable to shutdown state");
}
executorService.shutdownNow();
assertTrue(executorService.isShutdown(), "ExecutorService should have shut down!");
}
}