Search in sources :

Example 26 with PartitionData

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

the class SimpleLoadBalancerStateTest method testListValueInTransportClientProperties.

@Test(groups = { "small", "back-end" })
public void testListValueInTransportClientProperties() throws URISyntaxException {
    reset();
    URI uri = URI.create("http://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);
    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>();
    List<String> allowedClientOverrideKeys = new ArrayList<String>();
    allowedClientOverrideKeys.add(PropertyKeys.HTTP_REQUEST_TIMEOUT);
    allowedClientOverrideKeys.add(PropertyKeys.HTTP_RESPONSE_COMPRESSION_OPERATIONS);
    transportClientProperties.put(PropertyKeys.ALLOWED_CLIENT_OVERRIDE_KEYS, allowedClientOverrideKeys);
    List<String> compressionOperations = new ArrayList<String>();
    compressionOperations.add("get");
    compressionOperations.add("batch_get");
    compressionOperations.add("get_all");
    transportClientProperties.put(PropertyKeys.HTTP_RESPONSE_COMPRESSION_OPERATIONS, compressionOperations);
    transportClientProperties = Collections.unmodifiableMap(transportClientProperties);
    _serviceRegistry.put("service-1", new ServiceProperties("service-1", "cluster-1", "/test", Arrays.asList("random"), Collections.<String, Object>emptyMap(), transportClientProperties, null, schemes, null));
    assertNull(_state.getClient("service-1", uri));
    _uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
    TrackerClient client = _state.getClient("service-1", uri);
    assertNotNull(client);
    assertEquals(client.getUri(), uri);
}
Also used : HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) URI(java.net.URI) NullStateListenerCallback(com.linkedin.d2.balancer.LoadBalancerState.NullStateListenerCallback) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) TrackerClient(com.linkedin.d2.balancer.clients.TrackerClient) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) Map(java.util.Map) HashMap(java.util.HashMap) Test(org.testng.annotations.Test) DegraderLoadBalancerTest(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerTest)

Example 27 with PartitionData

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

the class ZooKeeperServerTest method testZkServer.

