Examples with RewriteLoadBalancerClient com.linkedin.d2.balancer.clients.RewriteLoadBalancerClient used on opensource projects

Example 1 with RewriteLoadBalancerClient

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

the class SimpleLoadBalancer method getClient.

/**
 * Given a Request, returns a TransportClient that can handle requests for the Request.
 * The callback is given a client that can be called to retrieve data for the URN.
 *
 * @param request
 *          A request whose URI is a URL of the format "d2://>servicename</optional/path".
 * @param requestContext context for this request
 * @throws ServiceUnavailableException
 *           If the load balancer can't figure out how to reach a service for the given
 *           URN, an ServiceUnavailableException will be thrown.
 */
@Override
public void getClient(Request request, RequestContext requestContext, Callback<TransportClient> clientCallback) {
    URI uri = request.getURI();
    debug(_log, "get client for uri: ", uri);
    if (!D2_SCHEME_NAME.equalsIgnoreCase(uri.getScheme())) {
        throw new IllegalArgumentException("Unsupported scheme in URI " + uri);
    }
    // get the service for this uri
    String extractedServiceName = LoadBalancerUtil.getServiceNameFromUri(uri);
    listenToServiceAndCluster(extractedServiceName, Callbacks.handle(service -> {
        String serviceName = service.getServiceName();
        String clusterName = service.getClusterName();
        try {
            ClusterProperties cluster = getClusterProperties(serviceName, clusterName);
            // Check if we want to override the service URL and bypass choosing among the existing
            // tracker clients. This is useful when the service we want is not announcing itself to
            // the cluster, ie a private service for a set of clients. This mechanism is deprecated;
            // use host override list instead.
            @SuppressWarnings("deprecation") URI targetService = LoadBalancerUtil.TargetHints.getRequestContextTargetService(requestContext);
            // Checks if we have a host override list provided in the request context. If present,
            // get the override URI available override for the current cluster and service names.
            HostOverrideList overrides = (HostOverrideList) requestContext.getLocalAttr(HOST_OVERRIDE_LIST);
            URI override = overrides == null ? null : overrides.getOverride(clusterName, serviceName);
            if (targetService == null && override == null) {
                LoadBalancerStateItem<UriProperties> uriItem = getUriItem(serviceName, clusterName, cluster);
                UriProperties uris = uriItem.getProperty();
                List<LoadBalancerState.SchemeStrategyPair> orderedStrategies = _state.getStrategiesForService(serviceName, service.getPrioritizedSchemes());
                TrackerClient trackerClient = chooseTrackerClient(request, requestContext, serviceName, clusterName, cluster, uriItem, uris, orderedStrategies, service);
                String clusterAndServiceUriString = trackerClient.getUri() + service.getPath();
                _serviceAvailableStats.inc();
                clientCallback.onSuccess(new RewriteLoadBalancerClient(serviceName, URI.create(clusterAndServiceUriString), trackerClient));
            } else {
                URI target = override == null ? targetService : URI.create(override + service.getPath());
                if (targetService != null && override != null) {
                    _log.warn("Both TargetHints and HostOverrideList are found. HostOverList will take precedence %s.", target);
                }
                if (_log.isDebugEnabled()) {
                    _log.debug("Rewrite URI as specified in the TargetHints/HostOverrideList {} for cluster {} and service {}.", target, clusterName, serviceName);
                }
                TransportClient transportClient = _state.getClient(serviceName, target.getScheme());
                if (transportClient == null) {
                    throw new ServiceUnavailableException(serviceName, String.format("PEGA_1001. Cannot find transportClient for service %s and scheme %s with URI specified in" + "TargetHints/HostOverrideList %s", serviceName, target.getScheme(), target));
                }
                clientCallback.onSuccess(new RewriteLoadBalancerClient(serviceName, target, transportClient));
            }
        } catch (ServiceUnavailableException e) {
            clientCallback.onError(e);
        }
    }, clientCallback));
}

Example 2 with RewriteLoadBalancerClient

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

the class RewriteLoadBalancerClientTest method testWithEverything.

