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

Example 1 with ZKPersistentConnection

use of com.linkedin.d2.discovery.stores.zk.ZKPersistentConnection 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 ZKPersistentConnection

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

the class ZookeeperConnectionManagerTest method testMarkUpDuringSessionExpirationManyCallbacks.

@Test(invocationCount = 10, timeOut = 10000, retryAnalyzer = ThreeRetries.class)
public void testMarkUpDuringSessionExpirationManyCallbacks() throws Exception {
    ZooKeeperAnnouncer announcer = getZooKeeperAnnouncer(_cluster, _uri, WEIGHT);
    ZKPersistentConnection zkPersistentConnection = getZkPersistentConnection();
    ZooKeeperConnectionManager manager = createManager(true, zkPersistentConnection, announcer);
    // set up many concurrent callbacks
    FutureCallback<None> allMarkupsSucceed = new FutureCallback<>();
    int count = 1000;
    Callback<None> markUpAllServersCallback = new MultiCallback(allMarkupsSucceed, 2 * count);
    ExecutorService executorService = Executors.newScheduledThreadPool(100);
    for (int i = 0; i < count; i++) {
        executorService.execute(() -> {
            manager.markDownAllServers(new IgnoreCancelledCallback(markUpAllServersCallback));
            manager.markUpAllServers(new IgnoreCancelledCallback(markUpAllServersCallback));
        });
    }
    // expiring the connection
    long oldSessionId = zkPersistentConnection.getZooKeeper().getSessionId();
    ZKTestUtil.expireSession("localhost:" + PORT, zkPersistentConnection.getZooKeeper(), 10, TimeUnit.SECONDS);
    ZKTestUtil.waitForNewSessionEstablished(oldSessionId, zkPersistentConnection, 10, TimeUnit.SECONDS);
    try {
        allMarkupsSucceed.get(1, TimeUnit.MILLISECONDS);
        Assert.fail("All the callbacks were resolved before expiring the connection, which means it won't test that callbacks are invoked even after session expiration");
    } catch (Throwable e) {
    // expected
    }
    allMarkupsSucceed.get();
    // making sure that a new connection has been established. There should be no need to wait, because at least one markup should have been run on
    // the new connection, which means that by this part of code it should already have been established
    ZKTestUtil.waitForNewSessionEstablished(oldSessionId, zkPersistentConnection, 0, TimeUnit.SECONDS);
    // data validation
    dataValidation(_uri, _cluster, WEIGHT);
    shutdownManager(manager);
    executorService.shutdown();
}

Example 3 with ZKPersistentConnection

use of com.linkedin.d2.discovery.stores.zk.ZKPersistentConnection 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 4 with ZKPersistentConnection

use of com.linkedin.d2.discovery.stores.zk.ZKPersistentConnection 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 ZKPersistentConnection

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