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

Example 21 with StreamRequest

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

the class TestRestClientRequestBuilder method testAction.

@Test(dataProvider = "data")
public void testAction(RestClient.ContentType contentType, String expectedContentTypeHeader, String expectedRequestBody, String expectedEntitiesBody, List<RestClient.AcceptType> acceptTypes, String expectedAcceptHeader, boolean acceptContentTypePerClient, boolean streamAttachments, boolean acceptResponseAttachments) throws URISyntaxException {
    //its a RestRequest.
    if (streamAttachments == false && acceptResponseAttachments == false) {
        //RestRequest with a request entity
        RestRequest restRequest = clientGeneratedRestRequest(ActionRequest.class, ResourceMethod.ACTION, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient);
        Assert.assertEquals(restRequest.getHeader(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        Assert.assertEquals(restRequest.getEntity().asAvroString(), expectedRequestBody);
        Assert.assertEquals(restRequest.getHeader(ACCEPT_TYPE_HEADER), expectedAcceptHeader);
        //RestRequest without a request entity
        RestRequest restRequestNoEntity = clientGeneratedRestRequest(ActionRequest.class, ResourceMethod.ACTION, new DataMap(), contentType, acceptTypes, acceptContentTypePerClient);
        Assert.assertEquals(restRequest.getHeader(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        //Verify that that there is an empty payload based on the expected content type
        if (expectedContentTypeHeader.equalsIgnoreCase(RestConstants.HEADER_VALUE_APPLICATION_PSON)) {
            Assert.assertEquals(restRequestNoEntity.getEntity().asAvroString(), SERIALIZED_EMPTY_PSON);
        } else {
            Assert.assertEquals(restRequestNoEntity.getEntity().asAvroString(), SERIALIZED_EMPTY_JSON);
        }
        Assert.assertEquals(restRequest.getHeader(ACCEPT_TYPE_HEADER), expectedAcceptHeader);
    } else {
        //StreamRequest with a request entity
        StreamRequest streamRequest = clientGeneratedStreamRequest(ActionRequest.class, ResourceMethod.ACTION, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient, streamAttachments ? generateRequestAttachments() : null, acceptResponseAttachments);
        verifyStreamRequest(streamRequest, acceptResponseAttachments, expectedAcceptHeader, streamAttachments, expectedContentTypeHeader, expectedRequestBody);
        //StreamRequest without a request entity
        StreamRequest streamRequestNoEntity = clientGeneratedStreamRequest(ActionRequest.class, ResourceMethod.ACTION, new DataMap(), contentType, acceptTypes, acceptContentTypePerClient, streamAttachments ? generateRequestAttachments() : null, acceptResponseAttachments);
        //Verify that that there is an empty payload based on the expected content type
        if (expectedContentTypeHeader.equalsIgnoreCase(RestConstants.HEADER_VALUE_APPLICATION_PSON)) {
            verifyStreamRequest(streamRequestNoEntity, acceptResponseAttachments, expectedAcceptHeader, streamAttachments, expectedContentTypeHeader, SERIALIZED_EMPTY_PSON);
        } else {
            verifyStreamRequest(streamRequestNoEntity, acceptResponseAttachments, expectedAcceptHeader, streamAttachments, expectedContentTypeHeader, SERIALIZED_EMPTY_JSON);
        }
    }
}

Example 22 with StreamRequest

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

the class BackupRequestsClient method streamRequest.

@Override
public void streamRequest(StreamRequest request, RequestContext requestContext, Callback<StreamResponse> callback) {
    // Currently only support backup requests with IS_FULL_REQUEST.
    if (!isFullRequest(requestContext)) {
        _d2Client.streamRequest(request, requestContext, callback);
        return;
    }
    if (!isBuffered(requestContext)) {
        final FullEntityObserver observer = new FullEntityObserver(new Callback<ByteString>() {

            @Override
            public void onError(Throwable e) {
                LOG.warn("Failed to record request's entity for retrying backup request.");
            }

            @Override
            public void onSuccess(ByteString result) {
                requestContext.putLocalAttr(R2Constants.BACKUP_REQUEST_BUFFERED_BODY, result);
            }
        });
        request.getEntityStream().addObserver(observer);
    }
    if (_isD2Async) {
        requestAsync(request, requestContext, _d2Client::streamRequest, callback);
        return;
    }
    _d2Client.streamRequest(request, requestContext, decorateCallbackSync(request, requestContext, _d2Client::streamRequest, callback));
}

Example 23 with StreamRequest

use of com.linkedin.r2.message.stream.StreamRequest 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 24 with StreamRequest

use of com.linkedin.r2.message.stream.StreamRequest 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 25 with StreamRequest

use of com.linkedin.r2.message.stream.StreamRequest 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());
    }
}