Search in sources :

Example 1 with HostToKeyMapper

use of com.linkedin.d2.balancer.util.HostToKeyMapper in project rest.li by linkedin.

the class SimpleLoadBalancer method getPartitionInformation.

/**
 * If given a collection of keys, the method will maps keys to partitions and
 * return the servers that belongs to that partition up to limitHostPerPartition.
 *
 * If no keys are specified, the method will return hosts in all partitions
 *
 * @param serviceUri for example d2://articles
 * @param keys all the keys we want to find the partition for
 * @param limitHostPerPartition the number of hosts that we should return for this partition. Must be larger than 0.
 * @param hash this will be used to create Iterator for the hosts in the hash ring
 * @return Number of hosts in requested partitions. See {@link com.linkedin.d2.balancer.util.HostToKeyMapper} for more details.
 * @throws ServiceUnavailableException
 */
@Override
public <K> HostToKeyMapper<K> getPartitionInformation(URI serviceUri, Collection<K> keys, int limitHostPerPartition, int hash) throws ServiceUnavailableException {
    if (limitHostPerPartition <= 0) {
        throw new IllegalArgumentException("limitHostPartition cannot be 0 or less");
    }
    ServiceProperties service = listenToServiceAndCluster(serviceUri);
    String serviceName = service.getServiceName();
    String clusterName = service.getClusterName();
    ClusterProperties cluster = getClusterProperties(serviceName, clusterName);
    LoadBalancerStateItem<UriProperties> uriItem = getUriItem(serviceName, clusterName, cluster);
    UriProperties uris = uriItem.getProperty();
    List<LoadBalancerState.SchemeStrategyPair> orderedStrategies = _state.getStrategiesForService(serviceName, service.getPrioritizedSchemes());
    Map<Integer, Integer> partitionWithoutEnoughHost = new HashMap<>();
    if (!orderedStrategies.isEmpty()) {
        // get the partitionId -> keys mapping
        final PartitionAccessor accessor = getPartitionAccessor(serviceName, clusterName);
        int maxPartitionId = accessor.getMaxPartitionId();
        List<K> unmappedKeys = new ArrayList<>();
        Map<Integer, Set<K>> partitionSet = getPartitionSet(keys, accessor, unmappedKeys);
        // get the partitionId -> host URIs list
        Map<Integer, KeysAndHosts<K>> partitionDataMap = new HashMap<>();
        for (Integer partitionId : partitionSet.keySet()) {
            for (LoadBalancerState.SchemeStrategyPair pair : orderedStrategies) {
                TrackerClientSubsetItem subsetItem = getPotentialClients(serviceName, service, cluster, uris, pair.getScheme(), partitionId, uriItem.getVersion());
                Map<URI, TrackerClient> trackerClients = subsetItem.getWeightedSubset();
                int size = Math.min(trackerClients.size(), limitHostPerPartition);
                List<URI> rankedUri = new ArrayList<>(size);
                Ring<URI> ring = pair.getStrategy().getRing(uriItem.getVersion(), partitionId, trackerClients, subsetItem.shouldForceUpdate());
                Iterator<URI> iterator = ring.getIterator(hash);
                while (iterator.hasNext() && rankedUri.size() < size) {
                    URI uri = iterator.next();
                    if (!rankedUri.contains(uri)) {
                        rankedUri.add(uri);
                    }
                }
                if (rankedUri.size() < limitHostPerPartition) {
                    partitionWithoutEnoughHost.put(partitionId, limitHostPerPartition - rankedUri.size());
                }
                KeysAndHosts<K> keysAndHosts = new KeysAndHosts<>(partitionSet.get(partitionId), rankedUri);
                partitionDataMap.put(partitionId, keysAndHosts);
                if (!rankedUri.isEmpty()) {
                    // don't go to the next strategy if there are already hosts in the current one
                    break;
                }
            }
        }
        return new HostToKeyMapper<>(unmappedKeys, partitionDataMap, limitHostPerPartition, maxPartitionId + 1, partitionWithoutEnoughHost);
    } else {
        throw new ServiceUnavailableException(serviceName, "PEGA_1009. Unable to find a load balancer strategy" + "Server Schemes: [" + String.join(", ", service.getPrioritizedSchemes()) + ']');
    }
}
Also used : Set(java.util.Set) HashSet(java.util.HashSet) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) KeysAndHosts(com.linkedin.d2.balancer.util.KeysAndHosts) ServiceUnavailableException(com.linkedin.d2.balancer.ServiceUnavailableException) URI(java.net.URI) LoadBalancerState(com.linkedin.d2.balancer.LoadBalancerState) TrackerClient(com.linkedin.d2.balancer.clients.TrackerClient) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) HostToKeyMapper(com.linkedin.d2.balancer.util.HostToKeyMapper) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) PartitionAccessor(com.linkedin.d2.balancer.util.partitions.PartitionAccessor) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties)

