Examples with URIMappingResult com.linkedin.d2.balancer.util.URIMappingResult used on opensource projects

Example 1 with URIMappingResult

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

the class RingBasedUriMapper method mapUris.

/**
 * To achieve scatter-gather, there will be two passes.
 *
 * Pass 1: All requests are assigned a partitionId based on partition properties
 *         If partitioning is not enabled, all requests will have default partitionId of 0;
 *
 * Pass 2: All requests in the same partition will be routed based on the ring of that partition.
 *         If sticky routing is not specified, ONE host on the ring of that partition will be assigned for all hosts assigned to that partition.
 *
 * Unmapped key in either step will be collected in the unmapped keySet in the result.
 *
 * @param <KEY> type of provided key
 * @param requestUriKeyPairs a list of URIKeyPair, each contains a d2 request uri and a unique resource key.
 * @return {@link URIMappingResult} that contains host to keySet mapping as well as unmapped keys.
 * @throws ServiceUnavailableException when the requested service is not available
 */
@Override
public <KEY> URIMappingResult<KEY> mapUris(List<URIKeyPair<KEY>> requestUriKeyPairs) throws ServiceUnavailableException {
    if (requestUriKeyPairs == null || requestUriKeyPairs.isEmpty()) {
        return new URIMappingResult<>(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
    }
    // API assumes that all requests will be made to the same service, just use the first request to get the service name and act as sample uri
    URI sampleURI = requestUriKeyPairs.get(0).getRequestUri();
    String serviceName = LoadBalancerUtil.getServiceNameFromUri(sampleURI);
    // To achieve scatter-gather, we require the following information
    PartitionAccessor accessor = _partitionInfoProvider.getPartitionAccessor(serviceName);
    Map<Integer, Ring<URI>> rings = _hashRingProvider.getRings(sampleURI);
    HashFunction<Request> hashFunction = _hashRingProvider.getRequestHashFunction(serviceName);
    Map<Integer, Set<KEY>> unmapped = new HashMap<>();
    // Pass One
    Map<Integer, List<URIKeyPair<KEY>>> requestsByPartition = distributeToPartitions(requestUriKeyPairs, accessor, unmapped);
    // Pass Two
    Map<URI, Integer> hostToParitionId = new HashMap<>();
    Map<URI, Set<KEY>> hostToKeySet = distributeToHosts(requestsByPartition, rings, hashFunction, hostToParitionId, unmapped);
    return new URIMappingResult<>(hostToKeySet, unmapped, hostToParitionId);
}

Example 2 with URIMappingResult

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

the class RingBasedURIMapperTest method testPartitionIdOverride.

@Test(dataProvider = "stickyPartitionPermutation")
public void testPartitionIdOverride(boolean sticky, boolean partitioned) throws Exception {
    int partitionCount = partitioned ? 10 : 1;
    int totalHostCount = 100;
    HashRingProvider ringProvider = createStaticHashRingProvider(totalHostCount, partitionCount, getHashFunction(sticky));
    PartitionInfoProvider infoProvider = createRangeBasedPartitionInfoProvider(partitionCount);
    URIMapper mapper = new RingBasedUriMapper(ringProvider, infoProvider);
    URIKeyPair<Integer> request = new URIKeyPair<>(new URI("d2://testService/1"), IntStream.range(0, partitionCount).boxed().collect(Collectors.toSet()));
    if (partitioned) {
        Assert.assertThrows(() -> mapper.mapUris(Arrays.asList(request, request)));
    }
    URIMappingResult<Integer> uriMapperResult = mapper.mapUris(Collections.singletonList(request));
    Map<URI, Set<Integer>> mappedKeys = uriMapperResult.getMappedKeys();
    Assert.assertTrue(uriMapperResult.getUnmappedKeys().isEmpty());
    Assert.assertEquals(mappedKeys.size(), partitionCount);
    Assert.assertEquals(mappedKeys.keySet().stream().map(URIMapperTestUtil::getPartitionIdForURI).collect(Collectors.toSet()).size(), partitionCount);
    for (Set<Integer> keys : mappedKeys.values()) {
        Assert.assertTrue(keys.isEmpty());
    }
}

Example 3 with URIMappingResult

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

the class RingBasedURIMapperTest method testUnmappedKeys.

@Test(dataProvider = "stickyPartitionPermutation")
public void testUnmappedKeys(boolean sticky, boolean partitioned) throws Exception {
    int partitionCount = partitioned ? 10 : 1;
    int requestPerPartition = 100;
    List<Ring<URI>> rings = new ArrayList<>();
    IntStream.range(0, partitionCount).forEach(i -> rings.add(new MPConsistentHashRing<>(Collections.emptyMap())));
    StaticRingProvider ringProvider = new StaticRingProvider(rings);
    ringProvider.setHashFunction(getHashFunction(sticky));
    PartitionInfoProvider infoProvider = createRangeBasedPartitionInfoProvider(partitionCount);
    URIMapper mapper = new RingBasedUriMapper(ringProvider, infoProvider);
    List<URIKeyPair<Integer>> requests = testUtil.generateRequests(partitionCount, requestPerPartition);
    URIMappingResult<Integer> uriMapperResultNormal = mapper.mapUris(requests);
    Assert.assertTrue(uriMapperResultNormal.getUnmappedKeys().size() == partitionCount);
    uriMapperResultNormal.getUnmappedKeys().forEach((key, value) -> Assert.assertTrue(value.size() == requestPerPartition));
    URIKeyPair<Integer> request = new URIKeyPair<>(new URI("d2://testService/1"), IntStream.range(0, partitionCount).boxed().collect(Collectors.toSet()));
    URIMappingResult<Integer> uriMapperResultCustom = mapper.mapUris(Collections.singletonList(request));
    Assert.assertTrue(uriMapperResultCustom.getUnmappedKeys().size() == partitionCount);
    uriMapperResultCustom.getUnmappedKeys().forEach((key, value) -> Assert.assertTrue(value.isEmpty()));
}

Example 4 with URIMappingResult

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

the class RingBasedURIMapperTest method testErrorHandling.

@Test
public void testErrorHandling() throws ServiceUnavailableException, URISyntaxException {
    int partitionCount = 10;
    int totalHostCount = 100;
    HashRingProvider ringProvider = createStaticHashRingProvider(totalHostCount, partitionCount, getHashFunction(true));
    PartitionInfoProvider infoProvider = createRangeBasedPartitionInfoProvider(partitionCount);
    URIMapper mapper = new RingBasedUriMapper(ringProvider, infoProvider);
    URIKeyPair<Integer> requestWithoutPartitionId = new URIKeyPair<>(42, new URI("d2://badService/2"));
    URIKeyPair<Integer> requestWithoutKey = new URIKeyPair<>(43, new URI("d2://badService/partitionId=3"));
    URIKeyPair<Integer> requestWithoutBoth = new URIKeyPair<>(44, new URI("d2://badService"));
    List<URIKeyPair<Integer>> requests = Arrays.asList(requestWithoutKey, requestWithoutPartitionId, requestWithoutBoth);
    URIMappingResult<Integer> result = mapper.mapUris(requests);
    Assert.assertTrue(result.getMappedKeys().isEmpty());
    Assert.assertTrue(result.getUnmappedKeys().get(-1).contains(42));
    Assert.assertTrue(result.getUnmappedKeys().get(-1).contains(43));
    Assert.assertTrue(result.getUnmappedKeys().get(-1).contains(44));
}

Example 5 with URIMappingResult

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

the class RingBasedURIMapperTest method testStickyAndPartitioning.

@Test
public void testStickyAndPartitioning() throws ServiceUnavailableException {
    int partitionCount = 10;
    int requestPerPartition = 100;
    int totalHostCount = 100;
    HashRingProvider ringProvider = createStaticHashRingProvider(totalHostCount, partitionCount, getHashFunction(true));
    PartitionInfoProvider infoProvider = createRangeBasedPartitionInfoProvider(partitionCount);
    URIMapper mapper = new RingBasedUriMapper(ringProvider, infoProvider);
    List<URIKeyPair<Integer>> requests = testUtil.generateRequests(partitionCount, requestPerPartition);
    URIMappingResult<Integer> results = mapper.mapUris(requests);
    Map<URI, Set<Integer>> mapping = results.getMappedKeys();
    Map<Integer, Set<Integer>> unmappedKeys = results.getUnmappedKeys();
    Map<URI, Integer> hostToPartition = results.getHostPartitionInfo();
    Assert.assertTrue(unmappedKeys.isEmpty());
    Assert.assertEquals(100, mapping.size());
    Assert.assertEquals(100, hostToPartition.size());
}