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

Example 16 with StreamRequestBuilder

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

the class TestBackupRequestsClient method testStreamRequestWithNoIsFullRequest.

@Test(invocationCount = 3, dataProvider = "isD2Async")
public void testStreamRequestWithNoIsFullRequest(boolean isD2Async) throws Exception {
    // 1s till response comes back
    int responseDelayNano = 100000000;
    // make backup request after 0.5 second
    int backupDelayNano = 50000000;
    Deque<URI> hostsReceivingRequest = new ConcurrentLinkedDeque<>();
    BackupRequestsClient client = createAlwaysBackupClientWithHosts(Arrays.asList("http://test1.com:123", "http://test2.com:123"), hostsReceivingRequest, responseDelayNano, backupDelayNano, isD2Async);
    URI uri = URI.create("d2://testService");
    // if there is no IS_FULL_REQUEST set, backup requests will not happen
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.newEntityStream(new ByteStringWriter(CONTENT)));
    RequestContext context = new RequestContext();
    context.putLocalAttr(R2Constants.OPERATION, "get");
    RequestContext context1 = context.clone();
    CountDownLatch latch = new CountDownLatch(1);
    AtomicReference<AssertionError> failure = new AtomicReference<>();
    client.streamRequest(streamRequest, context1, new Callback<StreamResponse>() {

        @Override
        public void onError(Throwable e) {
            failure.set(new AssertionError("Callback onError"));
            latch.countDown();
        }

        @Override
        public void onSuccess(StreamResponse result) {
            try {
                assertEquals(result.getStatus(), 200);
                assertEquals(result.getHeader("buffered"), "false");
                assertEquals(hostsReceivingRequest.size(), 1);
                assertEquals(new HashSet<>(hostsReceivingRequest).size(), 1);
                hostsReceivingRequest.clear();
            } catch (AssertionError e) {
                failure.set(e);
            }
            latch.countDown();
        }
    });
    latch.await(2, TimeUnit.SECONDS);
    if (failure.get() != null) {
        throw failure.get();
    }
}

Example 17 with StreamRequestBuilder

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

the class TestBackupRequestsClient method testStreamRequestWithIsFullRequest.

@Test(invocationCount = 3, dataProvider = "isD2Async")
public void testStreamRequestWithIsFullRequest(boolean isD2Async) throws Exception {
    // 5s till response comes back
    int responseDelayNano = 500000000;
    // make backup request after 1 second
    int backupDelayNano = 100000000;
    Deque<URI> hostsReceivingRequest = new ConcurrentLinkedDeque<>();
    BackupRequestsClient client = createAlwaysBackupClientWithHosts(Arrays.asList("http://test1.com:123", "http://test2.com:123"), hostsReceivingRequest, responseDelayNano, backupDelayNano, isD2Async);
    URI uri = URI.create("d2://testService");
    // if there is IS_FULL_REQUEST set, backup requests will happen
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.newEntityStream(new ByteStringWriter(CONTENT)));
    RequestContext context = new RequestContext();
    context.putLocalAttr(R2Constants.OPERATION, "get");
    context.putLocalAttr(R2Constants.IS_FULL_REQUEST, true);
    RequestContext context1 = context.clone();
    CountDownLatch latch = new CountDownLatch(1);
    AtomicReference<AssertionError> failure = new AtomicReference<>();
    client.streamRequest(streamRequest, context1, new Callback<StreamResponse>() {

        @Override
        public void onError(Throwable e) {
            failure.set(new AssertionError("Callback onError"));
            latch.countDown();
        }

        @Override
        public void onSuccess(StreamResponse result) {
            try {
                assertEquals(result.getStatus(), 200);
                assertEquals(result.getHeader("buffered"), "true");
                assertEquals(hostsReceivingRequest.size(), 2);
                assertEquals(new HashSet<>(hostsReceivingRequest).size(), 2);
                hostsReceivingRequest.clear();
            } catch (AssertionError e) {
                failure.set(e);
            }
            latch.countDown();
        }
    });
    latch.await(6, TimeUnit.SECONDS);
    if (failure.get() != null) {
        throw failure.get();
    }
}

Example 18 with StreamRequestBuilder

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

Example 19 with StreamRequestBuilder

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

the class TestReplaceableFilter method testStreamingPassThrough.

@Test
public void testStreamingPassThrough() {
    ReplaceableFilter filter = getFilter();
    CountRequestChunksFilter countRequestChunksFilter = new CountRequestChunksFilter();
    CountResponseChunksFilter countResponseChunksFilter = new CountResponseChunksFilter();
    FilterChain filterChain = FilterChains.empty().addLast(countResponseChunksFilter).addLast(filter).addLast(countRequestChunksFilter);
    int requestChunks = 100;
    int responseChunks = 200;
    StreamRequest streamRequest = new StreamRequestBuilder(URI.create("/test")).build(EntityStreams.newEntityStream(new Writer() {

        private WriteHandle _wh;

        int _count = 0;

        @Override
        public void onInit(WriteHandle wh) {
            _wh = wh;
        }

        @Override
        public void onWritePossible() {
            while (_wh.remaining() > 0) {
                if (_count < requestChunks) {
                    _wh.write(ByteString.copy(new byte[10]));
                    _count++;
                } else {
                    _wh.done();
                }
            }
        }

        @Override
        public void onAbort(Throwable e) {
        }
    }));
    StreamResponse streamResponse = new StreamResponseBuilder().build(EntityStreams.newEntityStream(new Writer() {

        private WriteHandle _wh;

        int _count = 0;

        @Override
        public void onInit(WriteHandle wh) {
            _wh = wh;
        }

        @Override
        public void onWritePossible() {
            while (_wh.remaining() > 0) {
                if (_count < responseChunks) {
                    _wh.write(ByteString.copy(new byte[10]));
                    _count++;
                } else {
                    _wh.done();
                }
            }
        }

        @Override
        public void onAbort(Throwable e) {
        }
    }));
    FilterUtil.fireStreamRequest(filterChain, streamRequest);
    FilterUtil.fireStreamResponse(filterChain, streamResponse, new RequestContext(), new HashMap<>());
    Assert.assertEquals(countRequestChunksFilter.getRequestChunks(), requestChunks);
    Assert.assertEquals(countResponseChunksFilter.getResponseChunks(), responseChunks);
}

Example 20 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 {
    final long totalBytes = SMALL_BYTES_NUM;
    EntityStream entityStream = EntityStreams.newEntityStream(new BytesWriter(totalBytes, BYTE));
    StreamRequestBuilder builder = new StreamRequestBuilder(_clientProvider.createHttpURI(_port, ERROR_RECEIVER_URI));
    StreamRequest request = builder.setMethod("POST").build(entityStream);
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicReference<Throwable> error = new AtomicReference<>();
    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());
}