Search in sources :

Example 16 with URIRequest

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

the class DegraderLoadBalancerTest method clusterRecovery1TC.

/**
   * helper method to test DegraderLoadBalancerStrategy recovery with 1 TrackerClient.
   *
   * We want to test DegraderV2 and V3 with 2 different strategies : LoadBalacing and Call Dropping.
   * So this method needs to able to handle all 4 permutations.
   *
   * @param myMap
   * @param clock
   * @param stepsToFullRecovery
   * @param timeInterval
   * @param strategy
   */
public void clusterRecovery1TC(Map<String, Object> myMap, TestClock clock, int stepsToFullRecovery, Long timeInterval, DegraderLoadBalancerStrategyAdapter strategy, DegraderLoadBalancerStrategyV2_1.DegraderLoadBalancerState.Strategy strategyV2, DegraderLoadBalancerStrategyV3.PartitionDegraderLoadBalancerState.Strategy strategyV3) {
    final int NUM_CHECKS = 5;
    final Long TIME_INTERVAL = timeInterval;
    int localStepsToFullRecovery = stepsToFullRecovery;
    DegraderLoadBalancerStrategyConfig config = DegraderLoadBalancerStrategyConfig.createHttpConfigFromMap(myMap);
    List<TrackerClient> clients = new ArrayList<TrackerClient>();
    URI uri1 = URI.create("http://test.linkedin.com:3242/fdsaf");
    URIRequest request = new URIRequest(uri1);
    TrackerClient client1 = new TrackerClient(uri1, getDefaultPartitionData(1d), new TestLoadBalancerClient(uri1), clock, null);
    clients.add(client1);
    // force client1 to be disabled
    DegraderControl dcClient1Default = client1.getDegraderControl(DEFAULT_PARTITION_ID);
    dcClient1Default.setOverrideMinCallCount(5);
    dcClient1Default.setMinCallCount(5);
    dcClient1Default.setMaxDropRate(1d);
    dcClient1Default.setUpStep(1.0d);
    List<CallCompletion> ccList = new ArrayList<CallCompletion>();
    CallCompletion cc;
    for (int j = 0; j < NUM_CHECKS; j++) {
        cc = client1.getCallTracker().startCall();
        ccList.add(cc);
    }
    // add high latency and errors to shut off traffic to this tracker client.
    // note: the default values for highError and lowError in the degrader are 1.1,
    // which means we don't use errorRates when deciding when to lb/degrade.
    // In addition, because we changed to use the
    clock.addMs(3500);
    //for (int j = 0; j < NUM_CHECKS; j++)
    for (Iterator<CallCompletion> iter = ccList.listIterator(); iter.hasNext(); ) {
        cc = iter.next();
        cc.endCallWithError();
        iter.remove();
    }
    // go to next time interval.
    clock.addMs(TIME_INTERVAL);
    Assert.assertEquals(dcClient1Default.getCurrentComputedDropRate(), 1.0);
    // trigger a state update
    TrackerClient resultTC = getTrackerClient(strategy, request, new RequestContext(), 1, clients);
    if (config.getInitialRecoveryLevel() < 0.01) {
        //the returned TrackerClient should be null
        assertNull(resultTC, "expected null trackerclient");
        // tracker client, so it's time to try it out. We need to enter this code at least once.
        do {
            // go to next time interval.
            clock.addMs(TIME_INTERVAL);
            // try adjusting the hash ring on this updateState
            if (strategyV3 != null) {
                strategy.setStrategyV3(DEFAULT_PARTITION_ID, strategyV3);
            } else if (strategyV2 != null) {
                strategy.setStrategyV2(strategyV2);
            } else {
                fail("should set strategy (either LoadBalance or Degrader");
            }
            resultTC = getTrackerClient(strategy, request, new RequestContext(), 1, clients);
            localStepsToFullRecovery--;
        } while (localStepsToFullRecovery > 0);
    }
    assertNotNull(resultTC, "expected non-null trackerclient");
    // make calls to the tracker client to verify that it's on the road to healthy status.
    for (int j = 0; j < NUM_CHECKS; j++) {
        cc = resultTC.getCallTracker().startCall();
        ccList.add(cc);
    }
    clock.addMs(10);
    for (Iterator<CallCompletion> iter = ccList.listIterator(); iter.hasNext(); ) {
        cc = iter.next();
        cc.endCall();
        iter.remove();
    }
    // go to next time interval.
    clock.addMs(TIME_INTERVAL);
    Assert.assertTrue(dcClient1Default.getCurrentComputedDropRate() < 1d);
    resultTC = getTrackerClient(strategy, request, new RequestContext(), 1, clients);
    assertNotNull(resultTC, "expected non-null trackerclient");
}
Also used : ArrayList(java.util.ArrayList) URIRequest(com.linkedin.d2.balancer.util.URIRequest) DegraderControl(com.linkedin.util.degrader.DegraderControl) URI(java.net.URI) TrackerClient(com.linkedin.d2.balancer.clients.TrackerClient) CallCompletion(com.linkedin.util.degrader.CallCompletion) AtomicLong(java.util.concurrent.atomic.AtomicLong) RequestContext(com.linkedin.r2.message.RequestContext)

