Search in sources :

Example 1 with SimpleLoadBalancerState

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancerState 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;
}
Also used : ZKPersistentConnection(com.linkedin.d2.discovery.stores.zk.ZKPersistentConnection) WarmUpLoadBalancer(com.linkedin.d2.balancer.util.WarmUpLoadBalancer) SimpleLoadBalancerState(com.linkedin.d2.balancer.simple.SimpleLoadBalancerState) SimpleLoadBalancer(com.linkedin.d2.balancer.simple.SimpleLoadBalancer) LastSeenLoadBalancerWithFacilities(com.linkedin.d2.balancer.zkfs.LastSeenLoadBalancerWithFacilities) D2ClientJmxManager(com.linkedin.d2.jmx.D2ClientJmxManager) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) ZKConnectionBuilder(com.linkedin.d2.discovery.stores.zk.ZKConnectionBuilder) PropertyEventBusImpl(com.linkedin.d2.discovery.event.PropertyEventBusImpl) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) LastSeenLoadBalancerWithFacilities(com.linkedin.d2.balancer.zkfs.LastSeenLoadBalancerWithFacilities)

Example 2 with SimpleLoadBalancerState

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancerState in project rest.li by linkedin.

the class ZKFSTogglingLoadBalancerFactoryImpl method createLoadBalancer.

@Override
public TogglingLoadBalancer createLoadBalancer(ZKConnection zkConnection, ScheduledExecutorService executorService) {
    _log.info("Using d2ServicePath: " + _d2ServicePath);
    ZooKeeperPermanentStore<ClusterProperties> zkClusterRegistry = createPermanentStore(zkConnection, ZKFSUtil.clusterPath(_baseZKPath), new ClusterPropertiesJsonSerializer(), executorService, _zookeeperReadWindowMs);
    _d2ClientJmxManager.setZkClusterRegistry(zkClusterRegistry);
    ZooKeeperPermanentStore<ServiceProperties> zkServiceRegistry = createPermanentStore(zkConnection, ZKFSUtil.servicePath(_baseZKPath, _d2ServicePath), new ServicePropertiesJsonSerializer(_clientServicesConfig), executorService, _zookeeperReadWindowMs);
    _d2ClientJmxManager.setZkServiceRegistry(zkServiceRegistry);
    String backupStoreFilePath = null;
    if (_enableSaveUriDataOnDisk) {
        backupStoreFilePath = _fsd2DirPath + File.separator + "urisValues";
    }
    ZooKeeperEphemeralStore<UriProperties> zkUriRegistry = createEphemeralStore(zkConnection, ZKFSUtil.uriPath(_baseZKPath), new UriPropertiesJsonSerializer(), new UriPropertiesMerger(), _useNewEphemeralStoreWatcher, backupStoreFilePath, executorService, _zookeeperReadWindowMs);
    _d2ClientJmxManager.setZkUriRegistry(zkUriRegistry);
    FileStore<ClusterProperties> fsClusterStore = createFileStore(FileSystemDirectory.getClusterDirectory(_fsd2DirPath), new ClusterPropertiesJsonSerializer());
    _d2ClientJmxManager.setFsClusterStore(fsClusterStore);
    FileStore<ServiceProperties> fsServiceStore = createFileStore(FileSystemDirectory.getServiceDirectory(_fsd2DirPath, _d2ServicePath), new ServicePropertiesJsonSerializer());
    _d2ClientJmxManager.setFsServiceStore(fsServiceStore);
    FileStore<UriProperties> fsUriStore = createFileStore(_fsd2DirPath + File.separator + "uris", new UriPropertiesJsonSerializer());
    _d2ClientJmxManager.setFsUriStore(fsUriStore);
    PropertyEventBus<ClusterProperties> clusterBus = new PropertyEventBusImpl<>(executorService);
    PropertyEventBus<ServiceProperties> serviceBus = new PropertyEventBusImpl<>(executorService);
    PropertyEventBus<UriProperties> uriBus = new PropertyEventBusImpl<>(executorService);
    // This ensures the filesystem store receives the events from the event bus so that
    // it can keep a local backup.
    clusterBus.register(fsClusterStore);
    serviceBus.register(fsServiceStore);
    uriBus.register(fsUriStore);
    TogglingPublisher<ClusterProperties> clusterToggle = _factory.createClusterToggle(zkClusterRegistry, fsClusterStore, clusterBus);
    TogglingPublisher<ServiceProperties> serviceToggle = _factory.createServiceToggle(zkServiceRegistry, fsServiceStore, serviceBus);
    TogglingPublisher<UriProperties> uriToggle = _factory.createUriToggle(zkUriRegistry, fsUriStore, uriBus);
    SimpleLoadBalancerState state = new SimpleLoadBalancerState(executorService, uriBus, clusterBus, serviceBus, _clientFactories, _loadBalancerStrategyFactories, _sslContext, _sslParameters, _isSSLEnabled, _partitionAccessorRegistry, _sslSessionValidatorFactory, _deterministicSubsettingMetadataProvider);
    _d2ClientJmxManager.setSimpleLoadBalancerState(state);
    SimpleLoadBalancer balancer = new SimpleLoadBalancer(state, _lbTimeout, _lbTimeoutUnit, executorService);
    _d2ClientJmxManager.setSimpleLoadBalancer(balancer);
    TogglingLoadBalancer togLB = _factory.createBalancer(balancer, state, clusterToggle, serviceToggle, uriToggle);
    togLB.start(new Callback<None>() {

        @Override
        public void onError(Throwable e) {
            _log.warn("Failed to run start on the TogglingLoadBalancer, may not have registered " + "SimpleLoadBalancer and State with JMX.");
        }

        @Override
        public void onSuccess(None result) {
            _log.info("Registered SimpleLoadBalancer and State with JMX.");
        }
    });
    return togLB;
}
Also used : ServicePropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ServicePropertiesJsonSerializer) UriPropertiesMerger(com.linkedin.d2.balancer.properties.UriPropertiesMerger) TogglingLoadBalancer(com.linkedin.d2.balancer.util.TogglingLoadBalancer) UriPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.UriPropertiesJsonSerializer) ClusterPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ClusterPropertiesJsonSerializer) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) SimpleLoadBalancerState(com.linkedin.d2.balancer.simple.SimpleLoadBalancerState) SimpleLoadBalancer(com.linkedin.d2.balancer.simple.SimpleLoadBalancer) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) PropertyEventBusImpl(com.linkedin.d2.discovery.event.PropertyEventBusImpl) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) None(com.linkedin.common.util.None)

