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

Example 31 with Writer

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

the class TestMIMEIntegrationReaderWriter method testEachSingleBodyDataSource.

@Test(dataProvider = "eachSingleBodyDataSource", enabled = false)
public void testEachSingleBodyDataSource(final int chunkSize, final MIMEDataPart bodyPart) throws Exception {
    final MultiPartMIMEInputStream inputStreamDataSource = new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(bodyPart.getPartData().copyBytes()), scheduledExecutorService, bodyPart.getPartHeaders()).withWriteChunkSize(chunkSize).build();
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder("some preamble", "").appendDataSource(inputStreamDataSource).build();
    executeRequestAndAssert(writer, ImmutableList.of(bodyPart));
}

Example 32 with Writer

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

the class TestMIMEChainingReader method testMimeReaderDataSource.

//Verifies that a multi part mime reader can be used as a data source to the writer.
//To make the test easier to write, we simply chain back to the client in the form of simulating a response.
@Test(dataProvider = "chunkSizes")
public void testMimeReaderDataSource(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
    //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.
    final CountDownLatch latch = new CountDownLatch(1);
    ClientMultiPartMIMEReaderReceiverCallback _clientReceiver = new ClientMultiPartMIMEReaderReceiverCallback(latch);
    Callback<StreamResponse> callback = generateSuccessChainCallback(_clientReceiver);
    //Server side start
    MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    ServerMultiPartMIMEChainReaderWriterCallback _serverSender = new ServerMultiPartMIMEChainReaderWriterCallback(callback, reader);
    reader.registerReaderCallback(_serverSender);
    latch.await(_testTimeout, TimeUnit.MILLISECONDS);
    //Verify client. No need to verify the server.
    List<MIMETestUtils.SinglePartMIMEFullReaderCallback> singlePartMIMEReaderCallbacks = _clientReceiver.getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), 4);
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getFinishedData(), BODY_A.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getHeaders(), BODY_A.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getFinishedData(), BODY_B.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getHeaders(), BODY_B.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(2).getFinishedData(), BODY_C.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(2).getHeaders(), BODY_C.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(3).getFinishedData(), BODY_D.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(3).getHeaders(), BODY_D.getPartHeaders());
}

Example 33 with Writer

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

the class TestMIMEChainingMultipleSources method testSinglePartDataSource.

//Test breakdown:
//1. Main thread sends mime request
//2. Server A sends a simple POST request to Server B
//3. Server B sends back a mime response to Server A
//4. Server A takes the original incoming request from the main thread + a local input
//stream + the first part from the incoming mime response from Server B.
//5. Main thread then gets all of this and stores it.
//6. Server A then drains and stores the rest of the parts from Server B's response.
@Test(dataProvider = "chunkSizes", groups = TESTNG_GROUP_KNOWN_ISSUE)
public void testSinglePartDataSource(final int chunkSize) throws Exception {
    _chunkSize = chunkSize;
    final List<MultiPartMIMEDataSourceWriter> dataSources = generateInputStreamDataSources(chunkSize, _scheduledExecutorService);
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder().appendDataSources(dataSources).build();
    final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(Bootstrap.createHttpsURI(PORT_SERVER_A, SERVER_A_URI), "mixed", writer, Collections.<String, String>emptyMap());
    ClientMultiPartReceiver clientReceiver = new ClientMultiPartReceiver();
    Callback<StreamResponse> callback = generateSuccessChainCallback(clientReceiver);
    //Send the request to Server A
    _client.streamRequest(streamRequest, callback);
    _latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
    //Verify client
    List<MIMETestUtils.SinglePartMIMEFullReaderCallback> clientSinglePartCallbacks = clientReceiver.getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(clientReceiver.getSinglePartMIMEReaderCallbacks().size(), 6);
    Assert.assertEquals(clientSinglePartCallbacks.get(0).getFinishedData(), BODY_1.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(0).getHeaders(), BODY_1.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(1).getFinishedData(), BODY_5.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(1).getHeaders(), BODY_5.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(2).getFinishedData(), BODY_A.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(2).getHeaders(), BODY_A.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(3).getFinishedData(), BODY_B.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(3).getHeaders(), BODY_B.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(4).getFinishedData(), BODY_C.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(4).getHeaders(), BODY_C.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(5).getFinishedData(), BODY_D.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(5).getHeaders(), BODY_D.getPartHeaders());
    //Verify Server A
    List<MIMETestUtils.SinglePartMIMEFullReaderCallback> serverASinglePartCallbacks = _serverAMultiPartCallback.getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(serverASinglePartCallbacks.size(), 3);
    Assert.assertEquals(serverASinglePartCallbacks.get(0).getFinishedData(), BODY_2.getPartData());
    Assert.assertEquals(serverASinglePartCallbacks.get(0).getHeaders(), BODY_2.getPartHeaders());
    Assert.assertEquals(serverASinglePartCallbacks.get(1).getFinishedData(), BODY_3.getPartData());
    Assert.assertEquals(serverASinglePartCallbacks.get(1).getHeaders(), BODY_3.getPartHeaders());
    Assert.assertEquals(serverASinglePartCallbacks.get(2).getFinishedData(), BODY_4.getPartData());
    Assert.assertEquals(serverASinglePartCallbacks.get(2).getHeaders(), BODY_4.getPartHeaders());
}

