Examples with TransportClientFactory com.linkedin.r2.transport.common.TransportClientFactory used on opensource projects

Example 26 with TransportClientFactory

use of com.linkedin.r2.transport.common.TransportClientFactory in project rest.li by linkedin.

the class TestDisruptor method testRestTimeoutDisrupt.

@Test
public void testRestTimeoutDisrupt() throws Exception {
    final Map<String, String> properties = new HashMap<>();
    properties.put(HttpClientFactory.HTTP_REQUEST_TIMEOUT, String.valueOf(REQUEST_TIMEOUT));
    final TransportClientFactory factory = new HttpClientFactory.Builder().build();
    final TransportClient client = factory.getClient(properties);
    final RequestContext requestContext = new RequestContext();
    requestContext.putLocalAttr(DISRUPT_CONTEXT_KEY, DisruptContexts.timeout());
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicBoolean success = new AtomicBoolean(false);
    client.restRequest(new RestRequestBuilder(new URI(REQUEST_URI)).build(), requestContext, new HashMap<>(), response -> {
        success.set(response.hasError() && response.getError() instanceof TimeoutException);
        latch.countDown();
    });
    Assert.assertTrue(latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS), "Test execution timeout");
    Assert.assertTrue(success.get(), "Unexpected transport response");
}

Example 27 with TransportClientFactory

use of com.linkedin.r2.transport.common.TransportClientFactory in project rest.li by linkedin.

the class SharedZkConnectionProviderTest method testAnnouncerAndClientSharing.

/**
 * Testing sharing connections between announcers and d2client
 * @throws Exception
 */
