Examples with ZKConnectionBuilder com.linkedin.d2.discovery.stores.zk.ZKConnectionBuilder used on opensource projects

Example 1 with ZKConnectionBuilder

use of com.linkedin.d2.discovery.stores.zk.ZKConnectionBuilder in project rest.li by linkedin.

the class LastSeenBalancerWithFacilitiesFactory method create.

@Override
public LoadBalancerWithFacilities create(D2ClientConfig config) {
    LOG.info("Creating D2 LoadBalancer based on LastSeenLoadBalancerWithFacilities");
    D2ClientJmxManager d2ClientJmxManager = new D2ClientJmxManager(config.d2JmxManagerPrefix, config.jmxManager);
    // init connection
    ZKConnectionBuilder zkConnectionBuilder = new ZKConnectionBuilder(config.zkHosts);
    zkConnectionBuilder.setShutdownAsynchronously(config.shutdownAsynchronously).setIsSymlinkAware(config.isSymlinkAware).setTimeout((int) config.zkSessionTimeoutInMs);
    ZKPersistentConnection zkPersistentConnection;
    if (config.zkConnectionToUseForLB != null) {
        LOG.info("LastSeenLoadBalancer using shared connection to zookeeper");
        zkPersistentConnection = config.zkConnectionToUseForLB;
    } else {
        LOG.info("LastSeenLoadBalancer using its own connection to zookeeper");
        zkPersistentConnection = new ZKPersistentConnection(zkConnectionBuilder);
    }
    // init all the stores
    LastSeenZKStore<ClusterProperties> lsClusterStore = getClusterPropertiesLastSeenZKStore(config, zkPersistentConnection, d2ClientJmxManager, config._executorService, config.zookeeperReadWindowMs);
    PropertyEventBus<ClusterProperties> clusterBus = new PropertyEventBusImpl<>(config._executorService);
    clusterBus.setPublisher(lsClusterStore);
    LastSeenZKStore<ServiceProperties> lsServiceStore = getServicePropertiesLastSeenZKStore(config, zkPersistentConnection, d2ClientJmxManager, config._executorService, config.zookeeperReadWindowMs);
    PropertyEventBus<ServiceProperties> serviceBus = new PropertyEventBusImpl<>(config._executorService);
    serviceBus.setPublisher(lsServiceStore);
    LastSeenZKStore<UriProperties> lsUrisStore = getUriPropertiesLastSeenZKStore(config, zkPersistentConnection, d2ClientJmxManager, config._executorService, config.zookeeperReadWindowMs);
    PropertyEventBus<UriProperties> uriBus = new PropertyEventBusImpl<>(config._executorService);
    uriBus.setPublisher(lsUrisStore);
    // create the simple load balancer
    SimpleLoadBalancerState state = new SimpleLoadBalancerState(config._executorService, uriBus, clusterBus, serviceBus, config.clientFactories, config.loadBalancerStrategyFactories, config.sslContext, config.sslParameters, config.isSSLEnabled, config.partitionAccessorRegistry, config.sslSessionValidatorFactory, config.deterministicSubsettingMetadataProvider);
    d2ClientJmxManager.setSimpleLoadBalancerState(state);
    SimpleLoadBalancer simpleLoadBalancer = new SimpleLoadBalancer(state, config.lbWaitTimeout, config.lbWaitUnit, config._executorService);
    d2ClientJmxManager.setSimpleLoadBalancer(simpleLoadBalancer);
    // add facilities
    LastSeenLoadBalancerWithFacilities lastSeenLoadBalancer = new LastSeenLoadBalancerWithFacilities(simpleLoadBalancer, config.basePath, config.d2ServicePath, zkPersistentConnection, lsClusterStore, lsServiceStore, lsUrisStore);
    LoadBalancerWithFacilities balancer = lastSeenLoadBalancer;
    if (config.warmUp) {
        balancer = new WarmUpLoadBalancer(balancer, lastSeenLoadBalancer, config.startUpExecutorService, config.fsBasePath, config.d2ServicePath, config.downstreamServicesFetcher, config.warmUpTimeoutSeconds, config.warmUpConcurrentRequests);
    }
    return balancer;
}

Example 2 with ZKConnectionBuilder

use of com.linkedin.d2.discovery.stores.zk.ZKConnectionBuilder in project rest.li by linkedin.

the class SharedZkConnectionProviderTest method getD2Client.

/**
 * Obtain the d2client with the same setup.
 */
