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

Example 56 with StreamRequest

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

the class TestMIMEChainingAlternate method testAlternateSinglePartDataSource.

// This test has the server alternate between consuming a part and sending a part as a data source
// to a writer.
// Since we have four parts, the server will consume the 2nd and 4th and send out the 1st and 3rd.
// To make the test easier we will have two callbacks to send to the server to indicate
// the presence of each data source.
// This violates the typical client/server http pattern, but accomplishes the purpose of this test
// and it makes it easier to write.
@Test(dataProvider = "chunkSizes")
public void testAlternateSinglePartDataSource(final int chunkSize) throws Exception {
    final List<MultiPartMIMEDataSourceWriter> dataSources = generateInputStreamDataSources(chunkSize, _scheduledExecutorService);
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder().appendDataSources(dataSources).build();
    final StreamRequest streamRequest = mock(StreamRequest.class);
    when(streamRequest.getEntityStream()).thenReturn(writer.getEntityStream());
    final String contentTypeHeader = "multipart/mixed; boundary=" + writer.getBoundary();
    when(streamRequest.getHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER)).thenReturn(contentTypeHeader);
    // Client side preparation to read the part back on the callback.
    // We have two callbacks here since we will get two responses.
    // Note the chunks size will carry over since the client is controlling how much data he gets back
    // based on the chunk size he writes.
    MIMETestUtils.MultiPartMIMEFullReaderCallback clientReceiverA = new MIMETestUtils.MultiPartMIMEFullReaderCallback();
    MIMETestUtils.MultiPartMIMEFullReaderCallback clientReceiverB = new MIMETestUtils.MultiPartMIMEFullReaderCallback();
    Callback<StreamResponse> callbackA = generateSuccessChainCallback(clientReceiverA);
    Callback<StreamResponse> callbackB = generateSuccessChainCallback(clientReceiverB);
    // Server side start
    MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    final CountDownLatch latch = new CountDownLatch(1);
    ServerMultiPartMIMEAlternatorCallback serverSender = new ServerMultiPartMIMEAlternatorCallback(latch, callbackA, callbackB);
    reader.registerReaderCallback(serverSender);
    latch.await(_testTimeout, TimeUnit.MILLISECONDS);
    // Verify client
    Assert.assertEquals(clientReceiverA.getSinglePartMIMEReaderCallbacks().size(), 1);
    Assert.assertEquals(clientReceiverA.getSinglePartMIMEReaderCallbacks().get(0).getFinishedData(), BODY_A.getPartData());
    Assert.assertEquals(clientReceiverA.getSinglePartMIMEReaderCallbacks().get(0).getHeaders(), BODY_A.getPartHeaders());
    Assert.assertEquals(clientReceiverB.getSinglePartMIMEReaderCallbacks().size(), 1);
    Assert.assertEquals(clientReceiverB.getSinglePartMIMEReaderCallbacks().get(0).getFinishedData(), BODY_C.getPartData());
    Assert.assertEquals(clientReceiverB.getSinglePartMIMEReaderCallbacks().get(0).getHeaders(), BODY_C.getPartHeaders());
    // Verify server
    Assert.assertEquals(serverSender.getSinglePartMIMEReaderCallbacks().size(), 2);
    Assert.assertEquals(serverSender.getSinglePartMIMEReaderCallbacks().get(0).getFinishedData(), BODY_B.getPartData());
    Assert.assertEquals(serverSender.getSinglePartMIMEReaderCallbacks().get(0).getHeaders(), BODY_B.getPartHeaders());
    Assert.assertEquals(serverSender.getSinglePartMIMEReaderCallbacks().get(1).getFinishedData(), BODY_D.getPartData());
    Assert.assertEquals(serverSender.getSinglePartMIMEReaderCallbacks().get(1).getHeaders(), BODY_D.getPartHeaders());
}

Example 57 with StreamRequest

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

the class TestChannelPoolBehavior method testChannelReuse.

@Test
public void testChannelReuse() throws Exception {
    _client2.streamRequest(new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, NOT_FOUND_URI)).build(EntityStreams.newEntityStream(new SlowWriter())), new Callback<StreamResponse>() {

        @Override
        public void onError(Throwable e) {
            if (e instanceof StreamException) {
                StreamException streamException = (StreamException) e;
                streamException.getResponse().getEntityStream().setReader(new CancelingReader());
            }
            throw new RuntimeException(e);
        }

        @Override
        public void onSuccess(StreamResponse result) {
            result.getEntityStream().setReader(new DrainReader());
        }
    });
    Future<RestResponse> responseFuture = _client2.restRequest(new RestRequestBuilder(Bootstrap.createHttpURI(PORT, NORMAL_URI)).build());
    RestResponse response = responseFuture.get(WRITER_DELAY * 1000, TimeUnit.MILLISECONDS);
    Assert.assertEquals(response.getStatus(), RestStatus.OK);
}

