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

Example 1 with HostSet

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

the class AllPartitionsRequestBuilder method sendRequests.

/**
   * A convenience function for caller to issue search request with one call.
   * If finer-grain control is required, users should call buildRequestContexts instead and send requests by themselves
   *
   * @param client - the RestClient to use
   * @param request - the query-all request
   * @param requestContext - the original request context
   * @param callback - callback to be used for each request
   * @return the partition cover which informs the caller the RequestContexts, the total number of partitions and the available partition number
   * @throws ServiceUnavailableException
   */
public HostSet sendRequests(RestClient client, Request<T> request, RequestContext requestContext, Callback<Response<T>> callback) throws ServiceUnavailableException {
    URI serviceUri;
    try {
        serviceUri = new URI(D2_URI_PREFIX + request.getServiceName());
    } catch (URISyntaxException e) {
        throw new IllegalArgumentException(e);
    }
    final Collection<RequestContext> queryAllRequestContext = new ArrayList<RequestContext>();
    final HostSet uriResult = _mapper.getAllPartitionsMultipleHosts(serviceUri, 1);
    for (URI targetHost : uriResult.getAllHosts()) {
        RequestContext context = requestContext.clone();
        KeyMapper.TargetHostHints.setRequestContextTargetHost(context, targetHost);
        queryAllRequestContext.add(context);
        client.sendRequest(request, context, callback);
    }
    return uriResult;
}

Example 2 with HostSet

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

the class TestAllPartitionsRequestBuilder method testSendAllPartitionsRequests.

@Test(dataProvider = com.linkedin.restli.internal.common.TestConstants.RESTLI_PROTOCOL_1_2_PREFIX + "restliRequestOptions")
public void testSendAllPartitionsRequests(RestliRequestOptions options, RingFactory<URI> ringFactory) throws ServiceUnavailableException, URISyntaxException, RestException, InterruptedException {
    final int PARTITION_NUM = 5;
    List<URI> expectedUris = new ArrayList<URI>();
    ConsistentHashKeyMapper mapper = getKeyToHostMapper(PARTITION_NUM, expectedUris, ringFactory);
    AllPartitionsRequestBuilder<Greeting> searchRB = new AllPartitionsRequestBuilder<Greeting>(mapper);
    ActionRequestBuilder<Long, Greeting> builder = new ActionRequestBuilder<Long, Greeting>(TEST_URI, Greeting.class, _COLL_SPEC, options);
    ActionRequest<Greeting> request = builder.name("updateTone").id(1L).setParam(new FieldDef<Tone>("newTone", Tone.class, DataTemplateUtil.getSchema(Tone.class)), Tone.FRIENDLY).build();
    final Map<String, Greeting> results = new ConcurrentHashMap<String, Greeting>();
    final CountDownLatch latch = new CountDownLatch(PARTITION_NUM);
    final List<Throwable> errors = new ArrayList<Throwable>();
    final List<Greeting> responses = new ArrayList<Greeting>();
    Callback<Response<Greeting>> cb = new Callback<Response<Greeting>>() {

        @Override
        public void onError(Throwable e) {
            synchronized (errors) {
                errors.add(e);
            }
            latch.countDown();
        }

        @Override
        public void onSuccess(Response<Greeting> response) {
            results.put(response.getEntity().toString(), response.getEntity());
            synchronized (responses) {
                responses.add(response.getEntity());
            }
            latch.countDown();
        }
    };
    HostSet hostsResult = searchRB.sendRequests(getClient(), request, new RequestContext(), cb);
    List<URI> uris = hostsResult.getAllHosts();
    Assert.assertTrue(uris.containsAll(expectedUris));
    Assert.assertTrue(expectedUris.containsAll(uris));
    latch.await();
    if (!errors.isEmpty()) {
        Assert.fail("I knew it: " + errors.toString());
    }
    Assert.assertEquals(PARTITION_NUM, responses.size());
}