Example 2 with HostToKeyMapper

use of com.linkedin.d2.balancer.util.HostToKeyMapper in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testStrategyFallbackInGetPartitionInformationAndRing.

/**
 * Test falling back of strategy if partition can't be found in the original one
 */
@Test
public void testStrategyFallbackInGetPartitionInformationAndRing() throws Exception {
    // setup 3 partitions. Partition 1 and Partition 2 both have server1 - server3. Partition 3 only has server1.
    // create HTTP strategy
    Map<URI, Map<Integer, PartitionData>> partitionDescriptionsPlain = new HashMap<>();
    final URI server1Plain = new URI("http://foo1.com");
    partitionDescriptionsPlain.put(server1Plain, generatePartitionData(1, 2, 3));
    LoadBalancerStrategy plainStrategy = new TestLoadBalancerStrategy(partitionDescriptionsPlain);
    // create HTTPS strategy
    Map<URI, Map<Integer, PartitionData>> partitionDescriptionsSSL = new HashMap<>();
    final URI server2Https = new URI("https://foo2.com");
    partitionDescriptionsSSL.put(server2Https, generatePartitionData(1, 2));
    final URI server3Https = new URI("https://foo3.com");
    partitionDescriptionsSSL.put(server3Https, generatePartitionData(1, 2));
    LoadBalancerStrategy SSLStrategy = new TestLoadBalancerStrategy(partitionDescriptionsSSL);
    // Prioritize HTTPS over HTTP
    List<LoadBalancerState.SchemeStrategyPair> orderedStrategies = new ArrayList<>();
    orderedStrategies.add(new LoadBalancerState.SchemeStrategyPair(PropertyKeys.HTTPS_SCHEME, SSLStrategy));
    orderedStrategies.add(new LoadBalancerState.SchemeStrategyPair(PropertyKeys.HTTP_SCHEME, plainStrategy));
    // setup the partition accessor which can only map keys from 1 - 3.
    PartitionAccessor accessor = new TestPartitionAccessor();
    HashMap<URI, Map<Integer, PartitionData>> allUris = new HashMap<>();
    allUris.putAll(partitionDescriptionsSSL);
    allUris.putAll(partitionDescriptionsPlain);
    String serviceName = "articles";
    String clusterName = "cluster";
    String path = "path";
    String strategyName = "degrader";
    URI serviceURI = new URI("d2://" + serviceName);
    SimpleLoadBalancer balancer = new SimpleLoadBalancer(new PartitionedLoadBalancerTestState(clusterName, serviceName, path, strategyName, allUris, orderedStrategies, accessor), _d2Executor);
    List<Integer> keys = Arrays.asList(1, 2, 3, 123);
    HostToKeyMapper<Integer> resultPartInfo = balancer.getPartitionInformation(serviceURI, keys, 3, 123);
    MapKeyResult<Ring<URI>, Integer> resultRing = balancer.getRings(serviceURI, keys);
    Assert.assertEquals(resultPartInfo.getLimitHostPerPartition(), 3);
    Assert.assertEquals(resultRing.getMapResult().size(), 3);
    Map<Integer, Ring<URI>> ringPerKeys = new HashMap<>();
    resultRing.getMapResult().forEach((uriRing, keysAssociated) -> keysAssociated.forEach(key -> ringPerKeys.put(key, uriRing)));
    // Important section
    // partition 1 and 2
    List<URI> ordering1 = resultPartInfo.getPartitionInfoMap().get(1).getHosts();
    Set<URI> ordering1Ring = iteratorToSet(ringPerKeys.get(1).getIterator(0));
    List<URI> ordering2 = resultPartInfo.getPartitionInfoMap().get(2).getHosts();
    Set<URI> ordering2Ring = iteratorToSet(ringPerKeys.get(2).getIterator(0));
    // partition 1 and 2. check that the HTTPS hosts are there
    // all the above variables should be the same, since all the hosts are in both partitions
    Assert.assertEqualsNoOrder(ordering1.toArray(), ordering2.toArray());
    Assert.assertEqualsNoOrder(ordering1.toArray(), ordering1Ring.toArray());
    Assert.assertEqualsNoOrder(ordering1.toArray(), ordering2Ring.toArray());
    Assert.assertEqualsNoOrder(ordering1.toArray(), Arrays.asList(server2Https, server3Https).toArray());
    // partition 3, test that is falling back to HTTP
    List<URI> ordering3 = resultPartInfo.getPartitionInfoMap().get(3).getHosts();
    Set<URI> ordering3Ring = iteratorToSet(ringPerKeys.get(3).getIterator(0));
    Assert.assertEquals(ordering3.size(), 1, "There should be just 1 http client in partition 3 (falling back from https)");
    Assert.assertEqualsNoOrder(ordering3.toArray(), ordering3Ring.toArray());
    Assert.assertEquals(ordering3.get(0), server1Plain);
}
Also used : Arrays(java.util.Arrays) DegraderTrackerClient(com.linkedin.d2.balancer.clients.DegraderTrackerClient) Test(org.testng.annotations.Test) RewriteClient(com.linkedin.d2.balancer.clients.RewriteClient) TransportCallback(com.linkedin.r2.transport.common.bridge.common.TransportCallback) RestResponse(com.linkedin.r2.message.rest.RestResponse) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) Map(java.util.Map) LoadBalancerTestState(com.linkedin.d2.balancer.LoadBalancerTestState) HashFunction(com.linkedin.d2.balancer.util.hashing.HashFunction) SynchronousExecutorService(com.linkedin.d2.discovery.event.SynchronousExecutorService) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) ServicePropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ServicePropertiesJsonSerializer) Set(java.util.Set) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) Executors(java.util.concurrent.Executors) ClusterPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ClusterPropertiesJsonSerializer) CountDownLatch(java.util.concurrent.CountDownLatch) PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) BeforeSuite(org.testng.annotations.BeforeSuite) FutureCallback(com.linkedin.common.callback.FutureCallback) PartitionAccessor(com.linkedin.d2.balancer.util.partitions.PartitionAccessor) DegraderImpl(com.linkedin.util.degrader.DegraderImpl) PropertyStoreFactory(com.linkedin.d2.balancer.simulator.SimpleLoadBalancerSimulation.PropertyStoreFactory) Callback(com.linkedin.common.callback.Callback) TransportClient(com.linkedin.r2.transport.common.bridge.client.TransportClient) ArrayList(java.util.ArrayList) URIRequest(com.linkedin.d2.balancer.util.URIRequest) Assert(org.testng.Assert) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) DarkClusterConfig(com.linkedin.d2.DarkClusterConfig) StreamSupport(java.util.stream.StreamSupport) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) IOException(java.io.IOException) FileUtils(org.apache.commons.io.FileUtils) File(java.io.File) HostToKeyMapper(com.linkedin.d2.balancer.util.HostToKeyMapper) ExecutionException(java.util.concurrent.ExecutionException) AtomicLong(java.util.concurrent.atomic.AtomicLong) UriPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.UriPropertiesJsonSerializer) RangeBasedPartitionProperties(com.linkedin.d2.balancer.properties.RangeBasedPartitionProperties) ServiceUnavailableException(com.linkedin.d2.balancer.ServiceUnavailableException) TrackerClient(com.linkedin.d2.balancer.clients.TrackerClient) LoadBalancerUtil(com.linkedin.d2.balancer.util.LoadBalancerUtil) Spliterators(java.util.Spliterators) URISyntaxException(java.net.URISyntaxException) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) ThreeRetries(com.linkedin.test.util.retry.ThreeRetries) Random(java.util.Random) PartitionedLoadBalancerTestState(com.linkedin.d2.balancer.PartitionedLoadBalancerTestState) ConsistentHashRing(com.linkedin.d2.balancer.util.hashing.ConsistentHashRing) KeysAndHosts(com.linkedin.d2.balancer.util.KeysAndHosts) Request(com.linkedin.r2.message.Request) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) SimpleLoadBalancerSimulation(com.linkedin.d2.balancer.simulator.SimpleLoadBalancerSimulation) URI(java.net.URI) PropertyStore(com.linkedin.d2.discovery.stores.PropertyStore) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) Collection(java.util.Collection) NullPartitionProperties(com.linkedin.d2.balancer.properties.NullPartitionProperties) Assert.assertNotNull(org.testng.Assert.assertNotNull) PropertyKeys(com.linkedin.d2.balancer.properties.PropertyKeys) Collectors(java.util.stream.Collectors) RewriteLoadBalancerClient(com.linkedin.d2.balancer.clients.RewriteLoadBalancerClient) List(java.util.List) DEFAULT_PARTITION_ID(com.linkedin.d2.balancer.util.partitions.DefaultPartitionAccessor.DEFAULT_PARTITION_ID) None(com.linkedin.common.util.None) LoadBalancerState(com.linkedin.d2.balancer.LoadBalancerState) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) RandomHash(com.linkedin.d2.balancer.util.hashing.RandomHash) Spliterator(java.util.Spliterator) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) PartitionAccessException(com.linkedin.d2.balancer.util.partitions.PartitionAccessException) DarkClustersConverter(com.linkedin.d2.balancer.config.DarkClustersConverter) NamedThreadFactory(com.linkedin.r2.util.NamedThreadFactory) DataProvider(org.testng.annotations.DataProvider) MD5Hash(com.linkedin.d2.balancer.util.hashing.MD5Hash) Assert.assertEquals(org.testng.Assert.assertEquals) HashMap(java.util.HashMap) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) PropertySerializer(com.linkedin.d2.discovery.PropertySerializer) HashSet(java.util.HashSet) Ring(com.linkedin.d2.balancer.util.hashing.Ring) Nonnull(javax.annotation.Nonnull) HostOverrideList(com.linkedin.d2.balancer.util.HostOverrideList) RestRequest(com.linkedin.r2.message.rest.RestRequest) Iterator(java.util.Iterator) FileSystemDirectory(com.linkedin.d2.balancer.util.FileSystemDirectory) FileStore(com.linkedin.d2.discovery.stores.file.FileStore) Assert.fail(org.testng.Assert.fail) HashBasedPartitionProperties(com.linkedin.d2.balancer.properties.HashBasedPartitionProperties) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) TimeUnit(java.util.concurrent.TimeUnit) RequestContext(com.linkedin.r2.message.RequestContext) KeyMapper(com.linkedin.d2.balancer.KeyMapper) MapKeyResult(com.linkedin.d2.balancer.util.MapKeyResult) Assert.assertTrue(org.testng.Assert.assertTrue) AfterSuite(org.testng.annotations.AfterSuite) Collections(java.util.Collections) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) URI(java.net.URI) LoadBalancerState(com.linkedin.d2.balancer.LoadBalancerState) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) PartitionedLoadBalancerTestState(com.linkedin.d2.balancer.PartitionedLoadBalancerTestState) PartitionAccessor(com.linkedin.d2.balancer.util.partitions.PartitionAccessor) ConsistentHashRing(com.linkedin.d2.balancer.util.hashing.ConsistentHashRing) Ring(com.linkedin.d2.balancer.util.hashing.Ring) Map(java.util.Map) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Example 3 with HostToKeyMapper

