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

Example 1 with EntityStream

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

the class RAPResponseDecoder method channelRead0.

@Override
protected void channelRead0(final ChannelHandlerContext ctx, HttpObject msg) throws Exception {
    if (msg instanceof HttpResponse) {
        HttpResponse m = (HttpResponse) msg;
        _shouldCloseConnection = !HttpUtil.isKeepAlive(m);
        if (HttpUtil.is100ContinueExpected(m)) {
            ctx.writeAndFlush(CONTINUE).addListener(new ChannelFutureListener() {

                @Override
                public void operationComplete(ChannelFuture future) throws Exception {
                    if (!future.isSuccess()) {
                        ctx.fireExceptionCaught(future.cause());
                    }
                }
            });
        }
        if (!m.decoderResult().isSuccess()) {
            ctx.fireExceptionCaught(m.decoderResult().cause());
            return;
        }
        // remove chunked encoding.
        if (HttpUtil.isTransferEncodingChunked(m)) {
            HttpUtil.setTransferEncodingChunked(m, false);
        }
        Timeout<None> timeout = ctx.channel().attr(TIMEOUT_ATTR_KEY).getAndRemove();
        if (timeout == null) {
            LOG.debug("dropped a response after channel inactive or exception had happened.");
            return;
        }
        final TimeoutBufferedWriter writer = new TimeoutBufferedWriter(ctx, _maxContentLength, BUFFER_HIGH_WATER_MARK, BUFFER_LOW_WATER_MARK, timeout);
        EntityStream entityStream = EntityStreams.newEntityStream(writer);
        _chunkedMessageWriter = writer;
        StreamResponseBuilder builder = new StreamResponseBuilder();
        builder.setStatus(m.status().code());
        for (Map.Entry<String, String> e : m.headers()) {
            String key = e.getKey();
            String value = e.getValue();
            if (key.equalsIgnoreCase(HttpConstants.RESPONSE_COOKIE_HEADER_NAME)) {
                builder.addCookie(value);
            } else {
                builder.unsafeAddHeaderValue(key, value);
            }
        }
        ctx.fireChannelRead(builder.build(entityStream));
    } else if (msg instanceof HttpContent) {
        HttpContent chunk = (HttpContent) msg;
        TimeoutBufferedWriter currentWriter = _chunkedMessageWriter;
        // Sanity check
        if (currentWriter == null) {
            throw new IllegalStateException("received " + HttpContent.class.getSimpleName() + " without " + HttpResponse.class.getSimpleName());
        }
        if (!chunk.decoderResult().isSuccess()) {
            this.exceptionCaught(ctx, chunk.decoderResult().cause());
        }
        currentWriter.processHttpChunk(chunk);
        if (chunk instanceof LastHttpContent) {
            _chunkedMessageWriter = null;
        }
    } else {
        // something must be wrong, but let's proceed so that
        // handler after us has a chance to process it.
        ctx.fireChannelRead(msg);
    }
}

Example 2 with EntityStream

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

the class TestStreamRequest method testErrorReceiver.

@Test
public void testErrorReceiver() throws Exception {
    for (Client client : clients()) {
        final long totalBytes = SMALL_BYTES_NUM;
        EntityStream entityStream = EntityStreams.newEntityStream(new BytesWriter(totalBytes, BYTE));
        StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, ERROR_RECEIVER_URI));
        StreamRequest request = builder.setMethod("POST").build(entityStream);
        final CountDownLatch latch = new CountDownLatch(1);
        final AtomicReference<Throwable> error = new AtomicReference<Throwable>();
        Callback<StreamResponse> callback = new Callback<StreamResponse>() {

            @Override
            public void onError(Throwable e) {
                error.set(e);
                latch.countDown();
            }

            @Override
            public void onSuccess(StreamResponse result) {
                latch.countDown();
            }
        };
        client.streamRequest(request, callback);
        latch.await();
        Assert.assertNotNull(error.get());
    }
}

Example 3 with EntityStream

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

the class TestResponseCompression method testResponseCompression.

private void testResponseCompression(URI uri, long bytes, String acceptEncoding, final StreamingCompressor compressor) throws InterruptedException, TimeoutException, ExecutionException {
    for (Client client : clients()) {
        StreamRequestBuilder builder = new StreamRequestBuilder((Bootstrap.createHttpURI(PORT, uri)));
        builder.addHeaderValue(HttpConstants.ACCEPT_ENCODING, acceptEncoding);
        StreamRequest request = builder.build(EntityStreams.emptyStream());
        final FutureCallback<StreamResponse> callback = new FutureCallback<StreamResponse>();
        client.streamRequest(request, callback);
        final StreamResponse response = callback.get(60, TimeUnit.SECONDS);
        Assert.assertEquals(response.getStatus(), RestStatus.OK);
        final FutureCallback<None> readerCallback = new FutureCallback<None>();
        final BytesReader reader = new BytesReader(BYTE, readerCallback);
        final EntityStream decompressedStream = compressor.inflate(response.getEntityStream());
        decompressedStream.setReader(reader);
        readerCallback.get(60, TimeUnit.SECONDS);
        Assert.assertEquals(reader.getTotalBytes(), bytes);
        Assert.assertTrue(reader.allBytesCorrect());
    }
}

Example 4 with EntityStream

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

the class TestStreamRequest method testBackPressure.

@Test
public void testBackPressure() throws Exception {
    for (Client client : clients()) {
        final long totalBytes = SMALL_BYTES_NUM;
        TimedBytesWriter writer = new TimedBytesWriter(totalBytes, BYTE);
        EntityStream entityStream = EntityStreams.newEntityStream(writer);
        StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, RATE_LIMITED_URI));
        StreamRequest request = builder.setMethod("POST").build(entityStream);
        final AtomicInteger status = new AtomicInteger(-1);
        final CountDownLatch latch = new CountDownLatch(1);
        Callback<StreamResponse> callback = expectSuccessCallback(latch, status);
        client.streamRequest(request, callback);
        latch.await(60000, TimeUnit.MILLISECONDS);
        Assert.assertEquals(status.get(), RestStatus.OK);
        TimedBytesReader reader = _rateLimitedRequestHandler.getReader();
        Assert.assertNotNull(reader);
        Assert.assertEquals(totalBytes, reader.getTotalBytes());
        Assert.assertTrue(reader.allBytesCorrect());
        long clientSendTimespan = writer.getStopTime() - writer.getStartTime();
        long serverReceiveTimespan = reader.getStopTime() - reader.getStartTime();
        Assert.assertTrue(serverReceiveTimespan > 1000);
        double diff = Math.abs(serverReceiveTimespan - clientSendTimespan);
        double diffRatio = diff / clientSendTimespan;
        // make it generous to reduce the chance occasional test failures
        Assert.assertTrue(diffRatio < 0.2);
    }
}

Example 5 with EntityStream

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

the class TestEntityStream method testRaceBetweenDoneAndCancel.

private void testRaceBetweenDoneAndCancel(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.done();
            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));
    // in any case, reader shouldn't fail
    Assert.assertEquals(reader.errorTimes(), 0);
    // if done wins the race
    if (reader.doneTimes() > 0) {
        Assert.assertEquals(reader.doneTimes(), 1);
        Assert.assertEquals(observer.doneTimes(), 1);
        Assert.assertEquals(observer.errorTimes(), 0);
        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);
    }
}