Examples with StreamRequestBuilder com.linkedin.r2.message.stream.StreamRequestBuilder used on opensource projects

Example 11 with StreamRequestBuilder

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

the class RestClient method buildStreamRequest.

private StreamRequest buildStreamRequest(URI uri, ResourceMethod method, DataMap dataMap, Map<String, String> headers, List<String> cookies, ProtocolVersion protocolVersion, ContentType contentType, List<AcceptType> acceptTypes, boolean acceptResponseAttachments, List<Object> streamingAttachments) throws Exception {
    StreamRequestBuilder requestBuilder = new StreamRequestBuilder(uri).setMethod(method.getHttpMethod().toString());
    requestBuilder.setHeaders(headers);
    requestBuilder.setCookies(cookies);
    addAcceptHeaders(requestBuilder, acceptTypes, acceptResponseAttachments);
    addProtocolVersionHeader(requestBuilder, protocolVersion);
    if (method.getHttpMethod() == HttpMethod.POST) {
        requestBuilder.setHeader(RestConstants.HEADER_RESTLI_REQUEST_METHOD, method.toString());
    }
    //This request builders enforce this invariant.
    if (streamingAttachments != null) {
        final ByteStringWriter firstPartWriter;
        final ContentType type = resolveContentType(requestBuilder, dataMap, contentType);
        //with empty action parameters will have an empty JSON ({}) as the body.
        assert (type != null);
        switch(type) {
            case PSON:
                firstPartWriter = new ByteStringWriter(ByteString.copy(PSON_DATA_CODEC.mapToBytes(dataMap)));
                break;
            case JSON:
                firstPartWriter = new ByteStringWriter(ByteString.copy(JACKSON_DATA_CODEC.mapToBytes(dataMap)));
                break;
            default:
                throw new IllegalStateException("Unknown ContentType:" + type);
        }
        //Our protocol does not use an epilogue or a preamble.
        final MultiPartMIMEWriter.Builder attachmentsBuilder = new MultiPartMIMEWriter.Builder();
        for (final Object dataSource : streamingAttachments) {
            assert (dataSource instanceof RestLiAttachmentDataSourceWriter || dataSource instanceof RestLiDataSourceIterator);
            if (dataSource instanceof RestLiAttachmentDataSourceWriter) {
                AttachmentUtils.appendSingleAttachmentToBuilder(attachmentsBuilder, (RestLiAttachmentDataSourceWriter) dataSource);
            } else {
                AttachmentUtils.appendMultipleAttachmentsToBuilder(attachmentsBuilder, (RestLiDataSourceIterator) dataSource);
            }
        }
        final MultiPartMIMEWriter multiPartMIMEWriter = AttachmentUtils.createMultiPartMIMEWriter(firstPartWriter, type.getHeaderKey(), attachmentsBuilder);
        final String contentTypeHeader = MultiPartMIMEUtils.buildMIMEContentTypeHeader(AttachmentUtils.RESTLI_MULTIPART_SUBTYPE, multiPartMIMEWriter.getBoundary(), Collections.<String, String>emptyMap());
        requestBuilder.setHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER, contentTypeHeader);
        return requestBuilder.build(multiPartMIMEWriter.getEntityStream());
    } else {
        //taken the RestRequest code path.
        assert (acceptResponseAttachments == true);
        return Messages.toStreamRequest(buildRestRequest(uri, method, dataMap, headers, cookies, protocolVersion, contentType, acceptTypes, acceptResponseAttachments));
    }
}

Example 12 with StreamRequestBuilder

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

the class DegraderTrackerClientTest method testCallTrackingStreamRequest.

@Test
public void testCallTrackingStreamRequest() throws Exception {
    URI uri = URI.create("http://test.qa.com:1234/foo");
    SettableClock clock = new SettableClock();
    AtomicInteger action = new AtomicInteger(0);
    TransportClient tc = new TransportClient() {

        @Override
        public void restRequest(RestRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<RestResponse> callback) {
        }

        @Override
        public void streamRequest(StreamRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
            clock.addDuration(5);
            switch(action.get()) {
                // success
                case 0:
                    callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().build(EntityStreams.emptyStream())));
                    break;
                // fail with stream exception
                case 1:
                    callback.onResponse(TransportResponseImpl.error(new StreamException(new StreamResponseBuilder().setStatus(500).build(EntityStreams.emptyStream()))));
                    break;
                // fail with timeout exception
                case 2:
                    callback.onResponse(TransportResponseImpl.error(new RemoteInvocationException(new TimeoutException())));
                    break;
                // fail with other exception
                default:
                    callback.onResponse(TransportResponseImpl.error(new RuntimeException()));
                    break;
            }
        }

        @Override
        public void shutdown(Callback<None> callback) {
        }
    };
    DegraderTrackerClientImpl client = (DegraderTrackerClientImpl) createTrackerClient(tc, clock, uri);
    CallTracker callTracker = client.getCallTracker();
    CallTracker.CallStats stats;
    DegraderControl degraderControl = client.getDegraderControl(DefaultPartitionAccessor.DEFAULT_PARTITION_ID);
    DelayConsumeCallback delayConsumeCallback = new DelayConsumeCallback();
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), delayConsumeCallback);
    clock.addDuration(5);
    // we only recorded the time when stream response arrives, but callcompletion.endcall hasn't been called yet.
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 0);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 0);
    // delay
    clock.addDuration(100);
    delayConsumeCallback.consume();
    clock.addDuration(5000);
    // now that we consumed the entity stream, callcompletion.endcall has been called.
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 0);
    Assert.assertEquals(stats.getCallCountTotal(), 1);
    Assert.assertEquals(stats.getErrorCountTotal(), 0);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.0, 0.001);
    action.set(1);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), delayConsumeCallback);
    clock.addDuration(5);
    // we endcall with error immediately for stream exception, even before the entity is consumed
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 2);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 1);
    delayConsumeCallback.consume();
    clock.addDuration(5000);
    // no change in tracking after entity is consumed
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 2);
    Assert.assertEquals(stats.getErrorCountTotal(), 1);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.2, 0.001);
    action.set(2);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), new TestTransportCallback<>());
    clock.addDuration(5);
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 3);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 2);
    clock.addDuration(5000);
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 3);
    Assert.assertEquals(stats.getErrorCountTotal(), 2);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.4, 0.001);
    action.set(3);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), new TestTransportCallback<>());
    clock.addDuration(5);
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 4);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 3);
    clock.addDuration(5000);
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 4);
    Assert.assertEquals(stats.getErrorCountTotal(), 3);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.2, 0.001);
}