Example 3 with SimpleLoadBalancerState

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancerState in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testGetClient.

/**
 * This tests getClient(). When TargetHints and scheme does not match, throw ServiceUnavailableException
 * @throws Exception
 */
@Test(expectedExceptions = ServiceUnavailableException.class)
@SuppressWarnings("deprecation")
public void testGetClient() throws Exception {
    Map<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>> loadBalancerStrategyFactories = new HashMap<>();
    Map<String, TransportClientFactory> clientFactories = new HashMap<>();
    List<String> prioritizedSchemes = new ArrayList<>();
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    ScheduledExecutorService executorService = new SynchronousExecutorService();
    // loadBalancerStrategyFactories.put("rr", new RandomLoadBalancerStrategyFactory());
    loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
    // PrpcClientFactory();
    // new
    clientFactories.put(PropertyKeys.HTTPS_SCHEME, new DoNothingClientFactory());
    // HttpClientFactory();
    SimpleLoadBalancerState state = new SimpleLoadBalancerState(executorService, uriRegistry, clusterRegistry, serviceRegistry, clientFactories, loadBalancerStrategyFactories);
    SimpleLoadBalancer loadBalancer = new SimpleLoadBalancer(state, 5, TimeUnit.SECONDS, _d2Executor);
    FutureCallback<None> balancerCallback = new FutureCallback<>();
    loadBalancer.start(balancerCallback);
    balancerCallback.get(5, TimeUnit.SECONDS);
    Map<Integer, PartitionData> partitionData = new HashMap<>(1);
    partitionData.put(DEFAULT_PARTITION_ID, new PartitionData(1d));
    Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<>(3);
    prioritizedSchemes.add(PropertyKeys.HTTPS_SCHEME);
    clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1"));
    serviceRegistry.put("foo", new ServiceProperties("foo", "cluster-1", "/foo", Arrays.asList("degrader"), Collections.<String, Object>emptyMap(), null, null, prioritizedSchemes, null));
    uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
    URI uri = URI.create("http://test.qd.com:1234/foo");
    RequestContext requestContextWithHint = new RequestContext();
    LoadBalancerUtil.TargetHints.setRequestContextTargetService(requestContextWithHint, uri);
    URIRequest uriRequest = new URIRequest("d2://foo");
    loadBalancer.getClient(uriRequest, requestContextWithHint);
}
Also used : HashMap(java.util.HashMap) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) ArrayList(java.util.ArrayList) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) URI(java.net.URI) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) RequestContext(com.linkedin.r2.message.RequestContext) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) FutureCallback(com.linkedin.common.callback.FutureCallback) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) SynchronousExecutorService(com.linkedin.d2.discovery.event.SynchronousExecutorService) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) URIRequest(com.linkedin.d2.balancer.util.URIRequest) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) None(com.linkedin.common.util.None) Map(java.util.Map) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Example 4 with SimpleLoadBalancerState

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancerState in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testGetClientHostOverrideList.

