Examples with ZooKeeperConnectionManager com.linkedin.d2.balancer.servers.ZooKeeperConnectionManager used on opensource projects

Example 11 with ZooKeeperConnectionManager

use of com.linkedin.d2.balancer.servers.ZooKeeperConnectionManager in project rest.li by linkedin.

the class SharedZkConnectionProviderTest method startConnectionManagers.

private void startConnectionManagers(List<ZooKeeperConnectionManager> managers) throws Exception {
    FutureCallback<None> markupCallback = new FutureCallback<>();
    Callback<None> markupMultiCallback = Callbacks.countDown(markupCallback, managers.size());
    for (ZooKeeperConnectionManager manager : managers) {
        _threadPoolExecutor.submit(() -> {
            try {
                // Using random sleep to introduce delay to simulate uncertainty during real environment.
                Thread.sleep(Math.abs(new Random().nextInt()) % 100);
                manager.start(markupMultiCallback);
            } catch (Exception e) {
                markupMultiCallback.onError(new RuntimeException("Announcing failed for host: " + manager.getAnnouncers()[0].getUri() + " due to: " + e.getMessage(), e));
            }
        });
    }
    markupCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
}

Example 12 with ZooKeeperConnectionManager

use of com.linkedin.d2.balancer.servers.ZooKeeperConnectionManager 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 13 with ZooKeeperConnectionManager

use of com.linkedin.d2.balancer.servers.ZooKeeperConnectionManager 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 14 with ZooKeeperConnectionManager

use of com.linkedin.d2.balancer.servers.ZooKeeperConnectionManager in project rest.li by linkedin.

the class ExampleD2Server method main.

public static void main(String[] args) throws IOException, ParseException {
    // get server configuration
    JSONObject json = parseConfig();
    System.out.println("Finished parsing the server config");
    List<Map<String, Object>> d2ServersConfigs = (List<Map<String, Object>>) json.get("d2Servers");
    List<Map<String, Object>> echoServerConfigs = (List<Map<String, Object>>) json.get("echoServers");
    Long timeout = (Long) json.get("announcerStartTimeout");
    Long shutdownTimeout = (Long) json.get("announcerShutdownTimeout");
    // create and start echo servers
    List<EchoServer> echoServers = createAndStartEchoServers(echoServerConfigs);
    System.out.println("Finished creating echo servers");
    // create d2 announcers (for announcing the existence of these servers to the world)
    ZooKeeperAnnouncer[] zkAnnouncers = createZkAnnouncers(d2ServersConfigs);
    // manager will keep track of the lifecycle of d2 announcers i.e. start publishing
    // once connected to zookeeper, reconnect if zookeeper is disconnected, etc
    ZooKeeperConnectionManager manager = createZkManager(json, zkAnnouncers);
    System.out.println("Starting zookeeper announcers");
    ExecutorService executorService = Executors.newSingleThreadExecutor();
    startAnnouncers(manager, executorService, timeout);
    // waiting for user to turn off this process
    System.in.read();
    System.out.println("shutting down...");
    shutdown(echoServers, manager, executorService, shutdownTimeout);
}

Example 15 with ZooKeeperConnectionManager

use of com.linkedin.d2.balancer.servers.ZooKeeperConnectionManager in project rest.li by linkedin.

the class ExampleD2Server method createZkManager.

private static ZooKeeperConnectionManager createZkManager(JSONObject json, ZooKeeperAnnouncer[] zkAnnouncers) {
    ZKUriStoreFactory factory = new ZKUriStoreFactory();
    String zkConnectString = (String) json.get("zkConnectString");
    int zkRetryLimit = ((Long) json.get("zkRetryLimit")).intValue();
    int zkSessionTimeout = ((Long) json.get("zkSessionTimeout")).intValue();
    String zkBasePath = (String) json.get("zkBasePath");
    return new ZooKeeperConnectionManager(zkConnectString, zkSessionTimeout, zkBasePath, factory, zkRetryLimit, zkAnnouncers);
}