Example 13 with StreamRequestBuilder

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

the class RetryClientTest method testStreamRetryOverLimit.

@Test
public void testStreamRetryOverLimit() throws Exception {
    SimpleLoadBalancer balancer = prepareLoadBalancer(Arrays.asList("http://test.linkedin.com/retry1", "http://test.linkedin.com/retry2"), HttpClientFactory.UNLIMITED_CLIENT_REQUEST_RETRY_RATIO);
    DynamicClient dynamicClient = new DynamicClient(balancer, null);
    RetryClient client = new RetryClient(dynamicClient, balancer, 1, RetryClient.DEFAULT_UPDATE_INTERVAL_MS, RetryClient.DEFAULT_AGGREGATED_INTERVAL_NUM, SystemClock.instance(), true, true);
    URI uri = URI.create("d2://retryService?arg1=empty&arg2=empty");
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.emptyStream());
    DegraderTrackerClientTest.TestCallback<StreamResponse> streamCallback = new DegraderTrackerClientTest.TestCallback<>();
    client.streamRequest(streamRequest, streamCallback);
    assertNull(streamCallback.t);
    assertNotNull(streamCallback.e);
    assertTrue(streamCallback.e.getMessage().contains("Data not available"));
}

Example 14 with StreamRequestBuilder

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

the class RetryClientTest method testStreamRetryNoAvailableHosts.

@Test
public void testStreamRetryNoAvailableHosts() throws Exception {
    SimpleLoadBalancer balancer = prepareLoadBalancer(Arrays.asList("http://test.linkedin.com/retry1", "http://test.linkedin.com/retry2"), HttpClientFactory.UNLIMITED_CLIENT_REQUEST_RETRY_RATIO);
    DynamicClient dynamicClient = new DynamicClient(balancer, null);
    RetryClient client = new RetryClient(dynamicClient, balancer, D2ClientConfig.DEFAULT_RETRY_LIMIT, RetryClient.DEFAULT_UPDATE_INTERVAL_MS, RetryClient.DEFAULT_AGGREGATED_INTERVAL_NUM, SystemClock.instance(), true, true);
    URI uri = URI.create("d2://retryService?arg1=empty&arg2=empty");
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.emptyStream());
    FutureCallback<StreamResponse> streamCallback = new FutureCallback<>();
    client.streamRequest(streamRequest, streamCallback);
    try {
        streamCallback.get();
    } catch (ExecutionException e) {
        assertTrue(e.toString().contains("retryService is in a bad state"), e.getMessage());
    }
}

Example 15 with StreamRequestBuilder

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

the class RetryClientTest method testStreamRetry.

@Test
public void testStreamRetry() throws Exception {
    SimpleLoadBalancer balancer = prepareLoadBalancer(Arrays.asList("http://test.linkedin.com/retry1", "http://test.linkedin.com/good"), HttpClientFactory.UNLIMITED_CLIENT_REQUEST_RETRY_RATIO);
    DynamicClient dynamicClient = new DynamicClient(balancer, null);
    RetryClient client = new RetryClient(dynamicClient, balancer, D2ClientConfig.DEFAULT_RETRY_LIMIT, RetryClient.DEFAULT_UPDATE_INTERVAL_MS, RetryClient.DEFAULT_AGGREGATED_INTERVAL_NUM, SystemClock.instance(), true, true);
    URI uri = URI.create("d2://retryService?arg1arg2");
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.newEntityStream(new ByteStringWriter(CONTENT)));
    DegraderTrackerClientTest.TestCallback<StreamResponse> restCallback = new DegraderTrackerClientTest.TestCallback<>();
    client.streamRequest(streamRequest, restCallback);
    assertNull(restCallback.e);
    assertNotNull(restCallback.t);
}