Examples with PartitionData com.linkedin.d2.balancer.properties.PartitionData used on opensource projects

Example 11 with PartitionData

use of com.linkedin.d2.balancer.properties.PartitionData in project rest.li by linkedin.

the class SimpleLoadBalancerStateTest method testSSLDisabledWithHttpsInstances.

@Test(groups = { "small", "back-end" })
public void testSSLDisabledWithHttpsInstances() throws URISyntaxException {
    reset();
    URI uri = URI.create("http://cluster-1/test");
    URI httpsUri = URI.create("https://cluster-1/test");
    List<String> schemes = new ArrayList<String>();
    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>>();
    uriData.put(uri, partitionData);
    uriData.put(httpsUri, partitionData);
    schemes.add("https");
    schemes.add("http");
    assertNull(_state.getClient("service-1", uri));
    // set up state
    _state.listenToCluster("cluster-1", new NullStateListenerCallback());
    assertNull(_state.getClient("service-1", uri));
    _state.listenToService("service-1", new NullStateListenerCallback());
    assertNull(_state.getClient("service-1", uri));
    Map<String, Object> transportClientProperties = new HashMap<String, Object>();
    transportClientProperties.put(HttpClientFactory.HTTP_SSL_CONTEXT, _sslContext);
    transportClientProperties.put(HttpClientFactory.HTTP_SSL_PARAMS, _sslParameters);
    transportClientProperties = Collections.unmodifiableMap(transportClientProperties);
    ServiceProperties serviceProperties = new ServiceProperties("service-1", "cluster-1", "/test", Arrays.asList("random"), Collections.<String, Object>emptyMap(), transportClientProperties, null, schemes, null);
    _serviceRegistry.put("service-1", serviceProperties);
    assertNull(_state.getClient("service-1", uri));
    _uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
    List<LoadBalancerState.SchemeStrategyPair> schemeStrategyPairs = _state.getStrategiesForService("service-1", schemes);
    for (LoadBalancerState.SchemeStrategyPair pair : schemeStrategyPairs) {
        assertFalse("https".equalsIgnoreCase(pair.getScheme()), "https shouldn't be in any schemeStrategyPair");
    }
    TrackerClient client = _state.getClient("service-1", uri);
    assertNotNull(client);
    assertEquals(client.getUri(), uri);
    // Unfortunately, I don't know a good way to intercept the logs
    client = _state.getClient("service-1", httpsUri);
    assertNull(client, "shouldn't pick an https uri");
    _state.refreshTransportClientsPerService(serviceProperties);
}

Example 12 with PartitionData

use of com.linkedin.d2.balancer.properties.PartitionData in project rest.li by linkedin.

the class SimpleLoadBalancerStateTest method testShutdown.

@Test(groups = { "small", "back-end" })
public void testShutdown() throws URISyntaxException, InterruptedException {
    reset();
    URI uri = URI.create("http://cluster-1/test");
    TestListener listener = new TestListener();
    List<String> schemes = new ArrayList<String>();
    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>>();
    uriData.put(uri, partitionData);
    schemes.add("http");
    _state.register(listener);
    assertNull(listener.scheme);
    assertNull(listener.strategy);
    assertNull(listener.serviceName);
    // set up state
    _state.listenToCluster("cluster-1", new NullStateListenerCallback());
    _state.listenToService("service-1", new NullStateListenerCallback());
    _clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", schemes));
    _uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
    _serviceRegistry.put("service-1", new ServiceProperties("service-1", "cluster-1", "/test", Arrays.asList("random")));
    TrackerClient client = _state.getClient("cluster-1", uri);
    TestShutdownCallback callback = new TestShutdownCallback();
    _state.shutdown(callback);
    if (!callback.await(10, TimeUnit.SECONDS)) {
        fail("unable to shut down state");
    }
    for (TransportClientFactory factory : _clientFactories.values()) {
        SimpleLoadBalancerTest.DoNothingClientFactory f = (SimpleLoadBalancerTest.DoNothingClientFactory) factory;
        assertEquals(f.getRunningClientCount(), 0, "Not all clients were shut down");
    }
}

