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

Example 91 with StreamResponse

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

the class TestDisruptFilter method testStreamLatencyDisrupt.

@Test
public void testStreamLatencyDisrupt() throws Exception {
    final RequestContext requestContext = new RequestContext();
    requestContext.putLocalAttr(DISRUPT_CONTEXT_KEY, DisruptContexts.delay(REQUEST_LATENCY));
    final DisruptFilter filter = new DisruptFilter(_scheduler, _executor, REQUEST_TIMEOUT);
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicBoolean success = new AtomicBoolean(false);
    final NextFilter<StreamRequest, StreamResponse> next = new NextFilter<StreamRequest, StreamResponse>() {

        @Override
        public void onRequest(StreamRequest restRequest, RequestContext requestContext, Map<String, String> wireAttrs) {
            success.set(true);
            latch.countDown();
        }

        @Override
        public void onResponse(StreamResponse restResponse, RequestContext requestContext, Map<String, String> wireAttrs) {
            latch.countDown();
        }

        @Override
        public void onError(Throwable ex, RequestContext requestContext, Map<String, String> wireAttrs) {
            latch.countDown();
        }
    };
    filter.onStreamRequest(new StreamRequestBuilder(new URI(URI)).build(EntityStreams.emptyStream()), requestContext, Collections.emptyMap(), next);
    Assert.assertTrue(latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS), "Missing NextFilter invocation");
    Assert.assertTrue(success.get(), "Unexpected method invocation");
}

Example 92 with StreamResponse

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

the class TestMessages method testToStreamTransportCallbackStreamException.

@Test
public void testToStreamTransportCallbackStreamException() {
    TransportCallback<RestResponse> restCallback = response -> {
        Assert.assertTrue(response.hasError());
        Assert.assertNotNull(response.getError());
        Assert.assertTrue(response.getError() instanceof RestException);
        Assert.assertNotNull(response.getWireAttributes());
        Assert.assertEquals(response.getWireAttributes(), WIRE_ATTR);
    };
    TransportCallback<StreamResponse> streamCallback = Messages.toStreamTransportCallback(restCallback);
    StreamResponseBuilder builder = new StreamResponseBuilder();
    StreamResponse streamResponse = builder.build(EntityStreams.newEntityStream(new ByteStringWriter(DATA)));
    streamCallback.onResponse(TransportResponseImpl.error(new StreamException(streamResponse, new IllegalStateException()), WIRE_ATTR));
}

Example 93 with StreamResponse

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

the class TestClientTimeout method testReadAfterTimeout.

@Test(dataProvider = "clients")
public void testReadAfterTimeout(Client client) throws Exception {
    StreamRequest request = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, NORMAL_URI)).build(EntityStreams.emptyStream());
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicReference<StreamResponse> response = new AtomicReference<StreamResponse>();
    client.streamRequest(request, new Callback<StreamResponse>() {

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

        @Override
        public void onSuccess(StreamResponse result) {
            response.set(result);
            latch.countDown();
        }
    });
    latch.await(5000, TimeUnit.MILLISECONDS);
    Assert.assertNotNull(response.get());
    // let it timeout before we read
    Thread.sleep(5000);
    final AtomicReference<Throwable> throwable = new AtomicReference<Throwable>();
    final CountDownLatch errorLatch = new CountDownLatch(1);
    Reader reader = new DrainReader() {

        @Override
        public void onError(Throwable ex) {
            throwable.set(ex);
            errorLatch.countDown();
        }
    };
    response.get().getEntityStream().setReader(reader);
    errorLatch.await(5000, TimeUnit.MILLISECONDS);
    Assert.assertNotNull(throwable.get());
    Throwable rootCause = ExceptionUtils.getRootCause(throwable.get());
    Assert.assertTrue(rootCause instanceof TimeoutException);
    Assert.assertEquals(rootCause.getMessage(), "Timeout while receiving the response entity.");
}

Example 94 with StreamResponse

use of com.linkedin.r2.message.stream.StreamResponse 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());
}

Example 95 with StreamResponse

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

the class ServerStreamCompressionFilter method onStreamResponse.

/**
   * Optionally compresses outgoing response
   * */
@Override
public void onStreamResponse(final StreamResponse res, final RequestContext requestContext, final Map<String, String> wireAttrs, final NextFilter<StreamRequest, StreamResponse> nextFilter) {
    StreamResponse response = res;
    try {
        String responseCompression = (String) requestContext.getLocalAttr(HttpConstants.ACCEPT_ENCODING);
        if (responseCompression == null) {
            throw new CompressionException(HttpConstants.ACCEPT_ENCODING + " not in local attribute.");
        }
        List<AcceptEncoding> parsedEncodings = AcceptEncoding.parseAcceptEncodingHeader(responseCompression, _supportedEncoding);
        StreamEncodingType selectedEncoding = AcceptEncoding.chooseBest(parsedEncodings);
        //Check if there exists an acceptable encoding
        if (selectedEncoding == null) {
            //Not acceptable encoding status
            response = new StreamResponseBuilder().setStatus(HttpConstants.NOT_ACCEPTABLE).build(EntityStreams.emptyStream());
        } else if (selectedEncoding != StreamEncodingType.IDENTITY) {
            final int threshold = (Integer) requestContext.getLocalAttr(HttpConstants.HEADER_RESPONSE_COMPRESSION_THRESHOLD);
            final StreamingCompressor compressor = selectedEncoding.getCompressor(_executor);
            PartialReader reader = new PartialReader(threshold, new Callback<EntityStream[]>() {

                @Override
                public void onError(Throwable ex) {
                    nextFilter.onError(ex, requestContext, wireAttrs);
                }

                @Override
                public void onSuccess(EntityStream[] results) {
                    if (// entity stream is less than threshold
                    results.length == 1) {
                        StreamResponse response = res.builder().build(results[0]);
                        nextFilter.onResponse(response, requestContext, wireAttrs);
                    } else {
                        EntityStream compressedStream = compressor.deflate(EntityStreams.newEntityStream(new CompositeWriter(results)));
                        StreamResponseBuilder builder = res.builder();
                        // remove original content-length header if presents.
                        if (builder.getHeader(HttpConstants.CONTENT_LENGTH) != null) {
                            Map<String, String> headers = stripHeaders(builder.getHeaders(), HttpConstants.CONTENT_LENGTH);
                            builder.setHeaders(headers);
                        }
                        StreamResponse response = builder.addHeaderValue(HttpConstants.CONTENT_ENCODING, compressor.getContentEncodingName()).build(compressedStream);
                        nextFilter.onResponse(response, requestContext, wireAttrs);
                    }
                }
            });
            res.getEntityStream().setReader(reader);
            return;
        }
    } catch (CompressionException e) {
        LOG.error(e.getMessage(), e.getCause());
    }
    nextFilter.onResponse(response, requestContext, wireAttrs);
}