Example 34 with Writer

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

the class TestEntityStream method testEntityStream.

@Test
public void testEntityStream() throws Exception {
    TestWriter writer = new TestWriter();
    ControlReader reader = new ControlReader();
    TestObserver ob1 = new TestObserver();
    TestObserver ob2 = new TestObserver();
    EntityStream stream = EntityStreams.newEntityStream(writer);
    stream.addObserver(ob1);
    stream.addObserver(ob2);
    // write is not possible without a reader
    Assert.assertEquals(writer.getWritePossibleCount(), 0);
    stream.setReader(reader);
    // write is not possible before reader reads
    Assert.assertEquals(writer.getWritePossibleCount(), 0);
    reader.read(1);
    // write become possible
    Assert.assertEquals(writer.getWritePossibleCount(), 1);
    Assert.assertEquals(writer.remaining(), 1);
    writer.write();
    Assert.assertEquals(writer.remaining(), 0);
    reader.read(10);
    // write again become possible
    Assert.assertEquals(writer.getWritePossibleCount(), 2);
    Assert.assertEquals(writer.remaining(), 10);
    while (writer.remaining() > 1) {
        writer.write();
    }
    Assert.assertEquals(writer.remaining(), 1);
    reader.read(10);
    // write hasn't become impossible when reader reads again, so onWritePossible should not have been invoked again
    Assert.assertEquals(writer.getWritePossibleCount(), 2);
    while (writer.remaining() > 0) {
        writer.write();
    }
    Assert.assertEquals(ob1.getChunkCount(), 21);
    Assert.assertEquals(ob2.getChunkCount(), 21);
    Assert.assertEquals(reader.getChunkCount(), 21);
    try {
        writer.write();
        Assert.fail("should fail with IllegalStateException");
    } catch (IllegalStateException ex) {
    // expected
    }
}

Example 35 with Writer

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

the class TestEntityStream method testRaceBetweenErrorAndCancel.

private void testRaceBetweenErrorAndCancel(ExecutorService executor) throws Exception {
    final CountDownLatch startLatch = new CountDownLatch(1);
    final CountDownLatch finishLatch = new CountDownLatch(2);
    final CountDownLatch prepareLatch = new CountDownLatch(2);
    final TestWriter writer = new TestWriter();
    TestObserver observer = new TestObserver();
    final ControlReader reader = new ControlReader();
    EntityStream es = EntityStreams.newEntityStream(writer);
    es.addObserver(observer);
    es.setReader(reader);
    reader.read(1000);
    executor.submit(new Callable<Object>() {

        @Override
        public Object call() throws Exception {
            while (writer.remaining() > 100) {
                writer.write();
            }
            prepareLatch.countDown();
            startLatch.await();
            writer.error(new RuntimeException("writer has problem"));
            finishLatch.countDown();
            return null;
        }
    });
    executor.submit(new Callable<Object>() {

        @Override
        public Object call() throws Exception {
            prepareLatch.countDown();
            startLatch.await();
            reader.cancel();
            finishLatch.countDown();
            return null;
        }
    });
    prepareLatch.await();
    startLatch.countDown();
    Assert.assertTrue(finishLatch.await(1000, TimeUnit.MILLISECONDS));
    // if error wins the race
    if (reader.errorTimes() > 0) {
        Assert.assertEquals(reader.doneTimes(), 0);
        Assert.assertEquals(observer.doneTimes(), 0);
        Assert.assertEquals(observer.errorTimes(), 1);
        Assert.assertEquals(writer.abortedTimes(), 0);
    } else // if cancel wins the race
    {
        Assert.assertEquals(observer.doneTimes(), 0);
        Assert.assertEquals(observer.errorTimes(), 1);
        Assert.assertEquals(writer.abortedTimes(), 1);
        Assert.assertEquals(reader.doneTimes(), 0);
        Assert.assertEquals(reader.errorTimes(), 0);
    }
}