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

Example 16 with Writer

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

the class TestEntityStream method testRaceBetweenWriteAndCancel.

private void testRaceBetweenWriteAndCancel(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() > 0) {
                prepareLatch.countDown();
                startLatch.await();
                writer.write();
            }
            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);
    Assert.assertEquals(writer.abortedTimes(), 1);
    Assert.assertEquals(reader.getChunkCount(), observer.getChunkCount());
    Assert.assertEquals(observer.errorTimes(), 1);
    Assert.assertEquals(observer.getLastEvent(), "onError");
}

Example 17 with Writer

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

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

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

the class Http2FrameListener method onHeadersRead.

@Override
public void onHeadersRead(ChannelHandlerContext ctx, int streamId, Http2Headers headers, int padding, boolean endOfStream) throws Http2Exception {
    LOG.debug("Received HTTP/2 HEADERS frame, stream={}, end={}, headers={}, padding={}bytes", new Object[] { streamId, endOfStream, headers.size(), padding });
    // Ignores response for the upgrade request
    if (streamId == Http2CodecUtil.HTTP_UPGRADE_STREAM_ID) {
        return;
    }
    final StreamResponseBuilder builder = new StreamResponseBuilder();
    // Process HTTP/2 pseudo headers
    if (headers.status() != null) {
        builder.setStatus(Integer.parseInt(headers.status().toString()));
    }
    if (headers.authority() != null) {
        builder.addHeaderValue(HttpHeaderNames.HOST.toString(), headers.authority().toString());
    }
    // Process other HTTP headers
    for (Map.Entry<CharSequence, CharSequence> header : headers) {
        if (Http2Headers.PseudoHeaderName.isPseudoHeader(header.getKey())) {
            // Do no set HTTP/2 pseudo headers to response
            continue;
        }
        final String key = header.getKey().toString();
        final String value = header.getValue().toString();
        if (key.equalsIgnoreCase(HttpConstants.RESPONSE_COOKIE_HEADER_NAME)) {
            builder.addCookie(value);
        } else {
            builder.unsafeAddHeaderValue(key, value);
        }
    }
    // Gets async pool handle from stream properties
    Http2Connection.PropertyKey handleKey = ctx.channel().attr(Http2ClientPipelineInitializer.CHANNEL_POOL_HANDLE_ATTR_KEY).get();
    TimeoutAsyncPoolHandle<?> handle = _connection.stream(streamId).removeProperty(handleKey);
    if (handle == null) {
        _lifecycleManager.onError(ctx, Http2Exception.connectionError(Http2Error.PROTOCOL_ERROR, "No channel pool handle is associated with this stream", streamId));
        return;
    }
    final StreamResponse response;
    if (endOfStream) {
        response = builder.build(EntityStreams.emptyStream());
        ctx.fireChannelRead(handle);
    } else {
        // Associate an entity stream writer to the HTTP/2 stream
        final TimeoutBufferedWriter writer = new TimeoutBufferedWriter(ctx, streamId, _maxContentLength, handle);
        if (_connection.stream(streamId).setProperty(_writerKey, writer) != null) {
            _lifecycleManager.onError(ctx, Http2Exception.connectionError(Http2Error.PROTOCOL_ERROR, "Another writer has already been associated with current stream ID", streamId));
            return;
        }
        // Prepares StreamResponse for the channel pipeline
        EntityStream entityStream = EntityStreams.newEntityStream(writer);
        response = builder.build(entityStream);
    }
    // Gets callback from stream properties
    Http2Connection.PropertyKey callbackKey = ctx.channel().attr(Http2ClientPipelineInitializer.CALLBACK_ATTR_KEY).get();
    TransportCallback<?> callback = _connection.stream(streamId).removeProperty(callbackKey);
    if (callback != null) {
        ctx.fireChannelRead(new ResponseWithCallback<Response, TransportCallback<?>>(response, callback));
    }
}

Example 20 with Writer

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

the class TestCompressionEcho method testResponseCompression.

@Test(dataProvider = "compressionEchoData")
public void testResponseCompression(Client client, long bytes) throws InterruptedException, TimeoutException, ExecutionException {
    StreamRequestBuilder builder = new StreamRequestBuilder((Bootstrap.createHttpURI(PORT, ECHO_URI)));
    BytesWriter writer = new BytesWriter(bytes, BYTE);
    StreamRequest request = builder.build(EntityStreams.newEntityStream(writer));
    // add operation to enable sending accept encoding
    RequestContext requestContext = new RequestContext();
    requestContext.putLocalAttr(R2Constants.OPERATION, "get");
    final FutureCallback<StreamResponse> callback = new FutureCallback<StreamResponse>();
    client.streamRequest(request, requestContext, 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);
    response.getEntityStream().setReader(reader);
    readerCallback.get(60, TimeUnit.SECONDS);
    Assert.assertEquals(reader.getTotalBytes(), bytes);
    Assert.assertTrue(reader.allBytesCorrect());
}