Examples with MapKeyResult com.linkedin.d2.balancer.util.MapKeyResult used on opensource projects

Example 1 with MapKeyResult

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

the class SimpleLoadBalancer method getRings.

@Override
public <K> MapKeyResult<Ring<URI>, K> getRings(URI serviceUri, Iterable<K> keys) throws ServiceUnavailableException {
    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());
    if (!orderedStrategies.isEmpty()) {
        PartitionAccessor accessor = getPartitionAccessor(serviceName, clusterName);
        // first distribute keys to partitions
        Map<Integer, Set<K>> partitionSet = new HashMap<>();
        List<MapKeyResult.UnmappedKey<K>> unmappedKeys = new ArrayList<>();
        for (final K key : keys) {
            int partitionId;
            try {
                partitionId = accessor.getPartitionId(key.toString());
            } catch (PartitionAccessException e) {
                unmappedKeys.add(new MapKeyResult.UnmappedKey<>(key, MapKeyResult.ErrorType.FAIL_TO_FIND_PARTITION));
                continue;
            }
            Set<K> set = partitionSet.computeIfAbsent(partitionId, k -> new HashSet<>());
            set.add(key);
        }
        // then we find the ring for each partition and create a map of Ring<URI> to Set<K>
        final Map<Ring<URI>, Collection<K>> ringMap = new HashMap<>(partitionSet.size() * 2);
        for (Map.Entry<Integer, Set<K>> entry : partitionSet.entrySet()) {
            int partitionId = entry.getKey();
            Ring<URI> ring = null;
            for (LoadBalancerState.SchemeStrategyPair pair : orderedStrategies) {
                TrackerClientSubsetItem subsetItem = getPotentialClients(serviceName, service, cluster, uris, pair.getScheme(), partitionId, uriItem.getVersion());
                ring = pair.getStrategy().getRing(uriItem.getVersion(), partitionId, subsetItem.getWeightedSubset(), subsetItem.shouldForceUpdate());
                if (!ring.isEmpty()) {
                    // don't fallback to the next strategy if there are already hosts in the current one
                    break;
                }
            }
            // make sure the same ring is not used in other partition
            ringMap.put(ring, entry.getValue());
        }
        return new MapKeyResult<>(ringMap, unmappedKeys);
    } else {
        throw new ServiceUnavailableException(serviceName, "PEGA_1002. Unable to find a load balancer strategy. " + "Server Schemes: [" + String.join(", ", service.getPrioritizedSchemes()) + ']');
    }
}

Example 2 with MapKeyResult

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

the class ConsistentHashKeyMapper method doMapKeys.

private <K> MapKeyResult<URI, K> doMapKeys(Ring<URI> ring, Iterable<K> keys) throws ServiceUnavailableException {
    String[] keyTokens = new String[1];
    List<MapKeyResult.UnmappedKey<K>> unmappedKeys = new ArrayList<>();
    Map<URI, Collection<K>> result = new HashMap<>();
    for (K key : keys) {
        keyTokens[0] = key.toString();
        int hashCode = _hashFunction.hash(keyTokens);
        URI uri = ring.get(hashCode);
        if (uri == null) {
            unmappedKeys.add(new MapKeyResult.UnmappedKey<>(key, MapKeyResult.ErrorType.NO_HOST_AVAILABLE_IN_PARTITION));
            continue;
        }
        Collection<K> collection = result.get(uri);
        if (collection == null) {
            collection = new ArrayList<>();
            result.put(uri, collection);
        }
        collection.add(key);
    }
    return new MapKeyResult<>(result, unmappedKeys);
}

Example 3 with MapKeyResult

use of com.linkedin.d2.balancer.util.MapKeyResult 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);
}

Example 4 with MapKeyResult

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

the class ZKFSTest method testKeyMapper.

