Search in sources :

Example 11 with SimpleLoadBalancer

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testLoadBalancerDropRate.

/**
   * This test simulates dropping requests by playing with OverrideDropRate in config
   *
   */
@Test(groups = { "small", "back-end" })
public void testLoadBalancerDropRate() throws ServiceUnavailableException, ExecutionException, InterruptedException {
    final int RETRY = 10;
    for (int tryAgain = 0; tryAgain < RETRY; ++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 = new SynchronousExecutorService();
        //loadBalancerStrategyFactories.put("rr", new RandomLoadBalancerStrategyFactory());
        loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
        // PrpcClientFactory();
        // new
        clientFactories.put("http", new DoNothingClientFactory());
        // 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", Arrays.asList("degrader"), 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);
        Random random = new Random();
        for (int i = 0; i < 100; ++i) {
            try {
                RewriteClient client = (RewriteClient) loadBalancer.getClient(new URIRequest("d2://foo/52"), new RequestContext());
                TrackerClient tClient = (TrackerClient) client.getWrappedClient();
                DegraderImpl degrader = (DegraderImpl) tClient.getDegrader(DefaultPartitionAccessor.DEFAULT_PARTITION_ID);
                DegraderImpl.Config cfg = new DegraderImpl.Config(degrader.getConfig());
                // Change DropRate to 0.0 at the rate of 1/3
                cfg.setOverrideDropRate((random.nextInt(2) == 0) ? 1.0 : 0.0);
                degrader.setConfig(cfg);
                assertTrue(expectedUris.contains(client.getUri()));
                assertEquals(client.getUri().getScheme(), "http");
            } catch (ServiceUnavailableException e) {
                assertTrue(e.toString().contains("in a bad state (high latency/high error)"));
            }
        }
        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!");
    }
}
Also used : PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) HashMap(java.util.HashMap) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) ArrayList(java.util.ArrayList) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) ServiceUnavailableException(com.linkedin.d2.balancer.ServiceUnavailableException) URI(java.net.URI) Random(java.util.Random) TrackerClient(com.linkedin.d2.balancer.clients.TrackerClient) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) RequestContext(com.linkedin.r2.message.RequestContext) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) FutureCallback(com.linkedin.common.callback.FutureCallback) HashSet(java.util.HashSet) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) SynchronousExecutorService(com.linkedin.d2.discovery.event.SynchronousExecutorService) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) URIRequest(com.linkedin.d2.balancer.util.URIRequest) DegraderImpl(com.linkedin.util.degrader.DegraderImpl) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) CountDownLatch(java.util.concurrent.CountDownLatch) RewriteClient(com.linkedin.d2.balancer.clients.RewriteClient) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) None(com.linkedin.common.util.None) Map(java.util.Map) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Example 12 with SimpleLoadBalancer

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.

the class ConsistentHashKeyMapperTest method getConsistentHashKeyMapper.