/**
 * Tests getClient() when with host override list specified in the request context.
 * @throws Exception
 */
@Test
public void testGetClientHostOverrideList() throws Exception {
    Map<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>> loadBalancerStrategyFactories = new HashMap<>();
    Map<String, TransportClientFactory> clientFactories = new HashMap<>();
    List<String> prioritizedSchemes = new ArrayList<>();
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    ScheduledExecutorService executorService = new SynchronousExecutorService();
    loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
    clientFactories.put(PropertyKeys.HTTP_SCHEME, new DoNothingClientFactory());
    SimpleLoadBalancerState state = new SimpleLoadBalancerState(executorService, uriRegistry, clusterRegistry, serviceRegistry, clientFactories, loadBalancerStrategyFactories);
    SimpleLoadBalancer loadBalancer = new SimpleLoadBalancer(state, 5, TimeUnit.SECONDS, _d2Executor);
    FutureCallback<None> balancerCallback = new FutureCallback<>();
    loadBalancer.start(balancerCallback);
    balancerCallback.get();
    Map<Integer, PartitionData> partitionData = new HashMap<>(1);
    partitionData.put(DEFAULT_PARTITION_ID, new PartitionData(1d));
    Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<>(2);
    uriData.put(URI.create("http://host1/path"), partitionData);
    prioritizedSchemes.add(PropertyKeys.HTTP_SCHEME);
    String cluster1 = "Cluster1";
    String cluster2 = "Cluster2";
    String service1 = "service1";
    clusterRegistry.put(cluster1, new ClusterProperties(cluster1, Collections.emptyList(), Collections.emptyMap(), new HashSet<>(), NullPartitionProperties.getInstance()));
    serviceRegistry.put(service1, new ServiceProperties(service1, cluster1, "/service1Path", Arrays.asList("degrader"), Collections.<String, Object>emptyMap(), null, null, prioritizedSchemes, null));
    uriRegistry.put(cluster1, new UriProperties(cluster1, uriData));
    URI override = URI.create("http://override/path");
    URI expected = URI.create("http://override/path/service1Path");
    URIRequest uriRequest = new URIRequest("d2://service1");
    HostOverrideList clusterOverrides = new HostOverrideList();
    clusterOverrides.addClusterOverride(cluster1, override);
    RequestContext clusterContext = new RequestContext();
    clusterContext.putLocalAttr(HOST_OVERRIDE_LIST, clusterOverrides);
    Assert.assertEquals(((RewriteLoadBalancerClient) loadBalancer.getClient(uriRequest, clusterContext)).getUri(), URI.create("http://override/path/service1Path"));
    HostOverrideList serviceOverrides = new HostOverrideList();
    serviceOverrides.addServiceOverride(service1, override);
    RequestContext serviceContext = new RequestContext();
    serviceContext.putLocalAttr(HOST_OVERRIDE_LIST, serviceOverrides);
    Assert.assertEquals(((RewriteLoadBalancerClient) loadBalancer.getClient(uriRequest, serviceContext)).getUri(), URI.create("http://override/path/service1Path"));
    HostOverrideList overrides = new HostOverrideList();
    overrides.addOverride(override);
    RequestContext context = new RequestContext();
    context.putLocalAttr(HOST_OVERRIDE_LIST, overrides);
    Assert.assertEquals(((RewriteLoadBalancerClient) loadBalancer.getClient(uriRequest, context)).getUri(), URI.create("http://override/path/service1Path"));
    HostOverrideList unrelatedClusterOverrides = new HostOverrideList();
    unrelatedClusterOverrides.addClusterOverride(cluster2, override);
    RequestContext unrelatedClusterContext = new RequestContext();
    unrelatedClusterContext.putLocalAttr(HOST_OVERRIDE_LIST, unrelatedClusterOverrides);
    Assert.assertEquals(((RewriteLoadBalancerClient) loadBalancer.getClient(uriRequest, unrelatedClusterContext)).getUri(), URI.create("http://host1/path/service1Path"));
    HostOverrideList unrelatedServiceOverrides = new HostOverrideList();
    unrelatedServiceOverrides.addClusterOverride(cluster2, override);
    RequestContext unrelatedServiceContext = new RequestContext();
    unrelatedServiceContext.putLocalAttr(HOST_OVERRIDE_LIST, unrelatedServiceOverrides);
    Assert.assertEquals(((RewriteLoadBalancerClient) loadBalancer.getClient(uriRequest, unrelatedServiceContext)).getUri(), URI.create("http://host1/path/service1Path"));
}
Also used : HashMap(java.util.HashMap) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) ArrayList(java.util.ArrayList) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) URI(java.net.URI) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) HostOverrideList(com.linkedin.d2.balancer.util.HostOverrideList) RequestContext(com.linkedin.r2.message.RequestContext) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) FutureCallback(com.linkedin.common.callback.FutureCallback) HashSet(java.util.HashSet) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) SynchronousExecutorService(com.linkedin.d2.discovery.event.SynchronousExecutorService) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) URIRequest(com.linkedin.d2.balancer.util.URIRequest) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) None(com.linkedin.common.util.None) Map(java.util.Map) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Example 5 with SimpleLoadBalancerState