Example 58 with StreamRequest

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

the class TestCompressionEcho method testResponseCompression.

@Test(dataProvider = "compressionEchoData")
public void testResponseCompression(Client client, long bytes) throws InterruptedException, TimeoutException, ExecutionException {
    StreamRequestBuilder builder = new StreamRequestBuilder((_clientProvider.createHttpURI(_port, ECHO_URI)));
    BytesWriter writer = new BytesWriter(bytes, BYTE);
    StreamRequest request = builder.build(EntityStreams.newEntityStream(writer));
    // add operation to enable sending accept encoding
    RequestContext requestContext = new RequestContext();
    requestContext.putLocalAttr(R2Constants.OPERATION, "get");
    final FutureCallback<StreamResponse> callback = new FutureCallback<>();
    client.streamRequest(request, requestContext, callback);
    final StreamResponse response = callback.get(60, TimeUnit.SECONDS);
    Assert.assertEquals(response.getStatus(), RestStatus.OK);
    final FutureCallback<None> readerCallback = new FutureCallback<>();
    final BytesReader reader = new BytesReader(BYTE, readerCallback);
    response.getEntityStream().setReader(reader);
    readerCallback.get(60, TimeUnit.SECONDS);
    Assert.assertEquals(reader.getTotalBytes(), bytes);
    Assert.assertTrue(reader.allBytesCorrect());
}

Example 59 with StreamRequest

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

the class TestClientStreamCompressionFilter method testRequestCompressionRules.

@Test(dataProvider = "requestData")
public void testRequestCompressionRules(CompressionConfig requestCompressionConfig, CompressionOption requestCompressionOverride, boolean headerShouldBePresent, String operation) throws CompressionException, URISyntaxException, InterruptedException, ExecutionException, TimeoutException {
    Executor executor = Executors.newCachedThreadPool();
    ClientStreamCompressionFilter clientCompressionFilter = new ClientStreamCompressionFilter(StreamEncodingType.GZIP.getHttpName(), requestCompressionConfig, ACCEPT_COMPRESSIONS, new CompressionConfig(Integer.MAX_VALUE), Arrays.asList(ClientCompressionHelper.COMPRESS_ALL_RESPONSES_INDICATOR), executor);
    // The entity should be compressible for this test.
    int original = 100;
    byte[] entity = new byte[original];
    Arrays.fill(entity, (byte) 'A');
    StreamRequest streamRequest = new StreamRequestBuilder(new URI(URI)).setMethod(RestMethod.POST).build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(entity))));
    int compressed = EncodingType.GZIP.getCompressor().deflate(new ByteArrayInputStream(entity)).length;
    RequestContext context = new RequestContext();
    if (operation != null) {
        context.putLocalAttr(R2Constants.OPERATION, operation);
    }
    context.putLocalAttr(R2Constants.REQUEST_COMPRESSION_OVERRIDE, requestCompressionOverride);
    int entityLength = headerShouldBePresent ? compressed : original;
    FutureCallback<ByteString> callback = new FutureCallback<>();
    FullEntityReader reader = new FullEntityReader(callback);
    HeaderCaptureFilter captureFilter = new HeaderCaptureFilter(HttpConstants.CONTENT_ENCODING, headerShouldBePresent, entityLength, reader);
    clientCompressionFilter.onStreamRequest(streamRequest, context, Collections.<String, String>emptyMap(), captureFilter);
    ByteString entityRead = callback.get(10, TimeUnit.SECONDS);
    Assert.assertEquals(entityRead.length(), entityLength);
}

Example 60 with StreamRequest

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

the class TestClientStreamCompressionFilter method testCompressionOperations.

@Test(dataProvider = "operationsData")
public void testCompressionOperations(String compressionConfig, String[] operations, boolean headerShouldBePresent) throws URISyntaxException {
    StreamRequest streamRequest = new StreamRequestBuilder(new URI(URI)).build(EntityStreams.emptyStream());
    ClientStreamCompressionFilter clientCompressionFilter = new ClientStreamCompressionFilter(StreamEncodingType.IDENTITY.getHttpName(), new CompressionConfig(Integer.MAX_VALUE), ACCEPT_COMPRESSIONS, new CompressionConfig(Integer.MAX_VALUE), Arrays.asList(compressionConfig.split(",")), Executors.newCachedThreadPool());
    for (String operation : operations) {
        RequestContext context = new RequestContext();
        context.putLocalAttr(R2Constants.OPERATION, operation);
        clientCompressionFilter.onStreamRequest(streamRequest, context, Collections.<String, String>emptyMap(), new HeaderCaptureFilter(HttpConstants.ACCEPT_ENCODING, headerShouldBePresent, null));
    }
}