private ConsistentHashKeyMapper getConsistentHashKeyMapper(RingFactory<URI> ringFactory) throws URISyntaxException {
    String serviceName = "articles";
    String clusterName = "cluster";
    String path = "path";
    String strategyName = "degrader";
    //setup partition
    Map<URI, Map<Integer, PartitionData>> partitionDescriptions = new HashMap<URI, Map<Integer, PartitionData>>();
    final URI foo1 = new URI("http://foo1.com");
    Map<Integer, PartitionData> foo1Data = new HashMap<Integer, PartitionData>();
    foo1Data.put(0, new PartitionData(1.0));
    partitionDescriptions.put(foo1, foo1Data);
    final URI foo2 = new URI("http://foo2.com");
    Map<Integer, PartitionData> foo2Data = new HashMap<Integer, PartitionData>();
    foo2Data.put(3, new PartitionData(1.0));
    foo2Data.put(4, new PartitionData(1.0));
    partitionDescriptions.put(foo2, foo2Data);
    final URI foo3 = new URI("http://foo3.com");
    Map<Integer, PartitionData> foo3Data = new HashMap<Integer, PartitionData>();
    foo3Data.put(0, new PartitionData(1.0));
    partitionDescriptions.put(foo3, foo3Data);
    final URI foo4 = new URI("http://foo4.com");
    Map<Integer, PartitionData> foo4Data = new HashMap<Integer, PartitionData>();
    foo4Data.put(1, new PartitionData(1.0));
    partitionDescriptions.put(foo4, foo4Data);
    final URI foo5 = new URI("http://foo5.com");
    Map<Integer, PartitionData> foo5Data = new HashMap<Integer, PartitionData>();
    foo5Data.put(1, new PartitionData(1.0));
    partitionDescriptions.put(foo5, foo5Data);
    final URI foo6 = new URI("http://foo6.com");
    Map<Integer, PartitionData> foo6Data = new HashMap<Integer, PartitionData>();
    foo6Data.put(1, new PartitionData(1.0));
    partitionDescriptions.put(foo6, foo6Data);
    //setup strategy which involves tweaking the hash ring to get partitionId -> URI host
    List<LoadBalancerState.SchemeStrategyPair> orderedStrategies = new ArrayList<LoadBalancerState.SchemeStrategyPair>();
    LoadBalancerStrategy strategy = new TestLoadBalancerStrategy(partitionDescriptions, ringFactory);
    orderedStrategies.add(new LoadBalancerState.SchemeStrategyPair("http", 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));
    ConsistentHashKeyMapper mapper = new ConsistentHashKeyMapper(balancer, balancer);
    return mapper;
}
Also used : SimpleLoadBalancer(com.linkedin.d2.balancer.simple.SimpleLoadBalancer) 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) HashMap(java.util.HashMap) Map(java.util.Map)

Example 13 with SimpleLoadBalancer

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancer in project rest.li by linkedin.

the class ConsistentHashKeyMapperTest method testMapKeysConcurrency.

@SuppressWarnings("rawtypes")
@Test
public void testMapKeysConcurrency() throws Exception {
    String serviceName = "articles";
    String clusterName = "cluster";
    String path = "path";
    String strategyName = "degrader";
    int numPartitions = 500;
    // setup partition
    Map<URI, Map<Integer, PartitionData>> partitionDescriptions = new HashMap<URI, Map<Integer, PartitionData>>();
    final URI foo1 = new URI("http://foo1.com");
    Map<Integer, PartitionData> foo1Data = new HashMap<Integer, PartitionData>();
    for (int i = 0; i < numPartitions; i++) {
        foo1Data.put(i, new PartitionData(1.0));
    }
    partitionDescriptions.put(foo1, foo1Data);
    DegraderLoadBalancerStrategyV3 strategy = new DegraderLoadBalancerStrategyV3(new DegraderLoadBalancerStrategyConfig(5000), serviceName, null);
    List<LoadBalancerState.SchemeStrategyPair> orderedStrategies = new ArrayList<LoadBalancerState.SchemeStrategyPair>();
    orderedStrategies.add(new LoadBalancerState.SchemeStrategyPair("http", strategy));
    PartitionAccessor accessor = new TestDeadlockPartitionAccessor(numPartitions);
    SimpleLoadBalancer balancer = new SimpleLoadBalancer(new PartitionedLoadBalancerTestState(clusterName, serviceName, path, strategyName, partitionDescriptions, orderedStrategies, accessor));
    ConsistentHashKeyMapper mapper = new ConsistentHashKeyMapper(balancer, balancer);
    CountDownLatch latch = new CountDownLatch(numPartitions);
    List<Runnable> runnables = createRunnables(numPartitions, mapper, serviceName, latch);
    final ExecutorService executor = Executors.newFixedThreadPool(numPartitions);
    List<Future> futures = new ArrayList<Future>();
    for (int i = 0; i < numPartitions; i++) {
        futures.add(executor.submit(runnables.get(i)));
    }
    for (Future future : futures) {
        future.get(30, TimeUnit.SECONDS);
    }
}
Also used : HashMap(java.util.HashMap) DegraderLoadBalancerStrategyConfig(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyConfig) ArrayList(java.util.ArrayList) URI(java.net.URI) LoadBalancerState(com.linkedin.d2.balancer.LoadBalancerState) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) SimpleLoadBalancer(com.linkedin.d2.balancer.simple.SimpleLoadBalancer) PartitionedLoadBalancerTestState(com.linkedin.d2.balancer.PartitionedLoadBalancerTestState) CountDownLatch(java.util.concurrent.CountDownLatch) DegraderLoadBalancerStrategyV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyV3) PartitionAccessor(com.linkedin.d2.balancer.util.partitions.PartitionAccessor) ExecutorService(java.util.concurrent.ExecutorService) Future(java.util.concurrent.Future) HashMap(java.util.HashMap) Map(java.util.Map) Test(org.testng.annotations.Test)