use of com.linkedin.d2.balancer.util.HostToKeyMapper in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testGetPartitionInfoOrdering.

/**
 * This tests the getPartitionInfo() when given a collection of keys (actually a test for KeyMapper.mapKeysV3()).
 */
@Test
public void testGetPartitionInfoOrdering() throws Exception {
    String serviceName = "articles";
    String clusterName = "cluster";
    String path = "path";
    String strategyName = "degrader";
    // setup 3 partitions. Partition 1 and Partition 2 both have server1 - server3. Partition 3 only has server1.
    Map<URI, Map<Integer, PartitionData>> partitionDescriptions = new HashMap<>();
    final URI server1 = new URI("http://foo1.com");
    Map<Integer, PartitionData> server1Data = new HashMap<>();
    server1Data.put(1, new PartitionData(1.0));
    server1Data.put(2, new PartitionData(1.0));
    server1Data.put(3, new PartitionData(1.0));
    partitionDescriptions.put(server1, server1Data);
    final URI server2 = new URI("http://foo2.com");
    Map<Integer, PartitionData> server2Data = new HashMap<>();
    server2Data.put(1, new PartitionData(1.0));
    server2Data.put(2, new PartitionData(1.0));
    partitionDescriptions.put(server2, server2Data);
    final URI server3 = new URI("http://foo3.com");
    Map<Integer, PartitionData> server3Data = new HashMap<>();
    server3Data.put(1, new PartitionData(1.0));
    server3Data.put(2, new PartitionData(1.0));
    partitionDescriptions.put(server3, server3Data);
    // setup strategy which involves tweaking the hash ring to get partitionId -> URI host
    List<LoadBalancerState.SchemeStrategyPair> orderedStrategies = new ArrayList<>();
    LoadBalancerStrategy strategy = new TestLoadBalancerStrategy(partitionDescriptions);
    orderedStrategies.add(new LoadBalancerState.SchemeStrategyPair(PropertyKeys.HTTP_SCHEME, strategy));
    // setup the partition accessor which can only map keys from 1 - 3.
    PartitionAccessor accessor = new TestPartitionAccessor();
    URI serviceURI = new URI("d2://" + serviceName);
    SimpleLoadBalancer balancer = new SimpleLoadBalancer(new PartitionedLoadBalancerTestState(clusterName, serviceName, path, strategyName, partitionDescriptions, orderedStrategies, accessor), _d2Executor);
    List<Integer> keys = new ArrayList<>();
    keys.add(1);
    keys.add(2);
    keys.add(3);
    keys.add(123);
    HostToKeyMapper<Integer> result = balancer.getPartitionInformation(serviceURI, keys, 3, 123);
    Assert.assertEquals(result.getLimitHostPerPartition(), 3);
    Assert.assertEquals(1, result.getUnmappedKeys().size());
    Assert.assertEquals(123, (int) result.getUnmappedKeys().iterator().next().getKey());
    // partition 0 should be null
    Assert.assertNull(result.getPartitionInfoMap().get(0));
    // results for partition 1 should contain server1, server2 and server3
    KeysAndHosts<Integer> keysAndHosts1 = result.getPartitionInfoMap().get(1);
    Assert.assertTrue(keysAndHosts1.getKeys().size() == 1);
    Assert.assertTrue(keysAndHosts1.getKeys().iterator().next() == 1);
    List<URI> ordering1 = keysAndHosts1.getHosts();
    // results for partition 2 should be the same as partition1.
    KeysAndHosts<Integer> keysAndHosts2 = result.getPartitionInfoMap().get(2);
    Assert.assertTrue(keysAndHosts2.getKeys().size() == 1);
    Assert.assertTrue(keysAndHosts2.getKeys().iterator().next() == 2);
    List<URI> ordering2 = keysAndHosts2.getHosts();
    // for partition 3
    KeysAndHosts<Integer> keysAndHosts3 = result.getPartitionInfoMap().get(3);
    Assert.assertTrue(keysAndHosts3.getKeys().size() == 1);
    Assert.assertTrue(keysAndHosts3.getKeys().iterator().next() == 3);
    List<URI> ordering3 = keysAndHosts3.getHosts();
    // Just compare the size and contents of the list, not the ordering.
    Assert.assertTrue(ordering1.size() == 3);
    List<URI> allServers = new ArrayList<>();
    allServers.add(server1);
    allServers.add(server2);
    allServers.add(server3);
    Assert.assertTrue(ordering1.containsAll(allServers));
    Assert.assertTrue(ordering2.containsAll(allServers));
    Assert.assertEquals(ordering1, ordering2);
    Assert.assertEquals(ordering3.get(0), server1);
    Assert.assertTrue(result.getPartitionsWithoutEnoughHosts().containsKey(3));
    Assert.assertEquals((int) result.getPartitionsWithoutEnoughHosts().get(3), 2);
}
Also used : HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) PartitionedLoadBalancerTestState(com.linkedin.d2.balancer.PartitionedLoadBalancerTestState) URI(java.net.URI) LoadBalancerState(com.linkedin.d2.balancer.LoadBalancerState) PartitionAccessor(com.linkedin.d2.balancer.util.partitions.PartitionAccessor) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) Map(java.util.Map) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Example 4 with HostToKeyMapper