use of com.linkedin.d2.balancer.simple.SimpleLoadBalancerState in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testLoadBalancerDropRate.

/**
 * This test simulates dropping requests by playing with OverrideDropRate in config
 */
@Test(groups = { "small", "back-end" })
public void testLoadBalancerDropRate() throws ServiceUnavailableException, ExecutionException, InterruptedException {
    final int RETRY = 10;
    for (int tryAgain = 0; tryAgain < RETRY; ++tryAgain) {
        Map<String, LoadBalancerStrategyFactory<? extends LoadBalancerStrategy>> loadBalancerStrategyFactories = new HashMap<>();
        Map<String, TransportClientFactory> clientFactories = new HashMap<>();
        List<String> prioritizedSchemes = new ArrayList<>();
        MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
        MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
        MockStore<UriProperties> uriRegistry = new MockStore<>();
        ScheduledExecutorService executorService = new SynchronousExecutorService();
        // loadBalancerStrategyFactories.put("rr", new RandomLoadBalancerStrategyFactory());
        loadBalancerStrategyFactories.put("degrader", new DegraderLoadBalancerStrategyFactoryV3());
        // PrpcClientFactory();
        // new
        clientFactories.put(PropertyKeys.HTTP_SCHEME, new DoNothingClientFactory());
        // HttpClientFactory();
        SimpleLoadBalancerState state = new SimpleLoadBalancerState(executorService, uriRegistry, clusterRegistry, serviceRegistry, clientFactories, loadBalancerStrategyFactories);
        SimpleLoadBalancer loadBalancer = new SimpleLoadBalancer(state, 5, TimeUnit.SECONDS, _d2Executor);
        FutureCallback<None> balancerCallback = new FutureCallback<>();
        loadBalancer.start(balancerCallback);
        balancerCallback.get();
        URI uri1 = URI.create("http://test.qa1.com:1234");
        URI uri2 = URI.create("http://test.qa2.com:2345");
        URI uri3 = URI.create("http://test.qa3.com:6789");
        Map<Integer, PartitionData> partitionData = new HashMap<>(1);
        partitionData.put(DEFAULT_PARTITION_ID, new PartitionData(1d));
        Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<>(3);
        uriData.put(uri1, partitionData);
        uriData.put(uri2, partitionData);
        uriData.put(uri3, partitionData);
        prioritizedSchemes.add(PropertyKeys.HTTP_SCHEME);
        clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1"));
        serviceRegistry.put("foo", new ServiceProperties("foo", "cluster-1", "/foo", Arrays.asList("degrader"), Collections.<String, Object>emptyMap(), null, null, prioritizedSchemes, null));
        uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
        URI expectedUri1 = URI.create("http://test.qa1.com:1234/foo");
        URI expectedUri2 = URI.create("http://test.qa2.com:2345/foo");
        URI expectedUri3 = URI.create("http://test.qa3.com:6789/foo");
        Set<URI> expectedUris = new HashSet<>();
        expectedUris.add(expectedUri1);
        expectedUris.add(expectedUri2);
        expectedUris.add(expectedUri3);
        Random random = new Random();
        for (int i = 0; i < 100; ++i) {
            try {
                RewriteLoadBalancerClient client = (RewriteLoadBalancerClient) loadBalancer.getClient(new URIRequest("d2://foo/52"), new RequestContext());
                assertTrue(client.getDecoratedClient() instanceof RewriteClient);
                RewriteClient rewriteClient = (RewriteClient) client.getDecoratedClient();
                assertTrue(rewriteClient.getDecoratedClient() instanceof TrackerClient);
                DegraderTrackerClient tClient = (DegraderTrackerClient) rewriteClient.getDecoratedClient();
                DegraderImpl degrader = (DegraderImpl) tClient.getDegrader(DEFAULT_PARTITION_ID);
                DegraderImpl.Config cfg = new DegraderImpl.Config(degrader.getConfig());
                // Change DropRate to 0.0 at the rate of 1/3
                cfg.setOverrideDropRate((random.nextInt(2) == 0) ? 1.0 : 0.0);
                degrader.setConfig(cfg);
                assertTrue(expectedUris.contains(client.getUri()));
                assertEquals(client.getUri().getScheme(), PropertyKeys.HTTP_SCHEME);
            } catch (ServiceUnavailableException e) {
                assertTrue(e.toString().contains("in a bad state (high latency/high error)"));
            }
        }
        final CountDownLatch latch = new CountDownLatch(1);
        PropertyEventShutdownCallback callback = new PropertyEventShutdownCallback() {

            @Override
            public void done() {
                latch.countDown();
            }
        };
        state.shutdown(callback);
        if (!latch.await(60, TimeUnit.SECONDS)) {
            fail("unable to shutdown state");
        }
        executorService.shutdownNow();
        assertTrue(executorService.isShutdown(), "ExecutorService should have shut down!");
    }
}
Also used : PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) HashMap(java.util.HashMap) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) DarkClusterConfig(com.linkedin.d2.DarkClusterConfig) ArrayList(java.util.ArrayList) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) ServiceUnavailableException(com.linkedin.d2.balancer.ServiceUnavailableException) URI(java.net.URI) Random(java.util.Random) DegraderTrackerClient(com.linkedin.d2.balancer.clients.DegraderTrackerClient) TrackerClient(com.linkedin.d2.balancer.clients.TrackerClient) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) RequestContext(com.linkedin.r2.message.RequestContext) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) FutureCallback(com.linkedin.common.callback.FutureCallback) HashSet(java.util.HashSet) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) SynchronousExecutorService(com.linkedin.d2.discovery.event.SynchronousExecutorService) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) DegraderTrackerClient(com.linkedin.d2.balancer.clients.DegraderTrackerClient) RewriteLoadBalancerClient(com.linkedin.d2.balancer.clients.RewriteLoadBalancerClient) URIRequest(com.linkedin.d2.balancer.util.URIRequest) DegraderImpl(com.linkedin.util.degrader.DegraderImpl) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) CountDownLatch(java.util.concurrent.CountDownLatch) RewriteClient(com.linkedin.d2.balancer.clients.RewriteClient) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) None(com.linkedin.common.util.None) Map(java.util.Map) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) HashMap(java.util.HashMap) Test(org.testng.annotations.Test)

