Examples with RestResponseBuilder com.linkedin.r2.message.rest.RestResponseBuilder used on opensource projects

Example 61 with RestResponseBuilder

use of com.linkedin.r2.message.rest.RestResponseBuilder in project rest.li by linkedin.

the class TestBackupRequestsClient method createAlwaysBackupClientWithHosts.

private BackupRequestsClient createAlwaysBackupClientWithHosts(List<String> uris, Deque<URI> hostsReceivingRequestList, int responseDelayNano, int backupDelayNano, boolean isD2Async) throws IOException {
    Map<URI, Map<Integer, PartitionData>> partitionDescriptions = new HashMap<>();
    uris.forEach(uri -> partitionDescriptions.put(URI.create(uri), Collections.singletonMap(0, new PartitionData(1))));
    StaticLoadBalancerState LbState = new StaticLoadBalancerState() {

        @Override
        public TrackerClient getClient(String serviceName, URI uri) {
            return new DegraderTrackerClientImpl(uri, partitionDescriptions.get(uri), null, SystemClock.instance(), null) {

                @Override
                public void restRequest(RestRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<RestResponse> callback) {
                    // whenever a trackerClient is used to make request, record down it's hostname
                    hostsReceivingRequestList.add(uri);
                    // delay response to allow backup request to happen
                    _executor.schedule(() -> callback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build())), responseDelayNano, TimeUnit.NANOSECONDS);
                }

                @Override
                public void streamRequest(StreamRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
                    // whenever a trackerClient is used to make request, record down it's hostname
                    hostsReceivingRequestList.add(uri);
                    if (null != requestContext.getLocalAttr(R2Constants.BACKUP_REQUEST_BUFFERED_BODY)) {
                        callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().setHeader(BUFFERED_HEADER, String.valueOf(requestContext.getLocalAttr(R2Constants.BACKUP_REQUEST_BUFFERED_BODY) != null)).build(EntityStreams.emptyStream())));
                        return;
                    }
                    request.getEntityStream().setReader(new DrainReader() {

                        public void onDone() {
                            // delay response to allow backup request to happen
                            _executor.schedule(() -> callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().setHeader(BUFFERED_HEADER, String.valueOf(requestContext.getLocalAttr(R2Constants.BACKUP_REQUEST_BUFFERED_BODY) != null)).build(EntityStreams.emptyStream()))), responseDelayNano, TimeUnit.NANOSECONDS);
                        }
                    });
                }
            };
        }
    };
    LbState.TEST_URIS_PARTITIONDESCRIPTIONS.putAll(partitionDescriptions);
    LbState.TEST_SERVICE_BACKUP_REQUEST_PROPERTIES.add(createBackupRequestsConfiguration(5, "get"));
    LbState.refreshDefaultProperties();
    LoadBalancer loadBalancer = new SimpleLoadBalancer(LbState, _executor);
    DynamicClient dynamicClient = new DynamicClient(loadBalancer, null);
    return new BackupRequestsClient(dynamicClient, loadBalancer, _executor, null, 10, TimeUnit.SECONDS, isD2Async) {

        @Override
        Optional<TrackingBackupRequestsStrategy> getStrategyAfterUpdate(final String serviceName, final String operation) {
            // constantly enable backup request after backupDelayNano time.
            BackupRequestsStrategy alwaysBackup = new TestTrackingBackupRequestsStrategy.MockBackupRequestsStrategy(() -> Optional.of((long) backupDelayNano), () -> true);
            return Optional.of(new TrackingBackupRequestsStrategy(alwaysBackup));
        }
    };
}

Example 62 with RestResponseBuilder

use of com.linkedin.r2.message.rest.RestResponseBuilder in project rest.li by linkedin.

the class Messages method toRestResponse.

/**
 * Converts a StreamResponse to RestResponse
 * @param streamResponse the stream request to be converted
 * @param callback the callback to be invoked when the rest response is constructed
 * @param addContentLengthHeader whether the rest response should have content-length header
 */