private D2Client getD2Client(Map<String, TransportClientFactory> transportClientFactoryMap) {
    ZKConnectionBuilder connectionBuilder = new ZKConnectionBuilder("localhost:" + ZK_PORT);
    connectionBuilder.setTimeout(ZK_TIMEOUT);
    ZKPersistentConnection zkConnectionToUse = _provider.getZKPersistentConnection(connectionBuilder);
    D2ClientBuilder d2ClientBuilder = new D2ClientBuilder();
    d2ClientBuilder.setZkHosts("localhost:" + ZK_PORT).setZkSessionTimeout(ZK_TIMEOUT, TimeUnit.MILLISECONDS).setZKConnectionForloadBalancer(zkConnectionToUse).setLoadBalancerWithFacilitiesFactory(new LastSeenBalancerWithFacilitiesFactory()).setClientFactories(transportClientFactoryMap);
    return d2ClientBuilder.build();
}

Example 3 with ZKConnectionBuilder

use of com.linkedin.d2.discovery.stores.zk.ZKConnectionBuilder in project rest.li by linkedin.

the class SharedZkConnectionProviderTest method setUp.

@BeforeMethod
public void setUp() throws Exception {
    _provider = new SharedZkConnectionProvider();
    try {
        _zkServer = new ZKServer(ZK_PORT);
        _zkServer.startup();
    } catch (IOException e) {
        fail("unable to instantiate real zk server on port " + ZK_PORT);
    }
    ZKConnection serviceZkConn = new ZKConnectionBuilder("localhost:" + ZK_PORT).setTimeout(5000).setWaitForConnected(true).build();
    ZKConnection clusterZkConn = new ZKConnectionBuilder("localhost:" + ZK_PORT).setTimeout(5000).setWaitForConnected(true).build();
    _serviceRegistry = new ZooKeeperPermanentStore<>(serviceZkConn, new ServicePropertiesJsonSerializer(), ZKFSUtil.servicePath(ZKBASE_PATH));
    _clusterRegistry = new ZooKeeperPermanentStore<>(clusterZkConn, new ClusterPropertiesJsonSerializer(), ZKFSUtil.clusterPath(ZKBASE_PATH));
    FutureCallback<None> storesStartupCallBack = new FutureCallback<>();
    Callback<None> multiStartupCallback = Callbacks.countDown(storesStartupCallBack, 2);
    serviceZkConn.start();
    clusterZkConn.start();
    _serviceRegistry.start(multiStartupCallback);
    _clusterRegistry.start(multiStartupCallback);
    storesStartupCallBack.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    FutureCallback<None> propertiesSetupCallback = new FutureCallback<>();
    Callback<None> multiPropertiesCallback = Callbacks.countDown(propertiesSetupCallback, 2);
    ServiceProperties serviceProps = new ServiceProperties(SERVICE_NAME, CLUSTER_NAME, "/testService", Arrays.asList("degrader"), Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap(), Arrays.asList("http"), Collections.emptySet());
    _serviceRegistry.put(SERVICE_NAME, serviceProps, multiPropertiesCallback);
    ClusterProperties clusterProps = new ClusterProperties(CLUSTER_NAME);
    _clusterRegistry.put(CLUSTER_NAME, clusterProps, multiPropertiesCallback);
    propertiesSetupCallback.get(BLOCKING_CALL_TIMEOUT, TimeUnit.MILLISECONDS);
    _verificationStore = createAndStartVerificationStore();
    _threadPoolExecutor = Executors.newFixedThreadPool(10);
}

Example 4 with ZKConnectionBuilder

use of com.linkedin.d2.discovery.stores.zk.ZKConnectionBuilder in project rest.li by linkedin.

the class SharedZkConnectionProviderTest method prepareConnectionManagers.

/**
 * For each given uri, generate a zookeeperConnectionManager for announcement
 */
private List<ZooKeeperConnectionManager> prepareConnectionManagers(List<URI> hostNames) throws Exception {
    List<ZooKeeperConnectionManager> connectionManagers = new ArrayList<>();
    for (URI uri : hostNames) {
        ZooKeeperServer server = new ZooKeeperServer();
        ZooKeeperAnnouncer announcer = new ZooKeeperAnnouncer(server, true);
        announcer.setCluster(CLUSTER_NAME);
        announcer.setUri(uri.toString());
        Map<Integer, PartitionData> partitionWeight = new HashMap<>();
        partitionWeight.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(0.5d));
        announcer.setPartitionData(partitionWeight);
        ZooKeeperConnectionManager.ZKStoreFactory<UriProperties, ZooKeeperEphemeralStore<UriProperties>> factory = new ZooKeeperUriStoreFactory();
        ZKConnectionBuilder connectionBuilder = new ZKConnectionBuilder("localhost:" + ZK_PORT);
        connectionBuilder.setTimeout(ZK_TIMEOUT);
        ZKPersistentConnection connection = _provider.getZKPersistentConnection(connectionBuilder);
        ZooKeeperConnectionManager connectionManager = new ZooKeeperConnectionManager(connection, ZKBASE_PATH, factory, announcer);
        connectionManagers.add(connectionManager);
    }
    return connectionManagers;
}

Example 5 with ZKConnectionBuilder

use of com.linkedin.d2.discovery.stores.zk.ZKConnectionBuilder 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());
}