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

Example 86 with StreamRequest

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

the class TestServerRetryFilter method testStreamRetryFilter.

@Test
public void testStreamRetryFilter() {
    String retryMessage = "this is a retry";
    ServerRetryFilter retryFilter = new ServerRetryFilter();
    StreamFilter captureFilter = new StreamFilter() {

        @Override
        public void onStreamError(Throwable ex, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<StreamRequest, StreamResponse> nextFilter) {
            Assert.assertEquals(wireAttrs.get(R2Constants.RETRY_MESSAGE_ATTRIBUTE_KEY), retryMessage);
        }
    };
    FilterChain filterChain = FilterChains.createStreamChain(captureFilter, retryFilter);
    FilterUtil.fireRestError(filterChain, new StreamException(null, new RetriableRequestException(retryMessage)), new HashMap<>());
}

Example 87 with StreamRequest

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

the class TestDisruptFilter method testStreamErrorDisrupt.

@Test
public void testStreamErrorDisrupt() throws Exception {
    final RequestContext requestContext = new RequestContext();
    requestContext.putLocalAttr(DISRUPT_CONTEXT_KEY, DisruptContexts.error(REQUEST_LATENCY));
    final DisruptFilter filter = new DisruptFilter(_scheduler, _executor, REQUEST_TIMEOUT, _clock);
    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) {
            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) {
            success.set(ex instanceof DisruptedException);
            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 88 with StreamRequest

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

the class TestDisruptFilter method testExecutorRejectExecution.

@Test
public void testExecutorRejectExecution() throws Exception {
    final AtomicBoolean success = new AtomicBoolean(false);
    final CountDownLatch latch = new CountDownLatch(1);
    ExecutorService rejectedExecutor = EasyMock.createStrictMock(ExecutorService.class);
    rejectedExecutor.execute(EasyMock.anyObject(Runnable.class));
    EasyMock.expectLastCall().andAnswer(() -> {
        success.set(true);
        latch.countDown();
        throw new RejectedExecutionException();
    });
    EasyMock.replay(rejectedExecutor);
    final RequestContext requestContext = new RequestContext();
    requestContext.putLocalAttr(DISRUPT_CONTEXT_KEY, DisruptContexts.error(REQUEST_LATENCY));
    final DisruptFilter filter = new DisruptFilter(_scheduler, rejectedExecutor, REQUEST_TIMEOUT, _clock);
    final NextFilter<StreamRequest, StreamResponse> next = new NextFilter<StreamRequest, StreamResponse>() {

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

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

        @Override
        public void onError(Throwable ex, RequestContext requestContext, Map<String, String> wireAttrs) {
            success.set(false);
            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");
    EasyMock.verify(rejectedExecutor);
    EasyMock.reset(rejectedExecutor);
}

Example 89 with StreamRequest

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

the class ClientStreamCompressionFilter method onStreamRequest.

/**
 * Optionally compresses outgoing Stream requests
 */
public void onStreamRequest(StreamRequest req, final RequestContext requestContext, final Map<String, String> wireAttrs, final NextFilter<StreamRequest, StreamResponse> nextFilter) {
    // Set accepted encoding for compressed response
    String operation = (String) requestContext.getLocalAttr(R2Constants.OPERATION);
    if (!_acceptEncodingHeader.isEmpty() && _helper.shouldCompressResponseForOperation(operation)) {
        CompressionOption responseCompressionOverride = (CompressionOption) requestContext.getLocalAttr(R2Constants.RESPONSE_COMPRESSION_OVERRIDE);
        req = addResponseCompressionHeaders(responseCompressionOverride, req);
    }
    if (_requestContentEncoding != StreamEncodingType.IDENTITY) {
        final StreamRequest request = req;
        final StreamingCompressor compressor = _requestContentEncoding.getCompressor(_executor);
        CompressionOption option = (CompressionOption) requestContext.getLocalAttr(R2Constants.REQUEST_COMPRESSION_OVERRIDE);
        if (option == null || option != CompressionOption.FORCE_OFF) {
            final int threshold = option == CompressionOption.FORCE_ON ? 0 : _requestCompressionConfig.getCompressionThreshold();
            PartialReader reader = new PartialReader(threshold, new Callback<EntityStream[]>() {

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

                @Override
                public void onSuccess(EntityStream[] result) {
                    if (result.length == 1) {
                        StreamRequest uncompressedRequest = request.builder().build(result[0]);
                        nextFilter.onRequest(uncompressedRequest, requestContext, wireAttrs);
                    } else {
                        StreamRequestBuilder builder = request.builder();
                        EntityStream compressedStream = compressor.deflate(EntityStreams.newEntityStream(new CompositeWriter(result)));
                        Map<String, String> headers = stripHeaders(builder.getHeaders(), HttpConstants.CONTENT_LENGTH);
                        StreamRequest compressedRequest = builder.setHeaders(headers).setHeader(HttpConstants.CONTENT_ENCODING, compressor.getContentEncodingName()).build(compressedStream);
                        nextFilter.onRequest(compressedRequest, requestContext, wireAttrs);
                    }
                }
            });
            req.getEntityStream().setReader(reader);
            return;
        }
    }
    nextFilter.onRequest(req, requestContext, wireAttrs);
}

Example 90 with StreamRequest

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

the class ClientStreamCompressionFilter method onStreamResponse.

/**
 *  Decompresses server response
 */
@Override
public void onStreamResponse(StreamResponse res, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<StreamRequest, StreamResponse> nextFilter) {
    Boolean decompressionOff = (Boolean) requestContext.getLocalAttr(R2Constants.RESPONSE_DECOMPRESSION_OFF);
    if (decompressionOff == null || !decompressionOff) {
        // Check for header encoding
        String compressionHeader = res.getHeader(HttpConstants.CONTENT_ENCODING);
        // decompress if necessary
        if (compressionHeader != null) {
            final StreamEncodingType encoding = StreamEncodingType.get(compressionHeader.trim().toLowerCase());
            if (encoding == null) {
                nextFilter.onError(new IllegalArgumentException("Server returned unrecognized content encoding: " + compressionHeader), requestContext, wireAttrs);
                return;
            }
            final StreamingCompressor compressor = encoding.getCompressor(_executor);
            EntityStream uncompressedStream = compressor.inflate(res.getEntityStream());
            StreamResponseBuilder builder = res.builder();
            Map<String, String> headers = stripHeaders(builder.getHeaders(), HttpConstants.CONTENT_ENCODING, HttpConstants.CONTENT_LENGTH);
            res = builder.setHeaders(headers).build(uncompressedStream);
        }
    }
    nextFilter.onResponse(res, requestContext, wireAttrs);
}