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

Example 6 with WriteHandle

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

the class TestMIMEInputStream method testAbortBeforeWrite.

// /////////////////////////////////////////////////////////////////////////////////////
// 
// These tests will verify that aborts work properly and close the input stream regardless of the
// current state.
// This test will simulate an abort before any writes requested.
@Test
public void testAbortBeforeWrite() throws Exception {
    final byte[] smallInputData = "b".getBytes();
    final StrictByteArrayInputStream inputStream = new StrictByteArrayInputStream(smallInputData);
    final StrictByteArrayInputStream spyInputStream = spy(inputStream);
    // Setup:
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(spyInputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).build();
    // Setup for the close on the input stream.
    // The close must happen for the test to finish.
    final CountDownLatch closeInputStreamLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            closeInputStreamLatch.countDown();
            return null;
        }
    }).when(spyInputStream).close();
    // /////////////////////////////////
    // Start things off
    // Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onAbort(new IOException());
    // Wait to finish
    try {
        boolean successful = closeInputStreamLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        if (!successful) {
            Assert.fail("Timeout when waiting for abort to happen!");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to be closed!");
    }
    // /////////////////////////////////
    // Mock verifies:
    verify(spyInputStream, times(1)).close();
    verify(spyInputStream, never()).read(isA(byte[].class));
    verify(writeHandle, never()).write(isA(ByteString.class));
    verify(writeHandle, never()).remaining();
    verify(writeHandle, never()).error(isA(Throwable.class));
    verify(writeHandle, never()).done();
    verifyNoMoreInteractions(writeHandle);
}

Example 7 with WriteHandle

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

the class MultiPartMIMEChainReaderCallback method onNewDataSource.

@Override
public void onNewDataSource(final MultiPartMIMEDataSourceWriter multiPartMIMEDataSourceWriter) {
    multiPartMIMEDataSourceWriter.onInit(new WriteHandle() {

        @Override
        public void write(ByteString data) {
            _writeHandle.write(data);
        }

        @Override
        public void done() {
        // We intentionally ignore this, since we will be responsible for calling writeHandle.done() when all the data
        // sources represented by the MultiPartMIMEDataSourceIterator have finished.
        }

        @Override
        public void error(Throwable throwable) {
            _writeHandle.error(throwable);
        }

        @Override
        public int remaining() {
            return _writeHandle.remaining();
        }
    });
    _currentDataSource = multiPartMIMEDataSourceWriter;
    ByteString serializedBoundaryAndHeaders = null;
    try {
        serializedBoundaryAndHeaders = MultiPartMIMEUtils.serializeBoundaryAndHeaders(_normalEncapsulationBoundary, multiPartMIMEDataSourceWriter);
    } catch (IOException ioException) {
        // Should never happen
        onStreamError(ioException);
    }
    _writeHandle.write(serializedBoundaryAndHeaders);
    if (_writeHandle.remaining() > 0) {
        multiPartMIMEDataSourceWriter.onWritePossible();
    }
}

Example 8 with WriteHandle

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

the class TestEntityStream method testWriterInitError.

@Test
public void testWriterInitError() throws Exception {
    final CountDownLatch latch = new CountDownLatch(1);
    final TestWriter writer = new TestWriter() {

        @Override
        public void onInit(WriteHandle wh) {
            throw new RuntimeException();
        }
    };
    TestObserver observer = new TestObserver();
    final EntityStream es = EntityStreams.newEntityStream(writer);
    es.addObserver(observer);
    final ControlReader reader = new ControlReader();
    ExecutorService executor = Executors.newFixedThreadPool(1);
    executor.submit(new Runnable() {

        @Override
        public void run() {
            es.setReader(reader);
            reader.read(5);
            latch.countDown();
        }
    });
    Assert.assertTrue(latch.await(1000, TimeUnit.MILLISECONDS));
    Assert.assertEquals(reader.errorTimes(), 1);
    Assert.assertEquals(reader.getChunkCount(), 0);
    Assert.assertEquals(observer.errorTimes(), 1);
    Assert.assertEquals(reader.getChunkCount(), 0);
    Assert.assertEquals(writer.abortedTimes(), 1);
    Assert.assertEquals(writer.getWritePossibleCount(), 0);
}

Example 9 with WriteHandle

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

the class TestEntityStream method testWriterAndReaderInitError.

@Test
public void testWriterAndReaderInitError() throws Exception {
    final CountDownLatch latch = new CountDownLatch(1);
    final TestWriter writer = new TestWriter() {

        @Override
        public void onInit(WriteHandle wh) {
            throw new RuntimeException();
        }
    };
    TestObserver observer = new TestObserver();
    final EntityStream es = EntityStreams.newEntityStream(writer);
    es.addObserver(observer);
    final ControlReader reader = new ControlReader() {

        @Override
        public void onInit(ReadHandle rh) {
            throw new RuntimeException();
        }
    };
    ExecutorService executor = Executors.newFixedThreadPool(1);
    executor.submit(new Runnable() {

        @Override
        public void run() {
            es.setReader(reader);
            latch.countDown();
        }
    });
    Assert.assertTrue(latch.await(1000, TimeUnit.MILLISECONDS));
    Assert.assertEquals(reader.errorTimes(), 1);
    Assert.assertEquals(reader.getChunkCount(), 0);
    Assert.assertEquals(observer.errorTimes(), 1);
    Assert.assertEquals(reader.getChunkCount(), 0);
    Assert.assertEquals(writer.abortedTimes(), 1);
    Assert.assertEquals(writer.getWritePossibleCount(), 0);
}

Example 10 with WriteHandle

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

the class TestStreamingTimeout method testStreamTimeoutWithStreamTimeoutInClientStream.

@Test
public void testStreamTimeoutWithStreamTimeoutInClientStream() throws Exception {
    final EntityStream entityStream = EntityStreams.newEntityStream(new BytesWriter(LARGE_BYTES_NUM, BYTE) {

        int count = 2;

        @Override
        protected void afterWrite(WriteHandle wh, long written) {
            count = count * 2;
            long delay = Math.min(count, HTTP_STREAMING_TIMEOUT);
            try {
                Thread.sleep(delay);
            } catch (Exception ex) {
            // Do Nothing
            }
        }
    });
    final StreamRequestBuilder builder = new StreamRequestBuilder(_clientProvider.createHttpURI(_port, NON_RATE_LIMITED_URI));
    final StreamRequest request = builder.setMethod("POST").build(entityStream);
    final AtomicReference<Throwable> throwable = new AtomicReference<>();
    final CountDownLatch latch = new CountDownLatch(1);
    final Callback<StreamResponse> callback = expectErrorCallback(latch, throwable);
    _client.streamRequest(request, callback);
    latch.await(30000, TimeUnit.MILLISECONDS);
    Assert.assertNotNull(throwable.get());
    final Throwable rootCause = ExceptionUtils.getRootCause(throwable.get());
    Assert.assertTrue(rootCause instanceof TimeoutException);
    final TimeoutException timeoutException = (TimeoutException) rootCause;
    Assert.assertEquals(timeoutException.getMessage(), String.format(StreamingTimeout.STREAMING_TIMEOUT_MESSAGE, HTTP_STREAMING_TIMEOUT));
}