use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.
the class SimpleLoadBalancerStrawMan method main.
public static void main(String[] args) throws URISyntaxException, ServiceUnavailableException {
// define the load balancing strategies that we support (round robin, etc)
Map<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>> loadBalancerStrategyFactories = new HashMap<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>>();
loadBalancerStrategyFactories.put("rr", new RandomLoadBalancerStrategyFactory());
loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
// define the clients that we support (http, etc)
Map<String, TransportClientFactory> clientFactories = new HashMap<String, TransportClientFactory>();
clientFactories.put("http", new HttpClientFactory());
// listen for service updates (could be a glu discovery client, zk discovery client,
// config discovery client, etc)
ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
MockStore<ServiceProperties> serviceRegistry = new MockStore<ServiceProperties>();
MockStore<ClusterProperties> clusterRegistry = new MockStore<ClusterProperties>();
MockStore<UriProperties> uriRegistry = new MockStore<UriProperties>();
SimpleLoadBalancerState state = new SimpleLoadBalancerState(executorService, uriRegistry, clusterRegistry, serviceRegistry, clientFactories, loadBalancerStrategyFactories);
// create the load balancer
SimpleLoadBalancer loadBalancer = new SimpleLoadBalancer(state);
final TransportClient tc = loadBalancer.getClient(new URIRequest("d2://browsemaps/52"), new RequestContext());
final Client c = new TransportClientAdapter(tc, true);
c.restRequest(null);
}
use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.
the class SimpleLoadBalancerTest method testLoadBalancerWithPartitionsSmoke.
// load balancer working with partitioned cluster
@Test(groups = { "small", "back-end" })
public void testLoadBalancerWithPartitionsSmoke() throws URISyntaxException, ServiceUnavailableException, InterruptedException, ExecutionException {
for (int tryAgain = 0; tryAgain < 12; ++tryAgain) {
Map<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>> loadBalancerStrategyFactories = new HashMap<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>>();
Map<String, TransportClientFactory> clientFactories = new HashMap<String, TransportClientFactory>();
List<String> prioritizedSchemes = new ArrayList<String>();
MockStore<ServiceProperties> serviceRegistry = new MockStore<ServiceProperties>();
MockStore<ClusterProperties> clusterRegistry = new MockStore<ClusterProperties>();
MockStore<UriProperties> uriRegistry = new MockStore<UriProperties>();
ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
clientFactories.put("http", new DoNothingClientFactory());
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<URI, Double> uris = new HashMap<URI, Double>();
uris.put(uri1, 1d);
uris.put(uri2, 1d);
uris.put(uri3, 1d);
Map<URI, Map<Integer, PartitionData>> partitionDesc = new HashMap<URI, Map<Integer, PartitionData>>();
Map<Integer, PartitionData> server1 = new HashMap<Integer, PartitionData>();
server1.put(0, new PartitionData(1d));
server1.put(1, new PartitionData(1d));
Map<Integer, PartitionData> server2 = new HashMap<Integer, PartitionData>();
server2.put(0, new PartitionData(1d));
Map<Integer, PartitionData> server3 = new HashMap<Integer, PartitionData>();
server3.put(1, new PartitionData(1d));
partitionDesc.put(uri1, server1);
partitionDesc.put(uri2, server2);
partitionDesc.put(uri3, server3);
prioritizedSchemes.add("http");
int partitionMethod = tryAgain % 4;
switch(partitionMethod) {
case 0:
clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<String, String>(), new HashSet<URI>(), new RangeBasedPartitionProperties("id=(\\d+)", 0, 50, 2)));
break;
case 1:
clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<String, String>(), new HashSet<URI>(), new HashBasedPartitionProperties("id=(\\d+)", 2, HashBasedPartitionProperties.HashAlgorithm.valueOf("MODULO"))));
break;
case 2:
clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<String, String>(), new HashSet<URI>(), new HashBasedPartitionProperties("id=(\\d+)", 2, HashBasedPartitionProperties.HashAlgorithm.valueOf("MD5"))));
break;
case 3:
// test getRings with gap. here, no server serves partition 2
clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<String, String>(), new HashSet<URI>(), new RangeBasedPartitionProperties("id=(\\d+)", 0, 50, 4)));
server3.put(3, new PartitionData(1d));
partitionDesc.put(uri3, server3);
break;
default:
break;
}
serviceRegistry.put("foo", new ServiceProperties("foo", "cluster-1", "/foo", Arrays.asList("degrader"), Collections.<String, Object>emptyMap(), null, null, prioritizedSchemes, null));
uriRegistry.put("cluster-1", new UriProperties("cluster-1", partitionDesc));
if (partitionMethod == 3) {
Map<Integer, Ring<URI>> ringMap = loadBalancer.getRings(URI.create("d2://foo"));
assertEquals(ringMap.size(), 4);
// the ring for partition 2 should be empty
assertEquals(ringMap.get(2).toString(), new ConsistentHashRing<URI>(Collections.emptyList()).toString());
continue;
}
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 < 1000; ++i) {
int ii = i % 100;
RewriteClient client = (RewriteClient) loadBalancer.getClient(new URIRequest("d2://foo/id=" + ii), new RequestContext());
String clientUri = client.getUri().toString();
HashFunction<String[]> hashFunction = null;
String[] str = new String[1];
// test KeyMapper target host hint: request is always to target host regardless of what's in d2 URI and whether it's hash-based or range-based partitions
RequestContext requestContextWithHint = new RequestContext();
KeyMapper.TargetHostHints.setRequestContextTargetHost(requestContextWithHint, uri1);
RewriteClient hintedClient1 = (RewriteClient) loadBalancer.getClient(new URIRequest("d2://foo/id=" + ii), requestContextWithHint);
String hintedUri1 = hintedClient1.getUri().toString();
Assert.assertEquals(hintedUri1, uri1.toString() + "/foo");
RewriteClient hintedClient2 = (RewriteClient) loadBalancer.getClient(new URIRequest("d2://foo/action=purge-all"), requestContextWithHint);
String hintedUri2 = hintedClient2.getUri().toString();
Assert.assertEquals(hintedUri2, uri1.toString() + "/foo");
if (partitionMethod == 2) {
hashFunction = new MD5Hash();
}
for (URI uri : expectedUris) {
if (clientUri.contains(uri.toString())) {
// check if only key belonging to partition 0 gets uri2
if (uri.equals(uri2)) {
if (partitionMethod == 0) {
assertTrue(ii < 50);
} else if (partitionMethod == 1) {
assertTrue(ii % 2 == 0);
} else {
str[0] = ii + "";
assertTrue(hashFunction.hash(str) % 2 == 0);
}
}
// check if only key belonging to partition 1 gets uri3
if (uri.equals(uri3)) {
if (partitionMethod == 0) {
assertTrue(ii >= 50);
} else if (partitionMethod == 1) {
assertTrue(ii % 2 == 1);
} else {
str[0] = ii + "";
assertTrue(hashFunction.hash(str) % 2 == 1);
}
}
}
}
}
// two rings for two partitions
Map<Integer, Ring<URI>> ringMap = loadBalancer.getRings(URI.create("d2://foo"));
assertEquals(ringMap.size(), 2);
if (partitionMethod != 2) {
Set<String> keys = new HashSet<String>();
for (int j = 0; j < 50; j++) {
if (partitionMethod == 0) {
keys.add(j + "");
} else {
keys.add(j * 2 + "");
}
}
// if it is range based partition, all keys from 0 ~ 49 belong to partition 0 according to the range definition
// if it is modulo based partition, all even keys belong to partition 0 because the partition count is 2
// only from partition 0
MapKeyResult<Ring<URI>, String> mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
Map<Ring<URI>, Collection<String>> keyToPartition = mapKeyResult.getMapResult();
assertEquals(keyToPartition.size(), 1);
for (Ring<URI> ring : keyToPartition.keySet()) {
assertEquals(ring, ringMap.get(0));
}
// now also from partition 1
keys.add("51");
mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
assertEquals(mapKeyResult.getMapResult().size(), 2);
assertEquals(mapKeyResult.getUnmappedKeys().size(), 0);
// now only from partition 1
keys.clear();
keys.add("99");
mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
keyToPartition = mapKeyResult.getMapResult();
assertEquals(keyToPartition.size(), 1);
assertEquals(mapKeyResult.getUnmappedKeys().size(), 0);
for (Ring<URI> ring : keyToPartition.keySet()) {
assertEquals(ring, ringMap.get(1));
}
keys.add("100");
mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
if (partitionMethod == 0) {
// key out of range
Collection<MapKeyResult.UnmappedKey<String>> unmappedKeys = mapKeyResult.getUnmappedKeys();
assertEquals(unmappedKeys.size(), 1);
}
try {
loadBalancer.getClient(new URIRequest("d2://foo/id=100"), new RequestContext());
if (partitionMethod == 0) {
// key out of range
fail("Should throw ServiceUnavailableException caused by PartitionAccessException");
}
} catch (ServiceUnavailableException e) {
}
}
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!");
}
}
use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.
the class SimpleLoadBalancerTest method testLoadBalancerWithWait.
@Test(groups = { "small", "back-end" })
public void testLoadBalancerWithWait() throws URISyntaxException, ServiceUnavailableException, InterruptedException {
URIRequest uriRequest = new URIRequest("d2://NonExistentService");
LoadBalancerTestState state = new LoadBalancerTestState();
SimpleLoadBalancer balancer = new SimpleLoadBalancer(state, 5, TimeUnit.SECONDS);
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.listenToService = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.isListeningToService = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.getServiceProperties = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.listenToCluster = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.isListeningToCluster = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.getClusterProperties = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.getUriProperties = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.getClient = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.getStrategy = true;
try {
balancer.getClient(uriRequest, new RequestContext());
fail("should have received a service unavailable exception");
} catch (ServiceUnavailableException e) {
}
state.getPartitionAccessor = true;
// victory
assertNotNull(balancer.getClient(uriRequest, new RequestContext()));
}
use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.
the class SimpleLoadBalancerSimulation method reset.
/**
* Reset the entire state of the simulation.
*
*/
public void reset() {
// simulation state
_random = new Random();
_possibleServices = Collections.synchronizedList(new ArrayList<String>());
_possibleClusters = Collections.synchronizedList(new ArrayList<String>());
_possiblePaths = Collections.synchronizedList(new ArrayList<String>());
_possibleSchemes = Collections.synchronizedList(new ArrayList<String>());
_possibleStrategies = Collections.synchronizedList(new ArrayList<String>());
_possibleUris = Collections.synchronizedList(new ArrayList<URI>());
// load balancer state
_executorService = Executors.newSingleThreadScheduledExecutor();
;
// pretend that these are zk stores
_serviceRegistry = new MockStore<ServiceProperties>();
_uriRegistry = new MockStore<UriProperties>();
_clusterRegistry = new MockStore<ClusterProperties>();
_loadBalancerStrategyFactories = new HashMap<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>>();
_clientFactories = new HashMap<String, TransportClientFactory>();
_state = new SimpleLoadBalancerState(_executorService, _uriRegistry, _clusterRegistry, _serviceRegistry, _clientFactories, _loadBalancerStrategyFactories);
_loadBalancer = new SimpleLoadBalancer(_state, 10, TimeUnit.SECONDS);
FutureCallback<None> callback = new FutureCallback<None>();
_loadBalancer.start(callback);
try {
callback.get();
} catch (Exception e) {
throw new RuntimeException("Balancer start failed", e);
}
// verification state
_expectedServiceProperties = new ConcurrentHashMap<String, ServiceProperties>();
_expectedClusterProperties = new ConcurrentHashMap<String, ClusterProperties>();
_expectedUriProperties = new ConcurrentHashMap<String, UriProperties>();
_totalMessages = 0;
// TODO parameterize this
for (int i = 0; i < 10; ++i) {
_possibleServices.add("service-" + i);
_possibleClusters.add("cluster-" + i);
_possiblePaths.add("/some/path/" + i);
_possibleSchemes.add("scheme" + i % 3);
_possibleStrategies.add("strategy-" + i);
_clientFactories.put("scheme" + i % 2, new DoNothingClientFactory());
_loadBalancerStrategyFactories.put("strategy-" + i, _loadBalancerStrategyFactoryToTest);
}
for (int i = 0; i < 1000; ++i) {
_possibleUris.add(URI.create(random(_possibleSchemes) + "://host" + i % 100 + ":" + (1000 + _random.nextInt(1000)) + random(_possiblePaths)));
}
// add bad stuff
// add a bad scheme to prioritized schemes
_possibleSchemes.add("BAD_PRIORITIZED_SCHEME");
// add a bad scheme to possible uris
_possibleUris.add(URI.create("BADSCHEME://host1001:" + (1000 + _random.nextInt(1000)) + random(_possiblePaths)));
// register jmx goodies
new JmxManager().registerLoadBalancer("SimpleLoadBalancer", _loadBalancer).registerLoadBalancerState("SimpleLoadBalancerState", _state);
}
use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.
the class SimpleLoadBalancerSimulation method verifyState.
/**
* Compare the simulator's view of reality with the load balancer's. This method should
* be called after every step is performed and all threads have finished.
*/
public void verifyState() {
// verify that we consumed all messages before we do anything
for (int i = 0; i < _queues.length; ++i) {
if (_queues[i].size() > 0) {
fail("there were messages left in the queue. all messages should have been consumed during this simulation step.");
}
}
// verify that all clients have been shut down
for (Map.Entry<String, TransportClientFactory> e : _clientFactories.entrySet()) {
DoNothingClientFactory factory = (DoNothingClientFactory) e.getValue();
if (factory.getRunningClientCount() != 0) {
fail("Not all clients were shut down from factory " + e.getKey());
}
}
try {
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");
}
} catch (InterruptedException e) {
fail("unable to shutdown state in verifyState.");
}
// New load balancer with no timeout; the code below checks for services that don't
// exist,
// and a load balancer with non-zero timeout will just timeout waiting for them to be
// registered, which will never happen because the PropertyEventThread is shut down.
_loadBalancer = new SimpleLoadBalancer(_state, 0, TimeUnit.SECONDS);
// verify services are as we expect
for (String possibleService : _possibleServices) {
// about it
if (!_expectedServiceProperties.containsKey(possibleService) || !_state.isListeningToService(possibleService)) {
LoadBalancerStateItem<ServiceProperties> serviceItem = _state.getServiceProperties(possibleService);
assertTrue(serviceItem == null || serviceItem.getProperty() == null);
} else {
ServiceProperties serviceProperties = _expectedServiceProperties.get(possibleService);
ClusterProperties clusterProperties = _expectedClusterProperties.get(serviceProperties.getClusterName());
UriProperties uriProperties = _expectedUriProperties.get(serviceProperties.getClusterName());
assertEquals(_state.getServiceProperties(possibleService).getProperty(), serviceProperties);
// verify round robin'ing of the hosts for this service
for (int i = 0; i < 100; ++i) {
try {
// this call will queue up messages if we're not listening to the service, but
// it's ok, because all of the messengers have been stopped.
final TransportClient client = _loadBalancer.getClient(new URIRequest("d2://" + possibleService + random(_possiblePaths)), new RequestContext());
// if we didn't receive service unavailable, we should
// get a client back
assertNotNull(client);
} catch (ServiceUnavailableException e) {
if (uriProperties != null && clusterProperties != null) {
// only way to get here is if the prioritized
// schemes could find no available uris in the
// cluster. let's see if we can find a URI that
// matches a prioritized scheme in the cluster.
Set<String> schemes = new HashSet<String>();
for (URI uri : uriProperties.Uris()) {
schemes.add(uri.getScheme());
}
for (String scheme : clusterProperties.getPrioritizedSchemes()) {
// the code.
if (schemes.contains(scheme) && _clientFactories.containsKey(scheme)) {
break;
}
assertFalse(schemes.contains(scheme) && _clientFactories.containsKey(scheme), "why couldn't a client be found for schemes " + clusterProperties.getPrioritizedSchemes() + " with URIs: " + uriProperties.Uris());
}
}
}
}
}
}
// verify clusters are as we expect
for (String possibleCluster : _possibleClusters) {
LoadBalancerStateItem<ClusterProperties> clusterItem = _state.getClusterProperties(possibleCluster);
if (!_expectedClusterProperties.containsKey(possibleCluster) || !_state.isListeningToCluster(possibleCluster)) {
assertTrue(clusterItem == null || clusterItem.getProperty() == null, "cluster item for " + possibleCluster + " is not null: " + clusterItem);
} else {
assertNotNull(clusterItem, "Item for cluster " + possibleCluster + " should not be null, listening: " + _state.isListeningToCluster(possibleCluster) + ", keys: " + _expectedClusterProperties.keySet());
assertEquals(clusterItem.getProperty(), _expectedClusterProperties.get(possibleCluster));
}
}
// verify uris are as we expect
for (String possibleCluster : _possibleClusters) {
LoadBalancerStateItem<UriProperties> uriItem = _state.getUriProperties(possibleCluster);
if (!_expectedUriProperties.containsKey(possibleCluster) || !_state.isListeningToCluster(possibleCluster)) {
assertTrue(uriItem == null || uriItem.getProperty() == null);
} else {
assertNotNull(uriItem);
assertEquals(uriItem.getProperty(), _expectedUriProperties.get(possibleCluster));
}
}
}
Aggregations