@Test(groups = "needZk")
public void testAnnouncerAndClientSharing() throws Exception {
    // connection shared to announcers
    List<URI> hostNames = prepareHostNames(20, "testAnnouncerAndClientSharing");
    List<ZooKeeperConnectionManager> connectionManagers = prepareConnectionManagers(hostNames);
    int l = 1;
    // set up a mock transport client
    Map<String, TransportClientFactory> transportClientMap = new HashMap<>();
    TestTransportClientFactory testClientFactory = new TestTransportClientFactory();
    transportClientMap.put("http", testClientFactory);
    // connection shared to d2client
    D2Client client = getD2Client(transportClientMap);
    // there should only be one connection
    assertEquals(_provider.getZkConnectionCount(), 1);
    // start both announcers and client
    FutureCallback<None> startUpCallback = new FutureCallback<>();
    Callback<None> startUpMultiCallback = Callbacks.countDown(startUpCallback, connectionManagers.size() + 1);
    _threadPoolExecutor.submit(() -> client.start(startUpMultiCallback));
    for (ZooKeeperConnectionManager manager : connectionManagers) {
        _threadPoolExecutor.submit(() -> manager.start(startUpMultiCallback));
    }
    startUpCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    // verify zookeeper is updated
    UriProperties properties = _verificationStore.get(CLUSTER_NAME);
    assertNotNull(properties);
    assertEquals(properties.Uris().size(), 20);
    // fire some requests to make sure announcement is successful and hosts properties can be retrieved successfully.
    int requestRepeat = 1000;
    FutureCallback<None> reqCallback = new FutureCallback<>();
    fireTestRequests(client, requestRepeat, reqCallback);
    reqCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    // verify d2client received the changes
    HostSet hosts = client.getFacilities().getKeyMapper().getAllPartitionsMultipleHosts(new URI("d2://testService"), 20);
    Assert.assertEquals(hosts.getAllHosts().size(), 20);
    Assert.assertEquals(testClientFactory.requestCount.get(), 1000);
    // Markdown half of the hosts and test the results
    FutureCallback<None> hostsMarkdownCallback = new FutureCallback<>();
    Callback<None> hostsMarkdownMultiCallback = Callbacks.countDown(hostsMarkdownCallback, 10);
    for (ZooKeeperConnectionManager manager : connectionManagers.subList(0, 10)) {
        _threadPoolExecutor.submit(() -> manager.getAnnouncers()[0].markDown(hostsMarkdownMultiCallback));
    }
    hostsMarkdownCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    // verify zookeeper is updated
    properties = _verificationStore.get(CLUSTER_NAME);
    assertNotNull(properties);
    assertEquals(properties.Uris().size(), 10);
    // fire some requests to make sure announcement is successful and hosts properties can be retrieved successfully.
    FutureCallback<None> secondReqCallback = new FutureCallback<>();
    fireTestRequests(client, requestRepeat, secondReqCallback);
    secondReqCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    // verify d2client can read the zookeeper updates.
    hosts = client.getFacilities().getKeyMapper().getAllPartitionsMultipleHosts(new URI("d2://testService"), 20);
    Assert.assertEquals(hosts.getAllHosts().size(), 10);
    Assert.assertEquals(testClientFactory.requestCount.get(), 2000);
    // Mix announcements with request firing to test connection robustness.
    FutureCallback<None> thirdReqCallback = new FutureCallback<>();
    Callback<None> thirdReqMultiCallback = Callbacks.countDown(thirdReqCallback, requestRepeat + 10);
    for (int i = 0; i < requestRepeat; i++) {
        _threadPoolExecutor.submit(() -> {
            try {
                RestRequestBuilder builder = new RestRequestBuilder(new URI("d2://testService"));
                client.restRequest(builder.build(), decorateNoneCallback(thirdReqMultiCallback));
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
        if (i % 100 == 0) {
            // markup one host every 100 requests
            ZooKeeperConnectionManager manager = connectionManagers.get(i / 100);
            _threadPoolExecutor.submit(() -> {
                try {
                    manager.getAnnouncers()[0].markUp(thirdReqMultiCallback);
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            });
        }
    }
    thirdReqCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    Assert.assertEquals(testClientFactory.requestCount.get(), 3000);
    // announcers can be shutdown after announcing, without affecting client. This should not happen though.
    FutureCallback<None> announcerShutdownCallback = new FutureCallback<>();
    Callback<None> announcersShutdownCallback = Callbacks.countDown(announcerShutdownCallback, connectionManagers.size());
    for (ZooKeeperConnectionManager manager : connectionManagers) {
        manager.shutdown(announcersShutdownCallback);
    }
    announcerShutdownCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    // fire some requests to make sure d2client is still usable.
    FutureCallback<None> fourthReqCallback = new FutureCallback<>();
    fireTestRequests(client, requestRepeat, fourthReqCallback);
    thirdReqCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    hosts = client.getFacilities().getKeyMapper().getAllPartitionsMultipleHosts(new URI("d2://testService"), 20);
    Assert.assertEquals(hosts.getAllHosts().size(), 20);
    Assert.assertEquals(testClientFactory.requestCount.get(), 4000);
    // test done!
    FutureCallback<None> clientShutdownCallback = new FutureCallback<>();
    client.shutdown(clientShutdownCallback);
    clientShutdownCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    // make sure the connection is properly stopped.
    ZKPersistentConnection connection = _provider.getZKPersistentConnection(new ZKConnectionBuilder("localhost:" + ZK_PORT).setTimeout(ZK_TIMEOUT));
    Assert.assertNotNull(connection);
    Assert.assertTrue(connection.isConnectionStopped());
}

Example 28 with TransportClientFactory

use of com.linkedin.r2.transport.common.TransportClientFactory in project rest.li by linkedin.

the class ZKFSTest method getBalancer.

private ZKFSLoadBalancer getBalancer() {
    ZKFSComponentFactory f = new ZKFSComponentFactory();
    Map<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>> loadBalancerStrategyFactories = new HashMap<>();
    loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
    Map<String, TransportClientFactory> clientFactories = new HashMap<>();
    clientFactories.put("http", new HttpClientFactory.Builder().build());
    // We rely on _tmpdir below being fresh for each test case.  Otherwise, leftover files in
    // _tmpdir from a previous test could affect another test.  This is accomplished with the
    // @BeforeMethod and @AfterMethod annotations.
    ZKFSTogglingLoadBalancerFactoryImpl f2 = new ZKFSTogglingLoadBalancerFactoryImpl(f, 5, TimeUnit.SECONDS, BASE_PATH, _tmpdir.getAbsolutePath(), clientFactories, loadBalancerStrategyFactories);
    ZKFSLoadBalancer balancer = new ZKFSLoadBalancer("localhost:" + PORT, 60000, 5000, f2, null, BASE_PATH);
    return balancer;
}

Example 29 with TransportClientFactory

use of com.linkedin.r2.transport.common.TransportClientFactory in project rest.li by linkedin.

the class SharedZkConnectionProviderTest method testZKPropertyUpdate.

/**
 * Test that when there is an zookeeper property update, d2client can receive the update correctly
 */
@Test(groups = "needZK")
public void testZKPropertyUpdate() throws Exception {
    List<URI> hosts = prepareHostNames(5, "testZKPropertyUpdate");
    List<ZooKeeperConnectionManager> connectionManagers = prepareConnectionManagers(hosts);
    Map<String, TransportClientFactory> transportClientMap = new HashMap<>();
    transportClientMap.put("http", new TestTransportClientFactory());
    // connection shared to d2client
    D2Client client = getD2Client(transportClientMap);
    FutureCallback<None> startupCallback = new FutureCallback<>();
    client.start(startupCallback);
    startupCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    startConnectionManagers(connectionManagers);
    Directory d2Directory = client.getFacilities().getDirectory();
    List<String> serviceList = new ArrayList<>();
    ServiceProperties serviceProps = new ServiceProperties("newTestService", CLUSTER_NAME, "/newTestService", Arrays.asList("degrader"), Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap(), Arrays.asList("http"), Collections.emptySet());
    FutureCallback<None> propertyCallback = new FutureCallback<>();
    _serviceRegistry.put("newTestService", serviceProps, propertyCallback);
    propertyCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    FutureCallback<None> finishCallback = new FutureCallback<>();
    d2Directory.getServiceNames(new Callback<List<String>>() {

        @Override
        public void onError(Throwable e) {
            finishCallback.onError(e);
        }

        @Override
        public void onSuccess(List<String> result) {
            serviceList.addAll(result);
            finishCallback.onSuccess(None.none());
        }
    });
    finishCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    Assert.assertEquals(serviceList.size(), 2);
    Assert.assertTrue(serviceList.contains("newTestService"));
    Assert.assertTrue(serviceList.contains("testService"));
    shutdownConnectionManagers(connectionManagers);
    FutureCallback<None> clientShutdownCallback = new FutureCallback<>();
    client.shutdown(clientShutdownCallback);
    clientShutdownCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
}

Example 30 with TransportClientFactory

use of com.linkedin.r2.transport.common.TransportClientFactory in project rest.li by linkedin.

the class SimpleLoadBalancerTest method setupLoadBalancer.

private SimpleLoadBalancer setupLoadBalancer(LoadBalancerState state, MockStore<ServiceProperties> serviceRegistry, MockStore<ClusterProperties> clusterRegistry, MockStore<UriProperties> uriRegistry) throws ExecutionException, InterruptedException {
    Map<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>> loadBalancerStrategyFactories = new HashMap<>();
    Map<String, TransportClientFactory> clientFactories = new HashMap<>();
    LoadBalancerState loadBalancerState = state;
    loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
    clientFactories.put(PropertyKeys.HTTP_SCHEME, new DoNothingClientFactory());
    clientFactories.put(PropertyKeys.HTTPS_SCHEME, new DoNothingClientFactory());
    if (loadBalancerState == null) {
        loadBalancerState = new SimpleLoadBalancerState(new SynchronousExecutorService(), uriRegistry, clusterRegistry, serviceRegistry, clientFactories, loadBalancerStrategyFactories);
    }
    SimpleLoadBalancer loadBalancer = new SimpleLoadBalancer(loadBalancerState, 5, TimeUnit.SECONDS, _d2Executor);
    FutureCallback<None> balancerCallback = new FutureCallback<>();
    loadBalancer.start(balancerCallback);
    balancerCallback.get();
    return loadBalancer;
}