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

Example 6 with ReadHandle

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

the class AbstractMIMEUnitTest method mockR2AndWrite.

//This is used when we need to mock out R2 and write a payload for our reader to read
protected void mockR2AndWrite(final ByteString payload, final int chunkSize, final String contentType) {
    _entityStream = mock(EntityStream.class);
    _readHandle = mock(ReadHandle.class);
    _streamRequest = mock(StreamRequest.class);
    //We have to use the AtomicReference holder technique to modify the current remaining buffer since the inner class
    //in doAnswer() can only access final variables.
    final AtomicReference<MultiPartMIMEReader.R2MultiPartMIMEReader> r2Reader = new AtomicReference<MultiPartMIMEReader.R2MultiPartMIMEReader>();
    //This takes the place of VariableByteStringWriter if we were to use R2 directly.
    final VariableByteStringViewer variableByteStringViewer = new VariableByteStringViewer(payload, chunkSize);
    //When data is requested, we write
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            final MultiPartMIMEReader.R2MultiPartMIMEReader reader = r2Reader.get();
            Object[] args = invocation.getArguments();
            //will always be 1 since MultiPartMIMEReader only does _rh.request(1)
            final int chunksRequested = (Integer) args[0];
            for (int i = 0; i < chunksRequested; i++) {
                //Our tests will run into a stack overflow unless we use a thread pool here to fire off the callbacks.
                //Especially in cases where the chunk size is 1. When the chunk size is one, the MultiPartMIMEReader
                //ends up doing many _rh.request(1) since each write is only 1 byte.
                //R2 uses a different technique to avoid stack overflows here which is unnecessary to emulate.
                _scheduledExecutorService.submit(new Runnable() {

                    @Override
                    public void run() {
                        ByteString clientData = variableByteStringViewer.onWritePossible();
                        if (clientData.equals(ByteString.empty())) {
                            reader.onDone();
                        } else {
                            reader.onDataAvailable(clientData);
                        }
                    }
                });
            }
            return null;
        }
    }).when(_readHandle).request(isA(Integer.class));
    //We need a final version of the read handle since its passed to an inner class below.
    final ReadHandle readHandleRef = _readHandle;
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            final MultiPartMIMEReader.R2MultiPartMIMEReader reader = (MultiPartMIMEReader.R2MultiPartMIMEReader) args[0];
            r2Reader.set(reader);
            //R2 calls init immediately upon setting the reader
            reader.onInit(readHandleRef);
            return null;
        }
    }).when(_entityStream).setReader(isA(MultiPartMIMEReader.R2MultiPartMIMEReader.class));
    when(_streamRequest.getEntityStream()).thenReturn(_entityStream);
    final String contentTypeHeader = contentType + ";somecustomparameter=somecustomvalue" + ";anothercustomparameter=anothercustomvalue";
    when(_streamRequest.getHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER)).thenReturn(contentTypeHeader);
}

Example 7 with ReadHandle

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

the class TestHttpNettyStreamClient method testCancelStreamRequests.