Example 17 with URIRequest

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

the class RingBasedURIMapperTest method testUniversalStickiness.

@Test
public void testUniversalStickiness() throws ServiceUnavailableException, URISyntaxException {
    int partitionCount = 4;
    int totalHostCount = 200;
    HashRingProvider ringProvider = createStaticHashRingProvider(totalHostCount, partitionCount, getHashFunction(true));
    HashFunction<Request> hashFunction = ringProvider.getRequestHashFunction(TEST_SERVICE);
    PartitionInfoProvider infoProvider = createRangeBasedPartitionInfoProvider(partitionCount);
    URIMapper mapper = new RingBasedUriMapper(ringProvider, infoProvider);
    // no partition, will be unmapped
    URIKeyPair<Integer> request1 = new URIKeyPair<>(1, new URI("d2://testService/1"));
    // partition 0
    URIKeyPair<Integer> request2 = new URIKeyPair<>(2, new URI("d2://testService/2?partition=0"));
    // partition 1
    URIKeyPair<Integer> request3 = new URIKeyPair<>(3, new URI("d2://testService/3?partition=1"));
    // partition 2
    URIKeyPair<Integer> request4 = new URIKeyPair<>(4, new URI("d2://testService/4?partition=2"));
    // partition 3
    URIKeyPair<Integer> request5 = new URIKeyPair<>(5, new URI("d2://testService/5?partition=3"));
    URIKeyPair<Integer> request6 = // partition 0 with different sticky key
    new URIKeyPair<>(6, new URI("d2://testService/6?partition=0"));
    URIKeyPair<Integer> request7 = // partition 1 with different sticky key
    new URIKeyPair<>(7, new URI("d2://testService/7?partition=1"));
    URIKeyPair<Integer> request8 = // partition 2 with different sticky key
    new URIKeyPair<>(8, new URI("d2://testService/8?partition=2"));
    URIKeyPair<Integer> request9 = // partition 3 with different sticky key
    new URIKeyPair<>(9, new URI("d2://testService/9?partition=3"));
    URIKeyPair<Integer> request10 = // with extra parameters
    new URIKeyPair<>(10, new URI("d2://testService/10?partition=0&uuid=1"));
    List<URIKeyPair<Integer>> requests = Arrays.asList(request1, request2, request3, request4, request5, request6, request7, request8, request9, request10);
    // uriMapper mapping
    URIMappingResult<Integer> uriMapperResult = mapper.mapUris(requests);
    // normal mapping
    Map<Integer, Set<Integer>> normalUnmapped = new HashMap<>();
    Map<URI, Set<Integer>> normalHostToKeySet = new HashMap<>();
    for (URIKeyPair<Integer> request : requests) {
        int partitionId = 0;
        try {
            partitionId = infoProvider.getPartitionAccessor(TEST_SERVICE).getPartitionId(request.getRequestUri());
        } catch (PartitionAccessException e) {
            normalUnmapped.computeIfAbsent(-1, k -> new HashSet<>()).add(request.getKey());
        }
        Ring<URI> ring = ringProvider.getRings(request.getRequestUri()).get(partitionId);
        URI uri = ring.get(hashFunction.hash(new URIRequest(request.getRequestUri())));
        normalHostToKeySet.computeIfAbsent(uri, k -> new HashSet<>());
        normalHostToKeySet.get(uri).add(request.getKey());
    }
    // they should have the same results
    Assert.assertEquals(uriMapperResult.getUnmappedKeys(), normalUnmapped);
    for (Map.Entry<URI, Set<Integer>> resolvedKeys : uriMapperResult.getMappedKeys().entrySet()) {
        Set<Integer> uriMapperKeySet = resolvedKeys.getValue();
        Assert.assertTrue(normalHostToKeySet.containsKey(resolvedKeys.getKey()));
        Set<Integer> normalKeySet = normalHostToKeySet.get(resolvedKeys.getKey());
        Assert.assertEquals(uriMapperKeySet, normalKeySet);
    }
}
Also used : Set(java.util.Set) HashSet(java.util.HashSet) PartitionInfoProvider(com.linkedin.d2.balancer.util.partitions.PartitionInfoProvider) HashMap(java.util.HashMap) URIRequest(com.linkedin.d2.balancer.util.URIRequest) Request(com.linkedin.r2.message.Request) URIRequest(com.linkedin.d2.balancer.util.URIRequest) URI(java.net.URI) URIKeyPair(com.linkedin.d2.balancer.util.URIKeyPair) URIMapper(com.linkedin.d2.balancer.URIMapper) PartitionAccessException(com.linkedin.d2.balancer.util.partitions.PartitionAccessException) HashMap(java.util.HashMap) Map(java.util.Map) Test(org.testng.annotations.Test)

