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

Example 1 with WriteHandle

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

the class TestEntityStream method testNoStackOverflow.

@Test
public void testNoStackOverflow() throws Exception {
    Writer dumbWriter = new Writer() {

        WriteHandle _wh;

        long _count = 0;

        final int _total = 1024 * 1024 * 1024;

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

        @Override
        public void onWritePossible() {
            while (_wh.remaining() > 0 && _count < _total) {
                byte[] bytes = new byte[(int) Math.min(4096, _total - _count)];
                _wh.write(ByteString.copy(bytes));
                _count += bytes.length;
            }
            if (_count >= _total) {
                _wh.done();
            }
        }

        @Override
        public void onAbort(Throwable ex) {
        // do nothing
        }
    };
    Reader dumbReader = new Reader() {

        ReadHandle _rh;

        @Override
        public void onInit(ReadHandle rh) {
            _rh = rh;
            _rh.request(1);
        }

        @Override
        public void onDataAvailable(ByteString data) {
            _rh.request(1);
        }

        @Override
        public void onDone() {
        }

        @Override
        public void onError(Throwable e) {
        }
    };
    EntityStreams.newEntityStream(dumbWriter).setReader(dumbReader);
}

Example 2 with WriteHandle

use of com.linkedin.r2.message.stream.entitystream.WriteHandle 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 3 with WriteHandle

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

the class TestMIMEInputStream method testSingleOnWritePossibleDataSources.

@Test(dataProvider = "singleOnWritePossibleDataSources")
public void testSingleOnWritePossibleDataSources(final byte[] inputData, final StrictByteArrayInputStream inputStream, final int writesRemainingPerOnWritePossibles, final int expectedTotalWrites) {
    // Setup:
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(inputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).build();
    // Simulate a write handle that offers decreasing writeHandle.remaining()
    final Integer[] remainingWriteHandleCount = simulateDecreasingWriteHandleCount(writesRemainingPerOnWritePossibles);
    when(writeHandle.remaining()).thenReturn(writesRemainingPerOnWritePossibles, remainingWriteHandleCount);
    final ByteArrayOutputStream byteArrayOutputStream = setupMockWriteHandleToOutputStream(writeHandle);
    // Setup for done()
    final CountDownLatch doneLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            doneLatch.countDown();
            return null;
        }
    }).when(writeHandle).done();
    // /////////////////////////////////
    // Start things off
    // Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onWritePossible();
    // Wait to finish
    try {
        boolean successful = doneLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        if (!successful) {
            Assert.fail("Timeout when waiting for input stream to completely transfer");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to completely transfer");
    }
    // /////////////////////////////////
    // Assert
    Assert.assertEquals(byteArrayOutputStream.toByteArray(), inputData, "All data from the input stream should have successfully been transferred");
    Assert.assertEquals(inputStream.isClosed(), true);
    // Mock verifies:
    verify(writeHandle, times(expectedTotalWrites)).write(isA(ByteString.class));
    // One extra remaining() call for knowing the stream has ended.
    verify(writeHandle, times(expectedTotalWrites + 1)).remaining();
    verify(writeHandle, never()).error(isA(Throwable.class));
    verify(writeHandle, times(1)).done();
    verifyNoMoreInteractions(writeHandle);
}

Example 4 with WriteHandle

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

the class TestMIMEInputStream method testExceptionDataSources.

@Test(dataProvider = "exceptionDataSources")
public void testExceptionDataSources(final ExceptionThrowingByteArrayInputStream exceptionThrowingByteArrayInputStream, final int expectedTotalWrites, final int expectedWriteHandleRemainingCalls, final byte[] expectedDataWritten) {
    // Setup:
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(exceptionThrowingByteArrayInputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).build();
    // Doesn't matter what we return here as long as its constant and above 0.
    when(writeHandle.remaining()).thenReturn(500);
    final ByteArrayOutputStream byteArrayOutputStream = setupMockWriteHandleToOutputStream(writeHandle);
    // Setup for error()
    final CountDownLatch errorLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            errorLatch.countDown();
            return null;
        }
    }).when(writeHandle).error(isA(Throwable.class));
    // /////////////////////////////////
    // Start things off
    // Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onWritePossible();
    // Wait to finish
    try {
        boolean successful = errorLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        if (!successful) {
            Assert.fail("Timeout when waiting for input stream to completely transfer");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to transfer");
    }
    // /////////////////////////////////
    // Assert
    Assert.assertEquals(byteArrayOutputStream.toByteArray(), expectedDataWritten, "Partial data should have been transferred in the case of an exception");
    Assert.assertEquals(exceptionThrowingByteArrayInputStream.isClosed(), true);
    // Mock verifies:
    verify(writeHandle, times(expectedTotalWrites)).write(isA(ByteString.class));
    verify(writeHandle, times(expectedWriteHandleRemainingCalls)).remaining();
    // Since we can't override equals
    verify(writeHandle, times(1)).error(isA(IOException.class));
    verify(writeHandle, never()).done();
    verifyNoMoreInteractions(writeHandle);
}

Example 5 with WriteHandle

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

the class TestMIMEInputStream method testDifferentDataSourceSizes.

@Test(dataProvider = "differentDataSourceSizes")
public void testDifferentDataSourceSizes(final byte[] inputData, final StrictByteArrayInputStream inputStream, final int writesRemainingPerOnWritePossibles, final int expectedWriteHandleRemaining, final int expectedTotalWrites) {
    // Setup:
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(inputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).build();
    // Simulate a write handle that offers decreasing writeHandle.remaining()
    final Integer[] remainingWriteHandleCount = simulateDecreasingWriteHandleCount(writesRemainingPerOnWritePossibles);
    when(writeHandle.remaining()).thenReturn(writesRemainingPerOnWritePossibles, remainingWriteHandleCount);
    final ByteArrayOutputStream byteArrayOutputStream = setupMockWriteHandleToOutputStream(writeHandle);
    // Setup for done()
    final CountDownLatch doneLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            doneLatch.countDown();
            return null;
        }
    }).when(writeHandle).done();
    // /////////////////////////////////
    // Start things off
    // Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onWritePossible();
    // Wait to finish
    try {
        boolean successful = doneLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        if (!successful) {
            Assert.fail("Timeout when waiting for input stream to completely transfer");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to completely transfer");
    }
    // /////////////////////////////////
    // Assert
    Assert.assertEquals(byteArrayOutputStream.toByteArray(), inputData, "All data from the input stream should have successfully been transferred");
    Assert.assertEquals(inputStream.isClosed(), true);
    // Mock verifies:
    verify(writeHandle, times(expectedTotalWrites)).write(isA(ByteString.class));
    verify(writeHandle, times(expectedWriteHandleRemaining)).remaining();
    verify(writeHandle, never()).error(isA(Throwable.class));
    verify(writeHandle, times(1)).done();
    verifyNoMoreInteractions(writeHandle);
}