use of com.linkedin.d2.balancer.util.HostToKeyMapper in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testGetAllPartitionMultipleHostsOrdering.

/**
 * This tests the getPartitionInfo() when keys are null (actually a test for KeyMapper.getAllPartitionMultipleHosts()).
 */
@Test
public void testGetAllPartitionMultipleHostsOrdering() throws Exception {
    String serviceName = "articles";
    String clusterName = "cluster";
    String path = "path";
    String strategyName = "degrader";
    // setup partition
    Map<URI, Map<Integer, PartitionData>> partitionDescriptions = new HashMap<>();
    final URI server1 = new URI("http://foo1.com");
    Map<Integer, PartitionData> server1Data = new HashMap<>();
    server1Data.put(1, new PartitionData(1.0));
    server1Data.put(2, new PartitionData(1.0));
    server1Data.put(3, new PartitionData(1.0));
    partitionDescriptions.put(server1, server1Data);
    final URI server2 = new URI("http://foo2.com");
    Map<Integer, PartitionData> server2Data = new HashMap<>();
    server2Data.put(1, new PartitionData(1.0));
    server2Data.put(2, new PartitionData(1.0));
    // server2Data.put(3, new PartitionData(1.0));
    partitionDescriptions.put(server2, server2Data);
    final URI server3 = new URI("http://foo3.com");
    Map<Integer, PartitionData> server3Data = new HashMap<>();
    server3Data.put(1, new PartitionData(1.0));
    server3Data.put(2, new PartitionData(1.0));
    // server3Data.put(3, new PartitionData(1.0));
    partitionDescriptions.put(server3, server3Data);
    // setup strategy which involves tweaking the hash ring to get partitionId -> URI host
    List<LoadBalancerState.SchemeStrategyPair> orderedStrategies = new ArrayList<>();
    LoadBalancerStrategy strategy = new TestLoadBalancerStrategy(partitionDescriptions);
    orderedStrategies.add(new LoadBalancerState.SchemeStrategyPair(PropertyKeys.HTTP_SCHEME, strategy));
    // setup the partition accessor which is used to get partitionId -> keys
    PartitionAccessor accessor = new TestPartitionAccessor();
    URI serviceURI = new URI("d2://" + serviceName);
    SimpleLoadBalancer balancer = new SimpleLoadBalancer(new PartitionedLoadBalancerTestState(clusterName, serviceName, path, strategyName, partitionDescriptions, orderedStrategies, accessor), _d2Executor);
    HostToKeyMapper<URI> result = balancer.getPartitionInformation(serviceURI, null, 3, 123);
    Assert.assertEquals(result.getPartitionInfoMap().size(), 4);
    Assert.assertEquals(4, result.getPartitionCount());
    // partition 0 should be empty
    Assert.assertTrue(result.getPartitionInfoMap().get(0).getHosts().isEmpty());
    // partition 1 should have server1, server2 and server3.
    List<URI> ordering1 = result.getPartitionInfoMap().get(1).getHosts();
    List<URI> allServers = new ArrayList<>();
    allServers.add(server1);
    allServers.add(server2);
    allServers.add(server3);
    Assert.assertTrue(ordering1.size() == 3);
    Assert.assertTrue(ordering1.containsAll(allServers));
    // partition 2 should be the same as partition 1
    List<URI> ordering2 = result.getPartitionInfoMap().get(2).getHosts();
    Assert.assertEquals(ordering1, ordering2);
    // partition 3 should only contain server1
    List<URI> ordering3 = result.getPartitionInfoMap().get(3).getHosts();
    Assert.assertEquals(ordering3.get(0), server1);
    // partition 0 and partition 3 should not have enough hosts: lacking 3 and 2 respectively.
    Assert.assertTrue(result.getPartitionsWithoutEnoughHosts().containsKey(3));
    Assert.assertTrue(result.getPartitionsWithoutEnoughHosts().containsKey(0));
    Assert.assertEquals((int) result.getPartitionsWithoutEnoughHosts().get(3), 2);
    Assert.assertEquals((int) result.getPartitionsWithoutEnoughHosts().get(0), 3);
}
Also used : HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) PartitionedLoadBalancerTestState(com.linkedin.d2.balancer.PartitionedLoadBalancerTestState) URI(java.net.URI) LoadBalancerState(com.linkedin.d2.balancer.LoadBalancerState) PartitionAccessor(com.linkedin.d2.balancer.util.partitions.PartitionAccessor) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) Map(java.util.Map) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Example 5 with HostToKeyMapper