@Test
public void testKeyMapper() throws Exception {
    final String TEST_SERVICE_NAME = "test-service";
    final String TEST_CLUSTER_NAME = "test-cluster";
    final URI TEST_SERVER_URI1 = URI.create("http://test-host-1/");
    final URI TEST_SERVER_URI2 = URI.create("http://test-host-2/");
    final int NUM_ITERATIONS = 5;
    startServer();
    try {
        ZKFSLoadBalancer balancer = getBalancer();
        FutureCallback<None> callback = new FutureCallback<>();
        balancer.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        ZKConnection conn = balancer.zkConnection();
        ZooKeeperPermanentStore<ServiceProperties> serviceStore = new ZooKeeperPermanentStore<>(conn, new ServicePropertiesJsonSerializer(), ZKFSUtil.servicePath(BASE_PATH));
        ServiceProperties props = new ServiceProperties(TEST_SERVICE_NAME, TEST_CLUSTER_NAME, "/test", Arrays.asList("degrader"), Collections.<String, Object>emptyMap(), null, null, Arrays.asList("http"), null);
        serviceStore.put(TEST_SERVICE_NAME, props);
        ClusterProperties clusterProperties = new ClusterProperties(TEST_CLUSTER_NAME);
        ZooKeeperPermanentStore<ClusterProperties> clusterStore = new ZooKeeperPermanentStore<>(conn, new ClusterPropertiesJsonSerializer(), ZKFSUtil.clusterPath(BASE_PATH));
        clusterStore.put(TEST_CLUSTER_NAME, clusterProperties);
        ZooKeeperEphemeralStore<UriProperties> uriStore = new ZooKeeperEphemeralStore<>(conn, new UriPropertiesJsonSerializer(), new UriPropertiesMerger(), ZKFSUtil.uriPath(BASE_PATH), false, true);
        Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<>();
        Map<Integer, PartitionData> partitionData = new HashMap<>(1);
        partitionData.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1.0d));
        uriData.put(TEST_SERVER_URI1, partitionData);
        uriData.put(TEST_SERVER_URI2, partitionData);
        UriProperties uriProps = new UriProperties(TEST_CLUSTER_NAME, uriData);
        callback = new FutureCallback<>();
        uriStore.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        uriStore.put(TEST_CLUSTER_NAME, uriProps);
        Set<Integer> keys = new HashSet<>();
        for (int ii = 0; ii < 100; ++ii) {
            keys.add(ii);
        }
        for (int ii = 0; ii < NUM_ITERATIONS; ++ii) {
            KeyMapper mapper = balancer.getKeyMapper();
            MapKeyResult<URI, Integer> batches = mapper.mapKeysV2(URI.create("d2://" + TEST_SERVICE_NAME), keys);
            Assert.assertEquals(batches.getMapResult().size(), 2);
            for (Map.Entry<URI, Collection<Integer>> oneBatch : batches.getMapResult().entrySet()) {
                Assert.assertTrue(oneBatch.getKey().toString().startsWith("http://test-host-"));
                Assert.assertTrue(keys.containsAll(oneBatch.getValue()));
            }
        }
    } finally {
        stopServer();
    }
}

Example 5 with MapKeyResult

use of com.linkedin.d2.balancer.util.MapKeyResult 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<>();
        Map<String, TransportClientFactory> clientFactories = new HashMap<>();
        List<String> prioritizedSchemes = new ArrayList<>();
        MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
        MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
        MockStore<UriProperties> uriRegistry = new MockStore<>();
        ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
        loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
        clientFactories.put(PropertyKeys.HTTP_SCHEME, new DoNothingClientFactory());
        SimpleLoadBalancerState state = new SimpleLoadBalancerState(executorService, uriRegistry, clusterRegistry, serviceRegistry, clientFactories, loadBalancerStrategyFactories);
        SimpleLoadBalancer loadBalancer = new SimpleLoadBalancer(state, 5, TimeUnit.SECONDS, executorService);
        FutureCallback<None> balancerCallback = new FutureCallback<>();
        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<>();
        uris.put(uri1, 1d);
        uris.put(uri2, 1d);
        uris.put(uri3, 1d);
        Map<URI, Map<Integer, PartitionData>> partitionDesc = new HashMap<>();
        Map<Integer, PartitionData> server1 = new HashMap<>();
        server1.put(0, new PartitionData(1d));
        server1.put(1, new PartitionData(1d));
        Map<Integer, PartitionData> server2 = new HashMap<>();
        server2.put(0, new PartitionData(1d));
        Map<Integer, PartitionData> server3 = new HashMap<>();
        server3.put(1, new PartitionData(1d));
        partitionDesc.put(uri1, server1);
        partitionDesc.put(uri2, server2);
        partitionDesc.put(uri3, server3);
        prioritizedSchemes.add(PropertyKeys.HTTP_SCHEME);
        int partitionMethod = tryAgain % 4;
        switch(partitionMethod) {
            case 0:
                clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<>(), new HashSet<>(), new RangeBasedPartitionProperties("id=(\\d+)", 0, 50, 2)));
                break;
            case 1:
                clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<>(), new HashSet<>(), new HashBasedPartitionProperties("id=(\\d+)", 2, HashBasedPartitionProperties.HashAlgorithm.valueOf("MODULO"))));
                break;
            case 2:
                clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<>(), new HashSet<>(), 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<>(), new HashSet<>(), 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.singletonMap(PropertyKeys.HTTP_LB_CONSISTENT_HASH_ALGORITHM, "pointBased"), 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<>(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<>();
        expectedUris.add(expectedUri1);
        expectedUris.add(expectedUri2);
        expectedUris.add(expectedUri3);
        for (int i = 0; i < 1000; ++i) {
            int ii = i % 100;
            RewriteLoadBalancerClient client = (RewriteLoadBalancerClient) 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);
            RewriteLoadBalancerClient hintedClient1 = (RewriteLoadBalancerClient) loadBalancer.getClient(new URIRequest("d2://foo/id=" + ii), requestContextWithHint);
            String hintedUri1 = hintedClient1.getUri().toString();
            Assert.assertEquals(hintedUri1, uri1.toString() + "/foo");
            RewriteLoadBalancerClient hintedClient2 = (RewriteLoadBalancerClient) 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<>();
            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!");
    }
}