Examples with LoadBalancer com.linkedin.d2.balancer.LoadBalancer used on opensource projects

Example 6 with LoadBalancer

use of com.linkedin.d2.balancer.LoadBalancer 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);
}

Example 7 with LoadBalancer

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

the class SimpleLoadBalancerTest method testClusterInfoProviderUnregisterClusterListener.

@Test
public void testClusterInfoProviderUnregisterClusterListener() throws InterruptedException, ExecutionException, ServiceUnavailableException {
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    SimpleLoadBalancer loadBalancer = setupLoadBalancer(null, serviceRegistry, clusterRegistry, uriRegistry);
    FutureCallback<None> balancerCallback = new FutureCallback<>();
    loadBalancer.start(balancerCallback);
    balancerCallback.get();
    MockClusterListener testClusterListener = new MockClusterListener();
    loadBalancer.registerClusterListener(testClusterListener);
    loadBalancer.listenToCluster(CLUSTER1_NAME, false, new LoadBalancerState.NullStateListenerCallback());
    clusterRegistry.put(CLUSTER1_NAME, new ClusterProperties(CLUSTER1_NAME, Collections.emptyList(), Collections.emptyMap(), Collections.emptySet(), NullPartitionProperties.getInstance(), Collections.emptyList(), new HashMap<>(), false));
    Assert.assertEquals(testClusterListener.getClusterAddedCount(CLUSTER1_NAME), 1, "expected add count of 1");
    Assert.assertEquals(testClusterListener.getClusterRemovedCount(CLUSTER1_NAME), 0, "expected remove count of 0");
    // now unregister, and we don't expect the counts to change.
    loadBalancer.unregisterClusterListener(testClusterListener);
    clusterRegistry.put(CLUSTER1_NAME, new ClusterProperties(CLUSTER1_NAME, Collections.emptyList(), Collections.emptyMap(), Collections.emptySet(), NullPartitionProperties.getInstance(), Collections.emptyList(), new HashMap<>(), false));
    Assert.assertEquals(testClusterListener.getClusterAddedCount(CLUSTER1_NAME), 1, "expected unchanged add count of 1 because unregistered ");
    Assert.assertEquals(testClusterListener.getClusterRemovedCount(CLUSTER1_NAME), 0, "expected unchanged remove count of 0 because unregistered");
}

Example 8 with LoadBalancer

use of com.linkedin.d2.balancer.LoadBalancer 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"));
}

Example 9 with LoadBalancer

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

the class SimpleLoadBalancerTest method testClusterInfoProviderGetDarkClusters.

@Test
public void testClusterInfoProviderGetDarkClusters() throws InterruptedException, ExecutionException, ServiceUnavailableException {
    int numHttp = 3;
    int numHttps = 4;
    int partitionIdForAdd = 0;
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    SimpleLoadBalancer loadBalancer = setupLoadBalancer(null, serviceRegistry, clusterRegistry, uriRegistry);
    DarkClusterConfig darkClusterConfig = new DarkClusterConfig().setMultiplier(1.0f).setDispatcherOutboundTargetRate(1).setDispatcherMaxRequestsToBuffer(1).setDispatcherBufferedRequestExpiryInSeconds(1);
    DarkClusterConfigMap darkClusterConfigMap = new DarkClusterConfigMap();
    darkClusterConfigMap.put(DARK_CLUSTER1_NAME, darkClusterConfig);
    clusterRegistry.put(CLUSTER1_NAME, new ClusterProperties(CLUSTER1_NAME, Collections.emptyList(), Collections.emptyMap(), Collections.emptySet(), NullPartitionProperties.getInstance(), Collections.emptyList(), DarkClustersConverter.toProperties(darkClusterConfigMap), false));
    populateUriRegistry(numHttp, numHttps, partitionIdForAdd, uriRegistry);
    loadBalancer.getDarkClusterConfigMap(CLUSTER1_NAME, new Callback<DarkClusterConfigMap>() {

        @Override
        public void onError(Throwable e) {
            Assert.fail("getDarkClusterConfigMap threw exception", e);
        }

        @Override
        public void onSuccess(DarkClusterConfigMap returnedDarkClusterConfigMap) {
            Assert.assertEquals(returnedDarkClusterConfigMap, darkClusterConfigMap, "dark cluster configs should be equal");
            Assert.assertEquals(returnedDarkClusterConfigMap.get(DARK_CLUSTER1_NAME).getMultiplier(), 1.0f, "multiplier should match");
        }
    });
}

Example 10 with LoadBalancer

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

the class SimpleLoadBalancerTest method testClusterInfoProviderRegisterClusterListener.

/**
 * The Register cluster Listener code is already tested in SimpleLoadBalancerStateTest, this is here for testing the
 * SimpleLoadBalancer API exposing this.
 */
@Test
public void testClusterInfoProviderRegisterClusterListener() throws InterruptedException, ExecutionException, ServiceUnavailableException {
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    SimpleLoadBalancer loadBalancer = setupLoadBalancer(null, serviceRegistry, clusterRegistry, uriRegistry);
    FutureCallback<None> balancerCallback = new FutureCallback<>();
    loadBalancer.start(balancerCallback);
    balancerCallback.get();
    MockClusterListener testClusterListener = new MockClusterListener();
    loadBalancer.registerClusterListener(testClusterListener);
    loadBalancer.listenToCluster(CLUSTER1_NAME, false, new LoadBalancerState.NullStateListenerCallback());
    clusterRegistry.put(CLUSTER1_NAME, new ClusterProperties(CLUSTER1_NAME, Collections.emptyList(), Collections.emptyMap(), Collections.emptySet(), NullPartitionProperties.getInstance(), Collections.emptyList(), new HashMap<>(), false));
    Assert.assertEquals(testClusterListener.getClusterAddedCount(CLUSTER1_NAME), 1, "expected add count of 1");
    Assert.assertEquals(testClusterListener.getClusterRemovedCount(CLUSTER1_NAME), 0, "expected remove count of 0");
    final CountDownLatch latch = new CountDownLatch(1);
    PropertyEventShutdownCallback callback = latch::countDown;
    loadBalancer.shutdown(callback);
    if (!latch.await(60, TimeUnit.SECONDS)) {
        fail("unable to shutdown state");
    }
    Assert.assertEquals(testClusterListener.getClusterAddedCount(CLUSTER1_NAME), 1, "expected add count of 1 after shutdown");
    Assert.assertEquals(testClusterListener.getClusterRemovedCount(CLUSTER1_NAME), 1, "expected remove count of 1 after shutdown");
}