Example 18 with URIRequest

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

the class ZKFSTest method testZKDown.

@Test
public void testZKDown() throws Exception {
    final String TEST_SERVICE_NAME = "testingService";
    final String TEST_CLUSTER_NAME = "someCluster";
    startServer();
    try {
        ZKFSLoadBalancer balancer = getBalancer();
        FutureCallback<None> callback = new FutureCallback<>();
        balancer.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        ZKConnection conn = new ZKConnection("localhost:" + PORT, 30000);
        conn.start();
        ZooKeeperPermanentStore<ServiceProperties> store = new ZooKeeperPermanentStore<>(conn, new ServicePropertiesJsonSerializer(), ZKFSUtil.servicePath(BASE_PATH));
        callback = new FutureCallback<>();
        store.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        ServiceProperties props = new ServiceProperties(TEST_SERVICE_NAME, TEST_CLUSTER_NAME, "/somePath", Arrays.asList("degrader"), Collections.<String, Object>emptyMap(), null, null, Arrays.asList("http"), null);
        store.put(TEST_SERVICE_NAME, props);
        ZooKeeperPermanentStore<ClusterProperties> clusterStore = new ZooKeeperPermanentStore<>(conn, new ClusterPropertiesJsonSerializer(), ZKFSUtil.clusterPath(BASE_PATH));
        callback = new FutureCallback<>();
        clusterStore.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        ClusterProperties clusterProps = new ClusterProperties("someCluster");
        clusterStore.put(TEST_CLUSTER_NAME, clusterProps);
        ZKConnection serverConn = new ZKConnection("localhost:" + PORT, 30000);
        serverConn.start();
        ZooKeeperEphemeralStore<UriProperties> uriStore = new ZooKeeperEphemeralStore<>(serverConn, new UriPropertiesJsonSerializer(), new UriPropertiesMerger(), ZKFSUtil.uriPath(BASE_PATH));
        callback = new FutureCallback<>();
        uriStore.start(callback);
        callback.get(30, TimeUnit.SECONDS);
        ZooKeeperServer server = new ZooKeeperServer(uriStore);
        callback = new FutureCallback<>();
        Map<Integer, PartitionData> partitionDataMap = new HashMap<>();
        partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1.0));
        server.markUp(TEST_CLUSTER_NAME, URI.create("http://test.uri"), partitionDataMap, callback);
        callback.get(30, TimeUnit.SECONDS);
        URIRequest request = new URIRequest("d2://" + TEST_SERVICE_NAME + "/foo");
        TransportClient client = balancer.getClient(request, new RequestContext());
        // Stop the server to cause a disconnect event
        stopServer();
        // Sleep to ensure the disconnect has propagated; ideally the Toggle should expose
        // some interface to allow detection that the toggle occurred
        Thread.sleep(1000);
        // Now see if it still works
        client = balancer.getClient(request, new RequestContext());
    } finally {
        stopServer();
    }
}
Also used : ServicePropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ServicePropertiesJsonSerializer) HashMap(java.util.HashMap) ZooKeeperPermanentStore(com.linkedin.d2.discovery.stores.zk.ZooKeeperPermanentStore) UriPropertiesMerger(com.linkedin.d2.balancer.properties.UriPropertiesMerger) UriPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.UriPropertiesJsonSerializer) ClusterPropertiesJsonSerializer(com.linkedin.d2.balancer.properties.ClusterPropertiesJsonSerializer) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) RequestContext(com.linkedin.r2.message.RequestContext) FutureCallback(com.linkedin.common.callback.FutureCallback) TransportClient(com.linkedin.r2.transport.common.bridge.client.TransportClient) ZKConnection(com.linkedin.d2.discovery.stores.zk.ZKConnection) URIRequest(com.linkedin.d2.balancer.util.URIRequest) ZooKeeperEphemeralStore(com.linkedin.d2.discovery.stores.zk.ZooKeeperEphemeralStore) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) None(com.linkedin.common.util.None) ZooKeeperServer(com.linkedin.d2.balancer.servers.ZooKeeperServer) Test(org.testng.annotations.Test)