Example 14 with SimpleLoadBalancer

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);
}
Also used : HashMap(java.util.HashMap) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) TransportClientAdapter(com.linkedin.r2.transport.common.bridge.client.TransportClientAdapter) RequestContext(com.linkedin.r2.message.RequestContext) TransportClient(com.linkedin.r2.transport.common.bridge.client.TransportClient) Client(com.linkedin.r2.transport.common.Client) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) HttpClientFactory(com.linkedin.r2.transport.http.client.HttpClientFactory) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) TransportClient(com.linkedin.r2.transport.common.bridge.client.TransportClient) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) URIRequest(com.linkedin.d2.balancer.util.URIRequest) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties)

Example 15 with SimpleLoadBalancer

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!");
    }
}
Also used : HashMap(java.util.HashMap) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) ArrayList(java.util.ArrayList) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) RangeBasedPartitionProperties(com.linkedin.d2.balancer.properties.RangeBasedPartitionProperties) RequestContext(com.linkedin.r2.message.RequestContext) FutureCallback(com.linkedin.common.callback.FutureCallback) HashSet(java.util.HashSet) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) HashBasedPartitionProperties(com.linkedin.d2.balancer.properties.HashBasedPartitionProperties) URIRequest(com.linkedin.d2.balancer.util.URIRequest) RewriteClient(com.linkedin.d2.balancer.clients.RewriteClient) Collection(java.util.Collection) Map(java.util.Map) HashMap(java.util.HashMap) PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) ServiceUnavailableException(com.linkedin.d2.balancer.ServiceUnavailableException) URI(java.net.URI) ConsistentHashRing(com.linkedin.d2.balancer.util.hashing.ConsistentHashRing) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) CountDownLatch(java.util.concurrent.CountDownLatch) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ConsistentHashRing(com.linkedin.d2.balancer.util.hashing.ConsistentHashRing) Ring(com.linkedin.d2.balancer.util.hashing.Ring) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) MD5Hash(com.linkedin.d2.balancer.util.hashing.MD5Hash) None(com.linkedin.common.util.None) Test(org.testng.annotations.Test)

Aggregations

SimpleLoadBalancer (com.linkedin.d2.balancer.simple.SimpleLoadBalancer)17 URI (java.net.URI)14 HashMap (java.util.HashMap)14 Test (org.testng.annotations.Test)14 LoadBalancerStrategy (com.linkedin.d2.balancer.strategies.LoadBalancerStrategy)12 ArrayList (java.util.ArrayList)11 ClusterProperties (com.linkedin.d2.balancer.properties.ClusterProperties)10 ServiceProperties (com.linkedin.d2.balancer.properties.ServiceProperties)10 UriProperties (com.linkedin.d2.balancer.properties.UriProperties)10 Map (java.util.Map)10 TransportClientFactory (com.linkedin.r2.transport.common.TransportClientFactory)9 PartitionData (com.linkedin.d2.balancer.properties.PartitionData)8 LoadBalancerStrategyFactory (com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory)8 LoadBalancerState (com.linkedin.d2.balancer.LoadBalancerState)7 DegraderLoadBalancerStrategyFactoryV3 (com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3)7 RequestContext (com.linkedin.r2.message.RequestContext)7 None (com.linkedin.common.util.None)6 PartitionedLoadBalancerTestState (com.linkedin.d2.balancer.PartitionedLoadBalancerTestState)6 URIRequest (com.linkedin.d2.balancer.util.URIRequest)6 PartitionAccessor (com.linkedin.d2.balancer.util.partitions.PartitionAccessor)6