Example 13 with PartitionData

use of com.linkedin.d2.balancer.properties.PartitionData in project rest.li by linkedin.

the class SimpleLoadBalancerStateTest method testShutdownWithListener.

@Test(groups = { "small", "back-end" })
public void testShutdownWithListener() throws URISyntaxException, InterruptedException {
    reset();
    URI uri = URI.create("http://cluster-1/test");
    TestListener listener = new TestListener();
    List<String> schemes = new ArrayList<String>();
    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>>();
    uriData.put(uri, partitionData);
    schemes.add("http");
    _state.register(listener);
    assertNull(listener.scheme);
    assertNull(listener.strategy);
    assertNull(listener.serviceName);
    // set up state
    _state.listenToCluster("cluster-1", new NullStateListenerCallback());
    _state.listenToService("service-1", new NullStateListenerCallback());
    List<String> strategyList = Arrays.asList("degraderV3");
    _state.refreshServiceStrategies(new ServiceProperties("service-1", "cluster-1", "/test", strategyList, Collections.<String, Object>emptyMap(), Collections.<String, Object>emptyMap(), Collections.<String, String>emptyMap(), schemes, Collections.<URI>emptySet()));
    assertEquals(listener.scheme, "http");
    assertNotNull(listener.strategy);
    assertEquals(listener.serviceName, "service-1");
    TrackerClient client = _state.getClient("cluster-1", uri);
    TestShutdownCallback callback = new TestShutdownCallback();
    _state.shutdown(callback);
    if (!callback.await(10, TimeUnit.SECONDS)) {
        fail("unable to shut down state");
    }
    for (TransportClientFactory factory : _clientFactories.values()) {
        SimpleLoadBalancerTest.DoNothingClientFactory f = (SimpleLoadBalancerTest.DoNothingClientFactory) factory;
        assertEquals(f.getRunningClientCount(), 0, "Not all clients were shut down");
    }
    assertNull(listener.scheme);
    assertNull(listener.strategy);
    assertNull(listener.serviceName);
}

Example 14 with PartitionData

use of com.linkedin.d2.balancer.properties.PartitionData 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<URI, Map<Integer, PartitionData>>();
    final URI server1 = new URI("http://foo1.com");
    Map<Integer, PartitionData> server1Data = new HashMap<Integer, PartitionData>();
    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<Integer, PartitionData>();
    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<Integer, PartitionData>();
    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<LoadBalancerState.SchemeStrategyPair>();
    LoadBalancerStrategy strategy = new TestLoadBalancerStrategy(partitionDescriptions);
    orderedStrategies.add(new LoadBalancerState.SchemeStrategyPair("http", 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));
    List<Integer> keys = new ArrayList<Integer>();
    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);
}

Example 15 with PartitionData

use of com.linkedin.d2.balancer.properties.PartitionData in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testLoadBalancerSmoke.

@Test(groups = { "small", "back-end" })
public void testLoadBalancerSmoke() throws URISyntaxException, ServiceUnavailableException, InterruptedException, ExecutionException {
    for (int tryAgain = 0; tryAgain < 1000; ++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);
        for (int i = 0; i < 100; ++i) {
            RewriteClient client = (RewriteClient) loadBalancer.getClient(new URIRequest("d2://foo/52"), new RequestContext());
            assertTrue(expectedUris.contains(client.getUri()));
            assertEquals(client.getUri().getScheme(), "http");
        }
        final CountDownLatch latch = new CountDownLatch(1);
        PropertyEventShutdownCallback callback = new PropertyEventShutdownCallback() {

            @Override
            public void done() {
                latch.countDown();
            }
        };
        state.shutdown(callback);
        if (!latch.await(60, TimeUnit.SECONDS)) {
            fail("unable to shutdown state");
        }
        executorService.shutdownNow();
        assertTrue(executorService.isShutdown(), "ExecutorService should have shut down!");
    }
}