Example 19 with URIRequest

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

the class SimpleLoadBalancerTest method testGetClientWithBannedURI.

@Test
public void testGetClientWithBannedURI() 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();
    URI uri1Banned = URI.create("http://test.qd.com:1234");
    URI uri2Usable = URI.create("http://test.qd.com:5678");
    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(uri1Banned, partitionData);
    uriData.put(uri2Usable, partitionData);
    prioritizedSchemes.add(PropertyKeys.HTTP_SCHEME);
    Set<URI> bannedSet = new HashSet<>();
    bannedSet.add(uri1Banned);
    clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", Collections.emptyList(), Collections.emptyMap(), bannedSet, NullPartitionProperties.getInstance()));
    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 expectedUri = URI.create("http://test.qd.com:5678/foo");
    URIRequest uriRequest = new URIRequest("d2://foo/52");
    for (int i = 0; i < 10; ++i) {
        RewriteLoadBalancerClient client = (RewriteLoadBalancerClient) loadBalancer.getClient(uriRequest, new RequestContext());
        Assert.assertEquals(client.getUri(), expectedUri);
    }
}
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) 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) RewriteLoadBalancerClient(com.linkedin.d2.balancer.clients.RewriteLoadBalancerClient) 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 20 with URIRequest

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

the class SimpleLoadBalancerTest method testLoadBalancerWithPartitionsSmoke.

