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

Example 66 with StreamRequest

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

the class TestStreamRequest method testRequestLarge.

@Test
public void testRequestLarge() throws Exception {
    final long totalBytes = LARGE_BYTES_NUM;
    EntityStream entityStream = EntityStreams.newEntityStream(new BytesWriter(totalBytes, BYTE));
    StreamRequestBuilder builder = new StreamRequestBuilder(_clientProvider.createHttpURI(_port, LARGE_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);
    BytesReader reader = _checkRequestHandler.getReader();
    Assert.assertNotNull(reader);
    Assert.assertEquals(totalBytes, reader.getTotalBytes());
    Assert.assertTrue(reader.allBytesCorrect());
}

Example 67 with StreamRequest

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

the class TestStreamResponseCompression method testResponseCompression.

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

Example 68 with StreamRequest

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

the class Http2StreamCodec method write.

@Override
public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
    if (!(msg instanceof RequestWithCallback)) {
        ctx.write(msg, promise);
        return;
    }
    Request request = ((RequestWithCallback) msg).request();
    Http2ConnectionEncoder encoder = encoder();
    int streamId = connection().local().incrementAndGetNextStreamId();
    final ChannelFuture headersFuture;
    if (request instanceof StreamRequest) {
        final StreamRequest streamRequest = (StreamRequest) request;
        final Http2Headers http2Headers = NettyRequestAdapter.toHttp2Headers(streamRequest);
        final BufferedReader bufferedReader = new BufferedReader(ctx, encoder, streamId, ((RequestWithCallback) msg).handle());
        final OrderedEntityStreamReader reader = new OrderedEntityStreamReader(ctx, bufferedReader);
        streamRequest.getEntityStream().setReader(reader);
        LOG.debug("Sent HTTP/2 HEADERS frame, stream={}, end={}, headers={}, padding={}bytes", new Object[] { streamId, NOT_END_STREAM, http2Headers.size(), NO_PADDING });
        headersFuture = encoder.writeHeaders(ctx, streamId, http2Headers, NO_PADDING, NOT_END_STREAM, promise);
        headersFuture.addListener(future -> {
            if (future.isSuccess()) {
                reader.request(BufferedReader.MAX_BUFFERED_CHUNKS);
            }
        });
    } else if (request instanceof RestRequest) {
        final RestRequest restRequest = (RestRequest) request;
        final Http2Headers headers = NettyRequestAdapter.toHttp2Headers(restRequest);
        LOG.debug("Sent HTTP/2 HEADERS frame, stream={}, end={}, headers={}, padding={}bytes", new Object[] { streamId, NOT_END_STREAM, headers.size(), NO_PADDING });
        headersFuture = encoder.writeHeaders(ctx, streamId, headers, NO_PADDING, NOT_END_STREAM, promise);
        headersFuture.addListener(future -> {
            if (future.isSuccess()) {
                final ByteBuf data = Unpooled.wrappedBuffer(restRequest.getEntity().asByteBuffer());
                LOG.debug("Sent HTTP/2 DATA frame, stream={}, end={}, data={}bytes, padding={}bytes", new Object[] { streamId, END_STREAM, data.readableBytes(), NO_PADDING });
                encoder.writeData(ctx, streamId, data, NO_PADDING, END_STREAM, ctx.newPromise());
                ctx.channel().flush();
            }
        });
    } else {
        // Release the handle to put the channel back to the pool
        ((RequestWithCallback) msg).handle().release();
        throw new IllegalArgumentException("Request is neither StreamRequest or RestRequest");
    }
    final TransportCallback<?> callback = ((RequestWithCallback) msg).callback();
    @SuppressWarnings("unchecked") final TimeoutAsyncPoolHandle<Channel> handle = (TimeoutAsyncPoolHandle<Channel>) ((RequestWithCallback) msg).handle();
    headersFuture.addListener(future -> {
        if (future.isSuccess()) {
            // Sets TransportCallback as a stream property to be retrieved later
            Http2PipelinePropertyUtil.set(ctx, connection(), streamId, Http2ClientPipelineInitializer.CALLBACK_ATTR_KEY, callback);
            // Sets AsyncPoolHandle as a stream property to be retrieved later
            Http2PipelinePropertyUtil.set(ctx, connection(), streamId, Http2ClientPipelineInitializer.CHANNEL_POOL_HANDLE_ATTR_KEY, handle);
            // Sets a timeout task to reset stream
            // Channel pool handle is also released at timeout
            handle.addTimeoutTask(() -> {
                LOG.debug("Reset stream upon timeout, stream={}", streamId);
                resetStream(ctx, streamId, Http2Error.CANCEL.code(), ctx.newPromise());
                ctx.flush();
            });
        } else {
            // Invokes callback onResponse with the error thrown during write header or data
            callback.onResponse(TransportResponseImpl.error(future.cause()));
            // Releases the handle to put the channel back to the pool
            handle.release();
            // Resets the stream if a stream is created after we sent header
            if (connection().stream(streamId) != null) {
                LOG.debug("Reset stream upon timeout, stream={}", streamId);
                resetStream(ctx, streamId, Http2Error.CANCEL.code(), ctx.newPromise());
                ctx.flush();
            }
        }
    });
}

Example 69 with StreamRequest

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

the class TestNettyRequestAdapter method testStreamToHttp2HeadersCookies.

@Test
public void testStreamToHttp2HeadersCookies() throws Exception {
    StreamRequestBuilder streamRequestBuilder = new StreamRequestBuilder(new URI(ANY_URI));
    IntStream.range(0, 10).forEach(i -> streamRequestBuilder.addCookie(ANY_COOKIE));
    StreamRequest request = streamRequestBuilder.build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(ANY_ENTITY.getBytes()))));
    Http2Headers headers = NettyRequestAdapter.toHttp2Headers(request);
    Assert.assertNotNull(headers);
    List<CharSequence> cookies = headers.getAll(HttpHeaderNames.COOKIE);
    Assert.assertNotNull(cookies);
    Assert.assertEquals(cookies.size(), 10);
}

Example 70 with StreamRequest

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

the class TestNettyRequestAdapter method testStreamToNettyRequestWithMultipleCookies.

@Test
public void testStreamToNettyRequestWithMultipleCookies() throws Exception {
    StreamRequestBuilder streamRequestBuilder = new StreamRequestBuilder(new URI(ANY_URI));
    streamRequestBuilder.setCookies(ANY_COOKIES);
    StreamRequest streamRequest = streamRequestBuilder.build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(ANY_ENTITY.getBytes()))));
    HttpRequest nettyRequest = NettyRequestAdapter.toNettyRequest(streamRequest);
    Assert.assertEquals(nettyRequest.headers().get("Cookie"), ENCODED_COOKIES_HEADER_VALUE);
}