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

Example 1 with StreamRequestBuilder

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

the class RewriteClientTestStreamRequest 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();
    RewriteClient client = new RewriteClient(serviceName, uri, wrappedClient);
    assertEquals(client.getUri(), uri);
    assertEquals(client.getServiceName(), serviceName);
    assertEquals(client.getWrappedClient(), wrappedClient);
    StreamRequest streamRequest = new StreamRequestBuilder(URI.create("d2://HistoryService/getCube")).build(EntityStreams.emptyStream());
    Map<String, String> restWireAttrs = new HashMap<String, String>();
    TestTransportCallback<StreamResponse> restCallback = new TestTransportCallback<StreamResponse>();
    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);
}

Example 2 with StreamRequestBuilder

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

the class TrackerClientTest 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) {
        }
    };
    TrackerClient client = 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 3 with StreamRequestBuilder

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

the class TestNettyRequestAdapter method testStreamToNettyRequest.

@Test
public void testStreamToNettyRequest() throws Exception {
    StreamRequestBuilder streamRequestBuilder = new StreamRequestBuilder(new URI(ANY_URI));
    streamRequestBuilder.setMethod("POST");
    streamRequestBuilder.setHeader("Content-Length", Integer.toString(ANY_ENTITY.length()));
    streamRequestBuilder.setHeader("Content-Type", "application/json");
    streamRequestBuilder.setCookies(Collections.singletonList("anyCookie"));
    StreamRequest streamRequest = streamRequestBuilder.build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(ANY_ENTITY.getBytes()))));
    HttpRequest nettyRequest = NettyRequestAdapter.toNettyRequest(streamRequest);
    Assert.assertEquals(nettyRequest.uri(), "/foo/bar?q=baz");
    Assert.assertEquals(nettyRequest.method(), HttpMethod.POST);
    Assert.assertEquals(nettyRequest.protocolVersion(), HttpVersion.HTTP_1_1);
    Assert.assertNull(nettyRequest.headers().get("Content-Length"));
    Assert.assertEquals(nettyRequest.headers().get("Content-Type"), "application/json");
    Assert.assertEquals(nettyRequest.headers().get("Cookie"), ANY_COOKIE);
}

Example 4 with StreamRequestBuilder

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

the class TestNettyRequestAdapter method testStreamToNettyRequestContentLengthIgnoreCase.

@Test
public void testStreamToNettyRequestContentLengthIgnoreCase() throws Exception {
    StreamRequestBuilder streamRequestBuilder = new StreamRequestBuilder(new URI(ANY_URI));
    streamRequestBuilder.setHeader("CONTENT-LENGTH", Integer.toString(ANY_ENTITY.length()));
    StreamRequest streamRequest = streamRequestBuilder.build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(ANY_ENTITY.getBytes()))));
    HttpRequest nettyRequest = NettyRequestAdapter.toNettyRequest(streamRequest);
    Assert.assertNull(nettyRequest.headers().get("Content-Length"));
}

Example 5 with StreamRequestBuilder

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

the class TestStreamRequest method testErrorReceiver.

@Test
public void testErrorReceiver() throws Exception {
    for (Client client : clients()) {
        final long totalBytes = SMALL_BYTES_NUM;
        EntityStream entityStream = EntityStreams.newEntityStream(new BytesWriter(totalBytes, BYTE));
        StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, ERROR_RECEIVER_URI));
        StreamRequest request = builder.setMethod("POST").build(entityStream);
        final CountDownLatch latch = new CountDownLatch(1);
        final AtomicReference<Throwable> error = new AtomicReference<Throwable>();
        Callback<StreamResponse> callback = new Callback<StreamResponse>() {

            @Override
            public void onError(Throwable e) {
                error.set(e);
                latch.countDown();
            }

            @Override
            public void onSuccess(StreamResponse result) {
                latch.countDown();
            }
        };
        client.streamRequest(request, callback);
        latch.await();
        Assert.assertNotNull(error.get());
    }
}