@Test(dataProvider = "requestResponseParameters", groups = TestGroupNames.TESTNG_GROUP_KNOWN_ISSUE)
public void testCancelStreamRequests(AbstractNettyStreamClient client, String method, int requestSize, int responseSize, boolean isFullRequest) throws Exception {
    AtomicInteger succeeded = new AtomicInteger(0);
    AtomicInteger failed = new AtomicInteger(0);
    Server server = new HttpServerBuilder().responseSize(responseSize).build();
    try {
        server.start();
        CountDownLatch latch = new CountDownLatch(REQUEST_COUNT);
        for (int i = 0; i < REQUEST_COUNT; i++) {
            StreamRequest request = new StreamRequestBuilder(new URI(URL)).setMethod(method).setHeader(HttpHeaderNames.HOST.toString(), HOST_NAME.toString()).build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(new byte[requestSize]))));
            RequestContext context = new RequestContext();
            context.putLocalAttr(R2Constants.IS_FULL_REQUEST, isFullRequest);
            client.streamRequest(request, context, new HashMap<>(), new TransportCallbackAdapter<>(new Callback<StreamResponse>() {

                @Override
                public void onSuccess(StreamResponse response) {
                    response.getEntityStream().setReader(new Reader() {

                        @Override
                        public void onDataAvailable(ByteString data) {
                        }

                        @Override
                        public void onDone() {
                            failed.incrementAndGet();
                            latch.countDown();
                        }

                        @Override
                        public void onError(Throwable e) {
                            failed.incrementAndGet();
                            latch.countDown();
                        }

                        @Override
                        public void onInit(ReadHandle rh) {
                            rh.cancel();
                            succeeded.incrementAndGet();
                            latch.countDown();
                        }
                    });
                }

                @Override
                public void onError(Throwable e) {
                    failed.incrementAndGet();
                    latch.countDown();
                }
            }));
        }
        if (!latch.await(30, TimeUnit.SECONDS)) {
            Assert.fail("Timeout waiting for responses. " + succeeded + " requests succeeded and " + failed + " requests failed out of total " + REQUEST_COUNT + " requests");
        }
        Assert.assertEquals(latch.getCount(), 0);
        Assert.assertEquals(failed.get(), 0);
        Assert.assertEquals(succeeded.get(), REQUEST_COUNT);
        FutureCallback<None> shutdownCallback = new FutureCallback<>();
        client.shutdown(shutdownCallback);
        shutdownCallback.get(30, TimeUnit.SECONDS);
    } finally {
        server.stop();
    }
}

Example 8 with ReadHandle

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

the class TestHttpNettyStreamClient method testResponseSize.

public void testResponseSize(AbstractNettyStreamClient client, int responseSize, int expectedResult) throws Exception {
    Server server = new HttpServerBuilder().responseSize(responseSize).build();
    try {
        server.start();
        RestRequest r = new RestRequestBuilder(new URI(URL)).build();
        FutureCallback<StreamResponse> cb = new FutureCallback<StreamResponse>();
        TransportCallback<StreamResponse> callback = new TransportCallbackAdapter<StreamResponse>(cb);
        client.streamRequest(Messages.toStreamRequest(r), new RequestContext(), new HashMap<String, String>(), callback);
        StreamResponse response = cb.get(30, TimeUnit.SECONDS);
        final CountDownLatch latch = new CountDownLatch(1);
        final AtomicReference<Throwable> error = new AtomicReference<Throwable>();
        response.getEntityStream().setReader(new Reader() {

            @Override
            public void onInit(ReadHandle rh) {
                rh.request(Integer.MAX_VALUE);
            }

            @Override
            public void onDataAvailable(ByteString data) {
            }

            @Override
            public void onDone() {
                latch.countDown();
            }

            @Override
            public void onError(Throwable e) {
                error.set(e);
                latch.countDown();
            }
        });
        if (!latch.await(30, TimeUnit.SECONDS)) {
            Assert.fail("Timeout waiting for response");
        }
        if (expectedResult == TOO_LARGE) {
            Assert.assertNotNull(error.get(), "Max response size exceeded, expected exception. ");
            verifyCauseChain(error.get(), TooLongFrameException.class);
        }
        if (expectedResult == RESPONSE_OK) {
            Assert.assertNull(error.get(), "Unexpected Exception: response size <= max size");
        }
    } catch (ExecutionException e) {
        if (expectedResult == RESPONSE_OK) {
            Assert.fail("Unexpected ExecutionException, response was <= max response size.", e);
        }
        verifyCauseChain(e, RemoteInvocationException.class, TooLongFrameException.class);
    } finally {
        server.stop();
    }
}

Example 9 with ReadHandle

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

the class TestHttpNettyStreamClient method testStreamRequests.

/**
   * Tests implementations of {@link AbstractNettyStreamClient} with different request dimensions.
   *
   * @param client Client implementation of {@link AbstractNettyStreamClient}
   * @param method HTTP request method
   * @param requestSize Request content size
   * @param responseSize Response content size
   * @param isFullRequest Whether to buffer a full request before stream
   * @throws Exception
   */