Aggregations

LoadBalancerStrategy (com.linkedin.d2.balancer.strategies.LoadBalancerStrategy)14 ClusterProperties (com.linkedin.d2.balancer.properties.ClusterProperties)12 ServiceProperties (com.linkedin.d2.balancer.properties.ServiceProperties)12 UriProperties (com.linkedin.d2.balancer.properties.UriProperties)12 LoadBalancerStrategyFactory (com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory)12 DegraderLoadBalancerStrategyFactoryV3 (com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3)12 TransportClientFactory (com.linkedin.r2.transport.common.TransportClientFactory)12 HashMap (java.util.HashMap)12 None (com.linkedin.common.util.None)11 RandomLoadBalancerStrategyFactory (com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory)11 FutureCallback (com.linkedin.common.callback.FutureCallback)10 SynchronousExecutorService (com.linkedin.d2.discovery.event.SynchronousExecutorService)9 ScheduledExecutorService (java.util.concurrent.ScheduledExecutorService)9 SimpleLoadBalancerState (com.linkedin.d2.balancer.simple.SimpleLoadBalancerState)8 ArrayList (java.util.ArrayList)8 URIRequest (com.linkedin.d2.balancer.util.URIRequest)7 MockStore (com.linkedin.d2.discovery.stores.mock.MockStore)7 RequestContext (com.linkedin.r2.message.RequestContext)7 Test (org.testng.annotations.Test)7 DarkClusterConfigMap (com.linkedin.d2.DarkClusterConfigMap)6