@Test(groups = { "small", "back-end" })
public void testZkServer() throws InterruptedException, URISyntaxException, IOException, PropertyStoreException, ExecutionException {
    URI uri1 = URI.create("http://cluster-1/test");
    URI uri2 = URI.create("http://cluster-1-again/test");
    ZKConnection zkClient = new ZKConnection("localhost:" + PORT, 5000);
    zkClient.start();
    ZooKeeperEphemeralStore<UriProperties> store = new ZooKeeperEphemeralStore<UriProperties>(zkClient, new UriPropertiesJsonSerializer(), new UriPropertiesMerger(), "/echo/lb/uris");
    FutureCallback<None> callback = new FutureCallback<None>();
    store.start(callback);
    callback.get();
    ZooKeeperServer server = new ZooKeeperServer(store);
    final String cluster = "cluster-1";
    assertNull(store.get(cluster));
    assertNull(store.get("cluster-2"));
    // bring up uri1
    markUp(server, cluster, uri1, 0.5d);
    UriProperties properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.getPartitionDataMap(uri1).get(DefaultPartitionAccessor.DEFAULT_PARTITION_ID).getWeight(), 0.5d);
    assertEquals(properties.Uris().size(), 1);
    // test mark up when already up call
    markUp(server, cluster, uri1, 2d);
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.getPartitionDataMap(uri1).get(DefaultPartitionAccessor.DEFAULT_PARTITION_ID).getWeight(), 2d);
    assertEquals(properties.Uris().size(), 1);
    // bring up uri 2
    markUp(server, cluster, uri2, 1.5d);
    properties = store.get(cluster);
    assertEquals(properties.getPartitionDataMap(uri1).get(DefaultPartitionAccessor.DEFAULT_PARTITION_ID).getWeight(), 2d);
    assertEquals(properties.getPartitionDataMap(uri2).get(DefaultPartitionAccessor.DEFAULT_PARTITION_ID).getWeight(), 1.5d);
    assertEquals(properties.Uris().size(), 2);
    // bring down uri 1
    markDown(server, cluster, uri1);
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.getPartitionDataMap(uri2).get(DefaultPartitionAccessor.DEFAULT_PARTITION_ID).getWeight(), 1.5d);
    assertEquals(properties.Uris().size(), 1);
    // test bring down when already down
    markDown(server, cluster, uri1);
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.getPartitionDataMap(uri2).get(DefaultPartitionAccessor.DEFAULT_PARTITION_ID).getWeight(), 1.5d);
    assertEquals(properties.Uris().size(), 1);
    // bring down uri 2
    markDown(server, cluster, uri2);
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.Uris().size(), 0);
    // test bad cluster doesn't exist
    markDown(server, "BAD CLUSTER", uri1);
    properties = store.get("BAD CLUSTER");
    assertNull(properties);
    // bring up uri1
    Map<Integer, PartitionData> partitionWeight = new HashMap<Integer, PartitionData>();
    partitionWeight.put(5, new PartitionData(0.3d));
    partitionWeight.put(15, new PartitionData(0.7d));
    markUp(server, cluster, uri1, partitionWeight, null);
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.getPartitionDataMap(uri1), partitionWeight);
    Map<String, Object> uri2SpecificProperties = new HashMap<String, Object>();
    uri2SpecificProperties.put("foo", "fooValue");
    uri2SpecificProperties.put("bar", 1);
    partitionWeight.put(10, new PartitionData(1d));
    // bring up uri2 with uri specific properties
    markUp(server, cluster, uri2, partitionWeight, uri2SpecificProperties);
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.Uris().size(), 2);
    assertEquals(properties.getPartitionDataMap(uri2), partitionWeight);
    assertNotNull(properties.getUriSpecificProperties());
    assertEquals(properties.getUriSpecificProperties().size(), 1);
    assertEquals(properties.getUriSpecificProperties().get(uri2), uri2SpecificProperties);
    // bring down uri1 and bring it back up again with properties
    markDown(server, cluster, uri1);
    Map<String, Object> uri1SpecificProperties = new HashMap<String, Object>();
    uri1SpecificProperties.put("baz", "bazValue");
    // use new partition data so that we can test the mapping later on
    Map<Integer, PartitionData> newUri1PartitionWeights = new HashMap<Integer, PartitionData>(partitionWeight);
    newUri1PartitionWeights.remove(10);
    markUp(server, cluster, uri1, newUri1PartitionWeights, uri1SpecificProperties);
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.Uris().size(), 2);
    assertEquals(properties.getPartitionDataMap(uri1), newUri1PartitionWeights);
    assertEquals(properties.getPartitionDataMap(uri2), partitionWeight);
    assertNotNull(properties.getUriSpecificProperties());
    assertEquals(properties.getUriSpecificProperties().size(), 2);
    assertEquals(properties.getUriSpecificProperties().get(uri1), uri1SpecificProperties);
    assertEquals(properties.getUriSpecificProperties().get(uri2), uri2SpecificProperties);
    Set<URI> uriSet = new HashSet<URI>();
    uriSet.add(uri1);
    uriSet.add(uri2);
    assertEquals(properties.getUriBySchemeAndPartition("http", 5), uriSet);
    uriSet.remove(uri1);
    assertEquals(properties.getUriBySchemeAndPartition("http", 10), uriSet);
}
Also used : ZKConnection(com.linkedin.d2.discovery.stores.zk.ZKConnection) HashMap(java.util.HashMap) UriPropertiesMerger(com.linkedin.d2.balancer.properties.UriPropertiesMerger) ZooKeeperEphemeralStore(com.linkedin.d2.discovery.stores.zk.ZooKeeperEphemeralStore) URI(java.net.URI) UriPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.UriPropertiesJsonSerializer) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) None(com.linkedin.common.util.None) FutureCallback(com.linkedin.common.callback.FutureCallback) HashSet(java.util.HashSet) Test(org.testng.annotations.Test)