@Test(dataProvider = "requestResponseParameters")
public void testStreamRequests(AbstractNettyStreamClient client, String method, int requestSize, int responseSize, boolean isFullRequest) throws Exception {
    AtomicInteger succeeded = new AtomicInteger(0);
    AtomicInteger failed = new AtomicInteger(0);
    Server server = new HttpServerBuilder().responseSize(responseSize).build();
    try {
        server.start();
        CountDownLatch latch = new CountDownLatch(REQUEST_COUNT);
        for (int i = 0; i < REQUEST_COUNT; i++) {
            StreamRequest request = new StreamRequestBuilder(new URI(URL)).setMethod(method).setHeader(HttpHeaderNames.HOST.toString(), HOST_NAME.toString()).build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(new byte[requestSize]))));
            RequestContext context = new RequestContext();
            context.putLocalAttr(R2Constants.IS_FULL_REQUEST, isFullRequest);
            client.streamRequest(request, context, new HashMap<>(), new TransportCallbackAdapter<>(new Callback<StreamResponse>() {

                @Override
                public void onSuccess(StreamResponse response) {
                    response.getEntityStream().setReader(new Reader() {

                        ReadHandle _rh;

                        int _consumed = 0;

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

                        @Override
                        public void onDone() {
                            succeeded.incrementAndGet();
                            latch.countDown();
                        }

                        @Override
                        public void onError(Throwable e) {
                            failed.incrementAndGet();
                            latch.countDown();
                        }

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

                @Override
                public void onError(Throwable e) {
                    failed.incrementAndGet();
                    latch.countDown();
                }
            }));
        }
        if (!latch.await(30, TimeUnit.SECONDS)) {
            Assert.fail("Timeout waiting for responses. " + succeeded + " requests succeeded and " + failed + " requests failed out of total " + REQUEST_COUNT + " requests");
        }
        Assert.assertEquals(latch.getCount(), 0);
        Assert.assertEquals(failed.get(), 0);
        Assert.assertEquals(succeeded.get(), REQUEST_COUNT);
        FutureCallback<None> shutdownCallback = new FutureCallback<>();
        client.shutdown(shutdownCallback);
        shutdownCallback.get(30, TimeUnit.SECONDS);
    } finally {
        server.stop();
    }
}

Example 10 with ReadHandle

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

the class TestStreamResponse method testBackpressure.

@Test
public void testBackpressure() throws Exception {
    for (Client client : clients()) {
        StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, SMALL_URI));
        StreamRequest request = builder.build(EntityStreams.emptyStream());
        final AtomicInteger status = new AtomicInteger(-1);
        final CountDownLatch latch = new CountDownLatch(1);
        final AtomicReference<Throwable> error = new AtomicReference<Throwable>();
        final Callback<None> readerCallback = getReaderCallback(latch, error);
        final TimedBytesReader reader = new TimedBytesReader(BYTE, readerCallback) {

            int count = 0;

            @Override
            protected void requestMore(final ReadHandle rh) {
                count++;
                if (count % 16 == 0) {
                    _scheduler.schedule(new Runnable() {

                        @Override
                        public void run() {
                            rh.request(1);
                        }
                    }, INTERVAL, TimeUnit.MILLISECONDS);
                } else {
                    rh.request(1);
                }
            }
        };
        Callback<StreamResponse> callback = getCallback(status, readerCallback, reader);
        client.streamRequest(request, callback);
        latch.await(60000, TimeUnit.MILLISECONDS);
        Assert.assertNull(error.get());
        Assert.assertEquals(status.get(), RestStatus.OK);
        long serverSendTimespan = _smallHandler.getWriter().getStopTime() - _smallHandler.getWriter().getStartTime();
        long clientReceiveTimespan = reader.getStopTime() - reader.getStartTime();
        Assert.assertTrue(clientReceiveTimespan > 1000);
        double diff = Math.abs(clientReceiveTimespan - serverSendTimespan);
        double diffRatio = diff / serverSendTimespan;
        // make it generous to reduce the chance occasional test failures
        Assert.assertTrue(diffRatio < 0.2);
    }
}