use of com.linkedin.d2.balancer.util.URIKeyPair in project rest.li by linkedin.
the class RingBasedUriMapper method distributeToHosts.
private <KEY> Map<URI, Set<KEY>> distributeToHosts(Map<Integer, List<URIKeyPair<KEY>>> requestsByParititonId, Map<Integer, Ring<URI>> rings, HashFunction<Request> hashFunction, Map<URI, Integer> hostToPartitionId, Map<Integer, Set<KEY>> unmapped) {
if (hashFunction instanceof RandomHash) {
return distributeToHostNonSticky(requestsByParititonId, rings, hostToPartitionId, unmapped);
}
Map<URI, Set<KEY>> hostToKeySet = new HashMap<>();
for (Map.Entry<Integer, List<URIKeyPair<KEY>>> entry : requestsByParititonId.entrySet()) {
int partitionId = entry.getKey();
for (URIKeyPair<KEY> request : entry.getValue()) {
int hashcode = hashFunction.hash(new URIRequest(request.getRequestUri()));
URI resolvedHost = rings.get(partitionId).get(hashcode);
if (resolvedHost == null) {
// under custom use case, key will be null, in which case we will just return a map from partition id to empty set
// Users should be able to understand what partitions do not have available hosts by examining the keys in "unmapped"
Set<KEY> unmappedKeys = convertURIKeyPairListToKeySet(entry.getValue());
unmapped.computeIfAbsent(entry.getKey(), k -> new HashSet<>()).addAll(unmappedKeys);
break;
} else {
// under custom use case, key will be null, in which case we will just return a map from uri to empty set
hostToPartitionId.putIfAbsent(resolvedHost, entry.getKey());
Set<KEY> newSet = hostToKeySet.computeIfAbsent(resolvedHost, host -> new HashSet<>());
if (request.getKey() != null) {
newSet.add(request.getKey());
}
}
}
}
return hostToKeySet;
}
use of com.linkedin.d2.balancer.util.URIKeyPair in project rest.li by linkedin.
the class RingBasedUriMapper method distributeToHostNonSticky.
/**
* if sticky is not enabled, map all uris of the same partition to ONE host. If the same host is picked for multiple partitions,
* keys to those partitions will be merged into one set.
*/
private <KEY> Map<URI, Set<KEY>> distributeToHostNonSticky(Map<Integer, List<URIKeyPair<KEY>>> requestsByParititonId, Map<Integer, Ring<URI>> rings, Map<URI, Integer> hostToPartitionId, Map<Integer, Set<KEY>> unmapped) {
Map<URI, Set<KEY>> hostToKeySet = new HashMap<>();
for (Map.Entry<Integer, List<URIKeyPair<KEY>>> entry : requestsByParititonId.entrySet()) {
URI resolvedHost = rings.get(entry.getKey()).get(ThreadLocalRandom.current().nextInt());
Set<KEY> allKeys = convertURIKeyPairListToKeySet(entry.getValue());
if (resolvedHost == null) {
unmapped.computeIfAbsent(entry.getKey(), k -> new HashSet<>()).addAll(allKeys);
} else {
hostToPartitionId.putIfAbsent(resolvedHost, entry.getKey());
hostToKeySet.computeIfAbsent(resolvedHost, host -> new HashSet<>()).addAll(allKeys);
}
}
return hostToKeySet;
}
use of com.linkedin.d2.balancer.util.URIKeyPair 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);
}
use of com.linkedin.d2.balancer.util.URIKeyPair 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());
}
}
use of com.linkedin.d2.balancer.util.URIKeyPair 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()));
}
Aggregations