use of com.linkedin.d2.balancer.util.HostToKeyMapper in project rest.li by linkedin.

the class ConsistentHashKeyMapperTest method verifyHostToMapperWithoutKeys.

private void verifyHostToMapperWithoutKeys(HostToKeyMapper<Integer> result) {
    Map<Integer, KeysAndHosts<Integer>> partitionInfoMap = result.getPartitionInfoMap();
    try {
        Assert.assertTrue(Arrays.asList(new URI[] { new URI("http://foo1.com"), new URI("http://foo3.com") }).containsAll(partitionInfoMap.get(0).getHosts()));
        Assert.assertTrue(Arrays.asList(new URI[] { new URI("http://foo5.com"), new URI("http://foo4.com"), new URI("http://foo6.com") }).containsAll(partitionInfoMap.get(1).getHosts()));
        Assert.assertTrue(Arrays.asList(new URI[] { new URI("http://foo2.com") }).containsAll(partitionInfoMap.get(3).getHosts()));
        Assert.assertTrue(Arrays.asList(new URI[] { new URI("http://foo2.com") }).containsAll(partitionInfoMap.get(4).getHosts()));
    } catch (URISyntaxException e) {
    }
}
Also used : KeysAndHosts(com.linkedin.d2.balancer.util.KeysAndHosts) URISyntaxException(java.net.URISyntaxException) URI(java.net.URI)