@Test
public void testWithEverything() {
    URI uri = URI.create("http://username:password@test.linkedin.com:9876/test");
    String serviceName = "HistoryService";
    TestClient wrappedClient = new TestClient();
    RewriteLoadBalancerClient client = new RewriteLoadBalancerClient(serviceName, uri, wrappedClient);
    assertEquals(client.getUri(), uri);
    assertEquals(client.getServiceName(), serviceName);
    RestRequest restRequest = new RestRequestBuilder(URI.create("d2://HistoryService/getCube?bar=baz#fragId")).build();
    Map<String, String> restWireAttrs = new HashMap<>();
    TestTransportCallback<RestResponse> restCallback = new TestTransportCallback<>();
    client.restRequest(restRequest, new RequestContext(), restWireAttrs, restCallback);
    assertFalse(restCallback.response.hasError());
    assertEquals(wrappedClient.restRequest.getHeaders(), restRequest.getHeaders());
    assertEquals(wrappedClient.restRequest.getEntity(), restRequest.getEntity());
    assertEquals(wrappedClient.restRequest.getMethod(), restRequest.getMethod());
    assertEquals(wrappedClient.restRequest.getURI(), URI.create("http://username:password@test.linkedin.com:9876/test/getCube?bar=baz#fragId"));
}

Example 3 with RewriteLoadBalancerClient

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

the class RewriteLoadBalancerClientTest method testClient.

@Test(groups = { "small", "back-end" })
public void testClient() throws URISyntaxException {
    URI uri = URI.create("http://test.linkedin.com/test");
    String serviceName = "HistoryService";
    TestClient wrappedClient = new TestClient();
    RewriteLoadBalancerClient client = new RewriteLoadBalancerClient(serviceName, uri, wrappedClient);
    assertEquals(client.getUri(), uri);
    assertEquals(client.getServiceName(), serviceName);
    RestRequest restRequest = new RestRequestBuilder(URI.create("d2://HistoryService/getCube")).build();
    Map<String, String> restWireAttrs = new HashMap<>();
    TestTransportCallback<RestResponse> restCallback = new TestTransportCallback<>();
    client.restRequest(restRequest, new RequestContext(), restWireAttrs, restCallback);
    assertFalse(restCallback.response.hasError());
    assertEquals(wrappedClient.restRequest.getHeaders(), restRequest.getHeaders());
    assertEquals(wrappedClient.restRequest.getEntity(), restRequest.getEntity());
    assertEquals(wrappedClient.restRequest.getMethod(), restRequest.getMethod());
    // check the rewrite
    assertEquals(wrappedClient.restRequest.getURI(), URI.create("http://test.linkedin.com/test/getCube"));
    assertEquals(wrappedClient.restWireAttrs, restWireAttrs);
}

Example 4 with RewriteLoadBalancerClient

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

the class RewriteLoadBalancerClientTestStreamRequest method testEscapingHelper.

private void testEscapingHelper(String hostUri, String serviceName, String path) {
    URI uri = URI.create(hostUri);
    TestClient wrappedClient = new TestClient();
    RewriteLoadBalancerClient client = new RewriteLoadBalancerClient(serviceName, uri, wrappedClient);
    assertEquals(client.getUri(), uri);
    assertEquals(client.getServiceName(), serviceName);
    StreamRequest streamRequest = getRequest("d2://" + serviceName + path);
    Map<String, String> restWireAttrs = new HashMap<>();
    TestTransportCallback<StreamResponse> restCallback = new TestTransportCallback<>();
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    assertFalse(restCallback.response.hasError());
    assertEquals(wrappedClient.streamRequest.getHeaders(), streamRequest.getHeaders());
    assertEquals(wrappedClient.streamRequest.getMethod(), streamRequest.getMethod());
    assertEquals(wrappedClient.streamRequest.getURI(), URI.create(hostUri + path));
}

Example 5 with RewriteLoadBalancerClient

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

the class RewriteLoadBalancerClientTestStreamRequest method testClient.

@Test(groups = { "small", "back-end" })
public void testClient() throws URISyntaxException {
    URI uri = URI.create("http://test.linkedin.com/test");
    String serviceName = "HistoryService";
    TestClient wrappedClient = new TestClient();
    RewriteLoadBalancerClient client = new RewriteLoadBalancerClient(serviceName, uri, wrappedClient);
    assertEquals(client.getUri(), uri);
    assertEquals(client.getServiceName(), serviceName);
    StreamRequest streamRequest = new StreamRequestBuilder(URI.create("d2://HistoryService/getCube")).build(EntityStreams.emptyStream());
    Map<String, String> restWireAttrs = new HashMap<>();
    TestTransportCallback<StreamResponse> restCallback = new TestTransportCallback<>();
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    assertFalse(restCallback.response.hasError());
    assertEquals(wrappedClient.streamRequest.getHeaders(), streamRequest.getHeaders());
    assertEquals(wrappedClient.streamRequest.getMethod(), streamRequest.getMethod());
    // check the rewrite
    assertEquals(wrappedClient.streamRequest.getURI(), URI.create("http://test.linkedin.com/test/getCube"));
    assertEquals(wrappedClient.restWireAttrs, restWireAttrs);
}