public static void toRestResponse(StreamResponse streamResponse, final Callback<RestResponse> callback, final boolean addContentLengthHeader) {
    final RestResponseBuilder builder = new RestResponseBuilder(streamResponse);
    Callback<ByteString> assemblyCallback = new Callback<ByteString>() {

        @Override
        public void onError(Throwable e) {
            callback.onError(e);
        }

        @Override
        public void onSuccess(ByteString result) {
            if (addContentLengthHeader) {
                builder.setHeader(HttpConstants.CONTENT_LENGTH, String.valueOf(result.length()));
            }
            RestResponse restResponse = builder.setEntity(result).build();
            callback.onSuccess(restResponse);
        }
    };
    streamResponse.getEntityStream().setReader(new FullEntityReader(assemblyCallback));
}

Example 63 with RestResponseBuilder

use of com.linkedin.r2.message.rest.RestResponseBuilder in project rest.li by linkedin.

the class TestMessages method testToRestTransportCallbackRestException.

@Test
public void testToRestTransportCallbackRestException() {
    TransportCallback<StreamResponse> streamCallback = response -> {
        Assert.assertTrue(response.hasError());
        Assert.assertNotNull(response.getError());
        Assert.assertTrue(response.getError() instanceof StreamException);
        Assert.assertNotNull(response.getWireAttributes());
        Assert.assertEquals(response.getWireAttributes(), WIRE_ATTR);
        Assert.assertEquals(response.getError().getStackTrace().length, 0);
    };
    TransportCallback<RestResponse> restCallback = Messages.toRestTransportCallback(streamCallback);
    RestResponseBuilder builder = new RestResponseBuilder();
    builder.setEntity(DATA);
    RestResponse restResponse = builder.build();
    restCallback.onResponse(TransportResponseImpl.error(new RestException(restResponse, new IllegalStateException()), WIRE_ATTR));
}

Example 64 with RestResponseBuilder

use of com.linkedin.r2.message.rest.RestResponseBuilder in project rest.li by linkedin.

the class TestMessages method testToRestTransportCallbackSuccess.

@Test
public void testToRestTransportCallbackSuccess() {
    TransportCallback<StreamResponse> streamCallback = response -> {
        Assert.assertFalse(response.hasError());
        Assert.assertNotNull(response.getResponse());
        response.getResponse().getEntityStream().setReader(ENTITY_VERIFIER);
        Assert.assertNotNull(response.getWireAttributes());
        Assert.assertEquals(response.getWireAttributes(), WIRE_ATTR);
    };
    TransportCallback<RestResponse> restCallback = Messages.toRestTransportCallback(streamCallback);
    RestResponseBuilder builder = new RestResponseBuilder();
    builder.setEntity(DATA);
    RestResponse restResponse = builder.build();
    restCallback.onResponse(TransportResponseImpl.success(restResponse, WIRE_ATTR));
}

Example 65 with RestResponseBuilder

use of com.linkedin.r2.message.rest.RestResponseBuilder in project rest.li by linkedin.

the class TestCapRepFilter method testRestException.

@Test
public void testRestException() throws IOException {
    Path dirPath = Files.createTempDirectory("caprep-test");
    CaptureLastCallFilter lastCallFilter = new CaptureLastCallFilter();
    FilterChain fc = FilterChains.createRestChain(lastCallFilter, _filter);
    RestRequest myRequest = new RestRequestBuilder(URI.create("/req1")).setEntity("123".getBytes()).build();
    RestResponse myErrorResponse = new RestResponseBuilder().setStatus(400).setEntity("321".getBytes()).build();
    RestException myRestException = new RestException(myErrorResponse);
    _filter.capture(dirPath.toString());
    RequestContext requestContext = new RequestContext();
    FilterUtil.fireRestRequest(fc, myRequest, requestContext, FilterUtil.emptyWireAttrs());
    FilterUtil.fireRestError(fc, myRestException, requestContext, FilterUtil.emptyWireAttrs());
    lastCallFilter.reset();
    _filter.replay(dirPath.toString());
    FilterUtil.fireSimpleRestRequest(fc);
    Assert.assertNull(lastCallFilter.getLastErr());
    FilterUtil.fireRestRequest(fc, myRequest);
    Assert.assertTrue(lastCallFilter.getLastErr() instanceof RestException);
    Assert.assertEquals(((RestException) lastCallFilter.getLastErr()).getResponse(), myErrorResponse);
}