Example 28 with PartitionData

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

the class ZooKeeperServerTest method markUp.

private void markUp(ZooKeeperServer server, String cluster, URI uri, double weight) {
    Map<Integer, PartitionData> partitionWeight = new HashMap<Integer, PartitionData>();
    partitionWeight.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(weight));
    markUp(server, cluster, uri, partitionWeight, null);
}
Also used : HashMap(java.util.HashMap) PartitionData(com.linkedin.d2.balancer.properties.PartitionData)

Example 29 with PartitionData

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

the class DegraderLoadBalancerTest method getDefaultPartitionData.

public static Map<Integer, PartitionData> getDefaultPartitionData(double weight, int numberOfPartitions) {
    PartitionData data = new PartitionData(weight);
    Map<Integer, PartitionData> partitionDataMap = new HashMap<Integer, PartitionData>(numberOfPartitions + 1);
    for (int p = 0; p < numberOfPartitions; ++p) partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID + p, data);
    return partitionDataMap;
}
Also used : PartitionData(com.linkedin.d2.balancer.properties.PartitionData) HashMap(java.util.HashMap)

Example 30 with PartitionData

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

the class LoadBalancerEchoServer method markUp.

public void markUp(Map<Integer, Double> partitionWeight) throws PropertyStoreException {
    FutureCallback<None> callback = new FutureCallback<None>();
    Map<Integer, PartitionData> partitionDataMap = new HashMap<Integer, PartitionData>();
    if (partitionWeight != null) {
        for (int partitionId : partitionWeight.keySet()) {
            partitionDataMap.put(partitionId, new PartitionData(partitionWeight.get(partitionId)));
        }
    } else {
        partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1d));
    }
    _announcer.markUp(_cluster, _uri, partitionDataMap, callback);
    try {
        callback.get(10, TimeUnit.SECONDS);
    } catch (Exception e) {
        throw new PropertyStoreException(e);
    }
}
Also used : PropertyStoreException(com.linkedin.d2.discovery.stores.PropertyStoreException) HashMap(java.util.HashMap) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) None(com.linkedin.common.util.None) FutureCallback(com.linkedin.common.callback.FutureCallback) URISyntaxException(java.net.URISyntaxException) TimeoutException(java.util.concurrent.TimeoutException) RestException(com.linkedin.r2.message.rest.RestException) PropertyStoreException(com.linkedin.d2.discovery.stores.PropertyStoreException) IOException(java.io.IOException) ExecutionException(java.util.concurrent.ExecutionException)

Aggregations

PartitionData (com.linkedin.d2.balancer.properties.PartitionData)55 HashMap (java.util.HashMap)53 Test (org.testng.annotations.Test)36 URI (java.net.URI)35 Map (java.util.Map)31 UriProperties (com.linkedin.d2.balancer.properties.UriProperties)28 ArrayList (java.util.ArrayList)25 TrackerClient (com.linkedin.d2.balancer.clients.TrackerClient)21 ServiceProperties (com.linkedin.d2.balancer.properties.ServiceProperties)20 None (com.linkedin.common.util.None)15 FutureCallback (com.linkedin.common.callback.FutureCallback)14 ClusterProperties (com.linkedin.d2.balancer.properties.ClusterProperties)13 DegraderLoadBalancerTest (com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerTest)13 NullStateListenerCallback (com.linkedin.d2.balancer.LoadBalancerState.NullStateListenerCallback)12 RequestContext (com.linkedin.r2.message.RequestContext)11 LoadBalancerStrategy (com.linkedin.d2.balancer.strategies.LoadBalancerStrategy)9 LoadBalancerState (com.linkedin.d2.balancer.LoadBalancerState)8 PartitionAccessor (com.linkedin.d2.balancer.util.partitions.PartitionAccessor)7 TransportClientFactory (com.linkedin.r2.transport.common.TransportClientFactory)7 HashSet (java.util.HashSet)7