// load balancer working with partitioned cluster
@Test(groups = { "small", "back-end" })
public void testLoadBalancerWithPartitionsSmoke() throws URISyntaxException, ServiceUnavailableException, InterruptedException, ExecutionException {
    for (int tryAgain = 0; tryAgain < 12; ++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 = Executors.newSingleThreadScheduledExecutor();
        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, executorService);
        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<URI, Double> uris = new HashMap<>();
        uris.put(uri1, 1d);
        uris.put(uri2, 1d);
        uris.put(uri3, 1d);
        Map<URI, Map<Integer, PartitionData>> partitionDesc = new HashMap<>();
        Map<Integer, PartitionData> server1 = new HashMap<>();
        server1.put(0, new PartitionData(1d));
        server1.put(1, new PartitionData(1d));
        Map<Integer, PartitionData> server2 = new HashMap<>();
        server2.put(0, new PartitionData(1d));
        Map<Integer, PartitionData> server3 = new HashMap<>();
        server3.put(1, new PartitionData(1d));
        partitionDesc.put(uri1, server1);
        partitionDesc.put(uri2, server2);
        partitionDesc.put(uri3, server3);
        prioritizedSchemes.add(PropertyKeys.HTTP_SCHEME);
        int partitionMethod = tryAgain % 4;
        switch(partitionMethod) {
            case 0:
                clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<>(), new HashSet<>(), new RangeBasedPartitionProperties("id=(\\d+)", 0, 50, 2)));
                break;
            case 1:
                clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<>(), new HashSet<>(), new HashBasedPartitionProperties("id=(\\d+)", 2, HashBasedPartitionProperties.HashAlgorithm.valueOf("MODULO"))));
                break;
            case 2:
                clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<>(), new HashSet<>(), new HashBasedPartitionProperties("id=(\\d+)", 2, HashBasedPartitionProperties.HashAlgorithm.valueOf("MD5"))));
                break;
            case 3:
                // test getRings with gap. here, no server serves partition 2
                clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1", null, new HashMap<>(), new HashSet<>(), new RangeBasedPartitionProperties("id=(\\d+)", 0, 50, 4)));
                server3.put(3, new PartitionData(1d));
                partitionDesc.put(uri3, server3);
                break;
            default:
                break;
        }
        serviceRegistry.put("foo", new ServiceProperties("foo", "cluster-1", "/foo", Arrays.asList("degrader"), Collections.singletonMap(PropertyKeys.HTTP_LB_CONSISTENT_HASH_ALGORITHM, "pointBased"), null, null, prioritizedSchemes, null));
        uriRegistry.put("cluster-1", new UriProperties("cluster-1", partitionDesc));
        if (partitionMethod == 3) {
            Map<Integer, Ring<URI>> ringMap = loadBalancer.getRings(URI.create("d2://foo"));
            assertEquals(ringMap.size(), 4);
            // the ring for partition 2 should be empty
            assertEquals(ringMap.get(2).toString(), new ConsistentHashRing<>(Collections.emptyList()).toString());
            continue;
        }
        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);
        for (int i = 0; i < 1000; ++i) {
            int ii = i % 100;
            RewriteLoadBalancerClient client = (RewriteLoadBalancerClient) loadBalancer.getClient(new URIRequest("d2://foo/id=" + ii), new RequestContext());
            String clientUri = client.getUri().toString();
            HashFunction<String[]> hashFunction = null;
            String[] str = new String[1];
            // test KeyMapper target host hint: request is always to target host regardless of what's in d2 URI and whether it's hash-based or range-based partitions
            RequestContext requestContextWithHint = new RequestContext();
            KeyMapper.TargetHostHints.setRequestContextTargetHost(requestContextWithHint, uri1);
            RewriteLoadBalancerClient hintedClient1 = (RewriteLoadBalancerClient) loadBalancer.getClient(new URIRequest("d2://foo/id=" + ii), requestContextWithHint);
            String hintedUri1 = hintedClient1.getUri().toString();
            Assert.assertEquals(hintedUri1, uri1.toString() + "/foo");
            RewriteLoadBalancerClient hintedClient2 = (RewriteLoadBalancerClient) loadBalancer.getClient(new URIRequest("d2://foo/action=purge-all"), requestContextWithHint);
            String hintedUri2 = hintedClient2.getUri().toString();
            Assert.assertEquals(hintedUri2, uri1.toString() + "/foo");
            if (partitionMethod == 2) {
                hashFunction = new MD5Hash();
            }
            for (URI uri : expectedUris) {
                if (clientUri.contains(uri.toString())) {
                    // check if only key belonging to partition 0 gets uri2
                    if (uri.equals(uri2)) {
                        if (partitionMethod == 0) {
                            assertTrue(ii < 50);
                        } else if (partitionMethod == 1) {
                            assertTrue(ii % 2 == 0);
                        } else {
                            str[0] = ii + "";
                            assertTrue(hashFunction.hash(str) % 2 == 0);
                        }
                    }
                    // check if only key belonging to partition 1 gets uri3
                    if (uri.equals(uri3)) {
                        if (partitionMethod == 0) {
                            assertTrue(ii >= 50);
                        } else if (partitionMethod == 1) {
                            assertTrue(ii % 2 == 1);
                        } else {
                            str[0] = ii + "";
                            assertTrue(hashFunction.hash(str) % 2 == 1);
                        }
                    }
                }
            }
        }
        // two rings for two partitions
        Map<Integer, Ring<URI>> ringMap = loadBalancer.getRings(URI.create("d2://foo"));
        assertEquals(ringMap.size(), 2);
        if (partitionMethod != 2) {
            Set<String> keys = new HashSet<>();
            for (int j = 0; j < 50; j++) {
                if (partitionMethod == 0) {
                    keys.add(j + "");
                } else {
                    keys.add(j * 2 + "");
                }
            }
            // if it is range based partition, all keys from 0 ~ 49 belong to partition 0 according to the range definition
            // if it is modulo based partition, all even keys belong to partition 0 because the partition count is 2
            // only from partition 0
            MapKeyResult<Ring<URI>, String> mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
            Map<Ring<URI>, Collection<String>> keyToPartition = mapKeyResult.getMapResult();
            assertEquals(keyToPartition.size(), 1);
            for (Ring<URI> ring : keyToPartition.keySet()) {
                assertEquals(ring, ringMap.get(0));
            }
            // now also from partition 1
            keys.add("51");
            mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
            assertEquals(mapKeyResult.getMapResult().size(), 2);
            assertEquals(mapKeyResult.getUnmappedKeys().size(), 0);
            // now only from partition 1
            keys.clear();
            keys.add("99");
            mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
            keyToPartition = mapKeyResult.getMapResult();
            assertEquals(keyToPartition.size(), 1);
            assertEquals(mapKeyResult.getUnmappedKeys().size(), 0);
            for (Ring<URI> ring : keyToPartition.keySet()) {
                assertEquals(ring, ringMap.get(1));
            }
            keys.add("100");
            mapKeyResult = loadBalancer.getRings(URI.create("d2://foo"), keys);
            if (partitionMethod == 0) {
                // key out of range
                Collection<MapKeyResult.UnmappedKey<String>> unmappedKeys = mapKeyResult.getUnmappedKeys();
                assertEquals(unmappedKeys.size(), 1);
            }
            try {
                loadBalancer.getClient(new URIRequest("d2://foo/id=100"), new RequestContext());
                if (partitionMethod == 0) {
                    // key out of range
                    fail("Should throw ServiceUnavailableException caused by PartitionAccessException");
                }
            } catch (ServiceUnavailableException e) {
            }
        }
        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 : HashMap(java.util.HashMap) DegraderLoadBalancerStrategyFactoryV3(com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3) ArrayList(java.util.ArrayList) PartitionData(com.linkedin.d2.balancer.properties.PartitionData) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) RangeBasedPartitionProperties(com.linkedin.d2.balancer.properties.RangeBasedPartitionProperties) RequestContext(com.linkedin.r2.message.RequestContext) FutureCallback(com.linkedin.common.callback.FutureCallback) HashSet(java.util.HashSet) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) RandomLoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategyFactory) LoadBalancerStrategyFactory(com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory) HashBasedPartitionProperties(com.linkedin.d2.balancer.properties.HashBasedPartitionProperties) RewriteLoadBalancerClient(com.linkedin.d2.balancer.clients.RewriteLoadBalancerClient) URIRequest(com.linkedin.d2.balancer.util.URIRequest) Collection(java.util.Collection) Map(java.util.Map) DarkClusterConfigMap(com.linkedin.d2.DarkClusterConfigMap) HashMap(java.util.HashMap) PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) ServiceUnavailableException(com.linkedin.d2.balancer.ServiceUnavailableException) URI(java.net.URI) ConsistentHashRing(com.linkedin.d2.balancer.util.hashing.ConsistentHashRing) TransportClientFactory(com.linkedin.r2.transport.common.TransportClientFactory) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) CountDownLatch(java.util.concurrent.CountDownLatch) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) ConsistentHashRing(com.linkedin.d2.balancer.util.hashing.ConsistentHashRing) Ring(com.linkedin.d2.balancer.util.hashing.Ring) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) MD5Hash(com.linkedin.d2.balancer.util.hashing.MD5Hash) None(com.linkedin.common.util.None) Test(org.testng.annotations.Test)