Aggregations

URI (java.net.URI)5 LoadBalancerState (com.linkedin.d2.balancer.LoadBalancerState)4 PartitionAccessor (com.linkedin.d2.balancer.util.partitions.PartitionAccessor)4 ArrayList (java.util.ArrayList)4 HashMap (java.util.HashMap)4 DarkClusterConfigMap (com.linkedin.d2.DarkClusterConfigMap)3 PartitionedLoadBalancerTestState (com.linkedin.d2.balancer.PartitionedLoadBalancerTestState)3 PartitionData (com.linkedin.d2.balancer.properties.PartitionData)3 LoadBalancerStrategy (com.linkedin.d2.balancer.strategies.LoadBalancerStrategy)3 KeysAndHosts (com.linkedin.d2.balancer.util.KeysAndHosts)3 Map (java.util.Map)3 Test (org.testng.annotations.Test)3 ServiceUnavailableException (com.linkedin.d2.balancer.ServiceUnavailableException)2 TrackerClient (com.linkedin.d2.balancer.clients.TrackerClient)2 ClusterProperties (com.linkedin.d2.balancer.properties.ClusterProperties)2 ServiceProperties (com.linkedin.d2.balancer.properties.ServiceProperties)2 UriProperties (com.linkedin.d2.balancer.properties.UriProperties)2 HostToKeyMapper (com.linkedin.d2.balancer.util.HostToKeyMapper)2 HashSet (java.util.HashSet)2 Set (java.util.Set)2