Search in sources :

Example 11 with FutureCallback

use of com.linkedin.common.callback.FutureCallback in project rest.li by linkedin.

the class ZookeeperConnectionManagerTest method testMarkUpAndMarkDown.

@Test
public void testMarkUpAndMarkDown() throws IOException, ExecutionException, InterruptedException, PropertyStoreException {
    final String uri = "http://cluster-2/test";
    final String cluster = "cluster-2";
    final double weight = 0.5d;
    ZooKeeperAnnouncer announcer = new ZooKeeperAnnouncer(new ZooKeeperServer());
    announcer.setCluster(cluster);
    announcer.setUri(uri);
    Map<Integer, PartitionData> partitionWeight = new HashMap<Integer, PartitionData>();
    partitionWeight.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(weight));
    announcer.setPartitionData(partitionWeight);
    ZooKeeperConnectionManager manager = createManager(announcer);
    FutureCallback<None> managerStartCallback = new FutureCallback<None>();
    manager.start(managerStartCallback);
    managerStartCallback.get();
    ZooKeeperEphemeralStore<UriProperties> store = createAndStartUriStore();
    UriProperties properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.getPartitionDataMap(URI.create(uri)).get(DefaultPartitionAccessor.DEFAULT_PARTITION_ID).getWeight(), weight);
    assertEquals(properties.Uris().size(), 1);
    FutureCallback<None> markDownCallback = new FutureCallback<None>();
    announcer.markDown(markDownCallback);
    markDownCallback.get();
    properties = store.get(cluster);
    assertNotNull(properties);
    assertEquals(properties.Uris().size(), 0);
}
Also used : HashMap(java.util.HashMap) 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) Test(org.testng.annotations.Test)

Example 12 with FutureCallback

use of com.linkedin.common.callback.FutureCallback 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 13 with FutureCallback

use of com.linkedin.common.callback.FutureCallback in project rest.li by linkedin.

the class LoadBalancerEchoServer method markDown.

public void markDown() throws PropertyStoreException {
    FutureCallback<None> callback = new FutureCallback<None>();
    _announcer.markDown(_cluster, _uri, callback);
    try {
        callback.get(10, TimeUnit.SECONDS);
    } catch (Exception e) {
        throw new PropertyStoreException(e);
    }
}
Also used : PropertyStoreException(com.linkedin.d2.discovery.stores.PropertyStoreException) 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)

Example 14 with FutureCallback

use of com.linkedin.common.callback.FutureCallback 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)

Example 15 with FutureCallback

use of com.linkedin.common.callback.FutureCallback 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<None>();
        balancer.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        ZKConnection conn = balancer.zkConnection();
        ZooKeeperPermanentStore<ServiceProperties> serviceStore = new ZooKeeperPermanentStore<ServiceProperties>(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<ClusterProperties>(conn, new ClusterPropertiesJsonSerializer(), ZKFSUtil.clusterPath(BASE_PATH));
        clusterStore.put(TEST_CLUSTER_NAME, clusterProperties);
        ZooKeeperEphemeralStore<UriProperties> uriStore = new ZooKeeperEphemeralStore<UriProperties>(conn, new UriPropertiesJsonSerializer(), new UriPropertiesMerger(), ZKFSUtil.uriPath(BASE_PATH), false, true);
        Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<URI, Map<Integer, PartitionData>>();
        Map<Integer, PartitionData> partitionData = new HashMap<Integer, PartitionData>(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<None>();
        uriStore.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        uriStore.put(TEST_CLUSTER_NAME, uriProps);
        Set<Integer> keys = new HashSet<Integer>();
        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();
    }
}
Also used : ServicePropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ServicePropertiesJsonSerializer) HashMap(java.util.HashMap) KeyMapper(com.linkedin.d2.balancer.KeyMapper) ZooKeeperPermanentStore(com.linkedin.d2.discovery.stores.zk.ZooKeeperPermanentStore) UriPropertiesMerger(com.linkedin.d2.balancer.properties.UriPropertiesMerger) URI(java.net.URI) UriPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.UriPropertiesJsonSerializer) ClusterPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ClusterPropertiesJsonSerializer) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) FutureCallback(com.linkedin.common.callback.FutureCallback) HashSet(java.util.HashSet) ZKConnection(com.linkedin.d2.discovery.stores.zk.ZKConnection) ZooKeeperEphemeralStore(com.linkedin.d2.discovery.stores.zk.ZooKeeperEphemeralStore) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) Collection(java.util.Collection) None(com.linkedin.common.util.None) Map(java.util.Map) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Aggregations

FutureCallback (com.linkedin.common.callback.FutureCallback)152 Test (org.testng.annotations.Test)116 None (com.linkedin.common.util.None)77 RestRequest (com.linkedin.r2.message.rest.RestRequest)50 ExecutionException (java.util.concurrent.ExecutionException)50 RequestContext (com.linkedin.r2.message.RequestContext)47 RestResponse (com.linkedin.r2.message.rest.RestResponse)43 RestRequestBuilder (com.linkedin.r2.message.rest.RestRequestBuilder)36 HashMap (java.util.HashMap)26 CountDownLatch (java.util.concurrent.CountDownLatch)25 TransportCallbackAdapter (com.linkedin.r2.transport.common.bridge.client.TransportCallbackAdapter)22 ByteString (com.linkedin.data.ByteString)21 RemoteInvocationException (com.linkedin.r2.RemoteInvocationException)21 URI (java.net.URI)21 TimeoutException (java.util.concurrent.TimeoutException)21 StreamResponse (com.linkedin.r2.message.stream.StreamResponse)20 UriProperties (com.linkedin.d2.balancer.properties.UriProperties)16 ArrayList (java.util.ArrayList)16 RestException (com.linkedin.r2.message.rest.RestException)15 AsyncSharedPoolImpl (com.linkedin.r2.transport.http.client.AsyncSharedPoolImpl)14