Aggregations

URIRequest (com.linkedin.d2.balancer.util.URIRequest)26 RequestContext (com.linkedin.r2.message.RequestContext)24 URI (java.net.URI)20 ArrayList (java.util.ArrayList)17 Test (org.testng.annotations.Test)15 HashMap (java.util.HashMap)14 TrackerClient (com.linkedin.d2.balancer.clients.TrackerClient)12 Map (java.util.Map)11 ClusterProperties (com.linkedin.d2.balancer.properties.ClusterProperties)9 ServiceProperties (com.linkedin.d2.balancer.properties.ServiceProperties)9 UriProperties (com.linkedin.d2.balancer.properties.UriProperties)9 LoadBalancerStrategy (com.linkedin.d2.balancer.strategies.LoadBalancerStrategy)9 HashSet (java.util.HashSet)9 FutureCallback (com.linkedin.common.callback.FutureCallback)8 None (com.linkedin.common.util.None)8 DegraderTrackerClient (com.linkedin.d2.balancer.clients.DegraderTrackerClient)8 TransportClientFactory (com.linkedin.r2.transport.common.TransportClientFactory)8 PartitionData (com.linkedin.d2.balancer.properties.PartitionData)7 LoadBalancerStrategyFactory (com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory)7 DegraderLoadBalancerStrategyFactoryV3 (com.linkedin.d2.balancer.strategies.degrader.DegraderLoadBalancerStrategyFactoryV3)7