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

Example 16 with EntityStream

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

the class LoggingFilter method onStreamResponse.

@Override
public void onStreamResponse(StreamResponse res, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<StreamRequest, StreamResponse> nextFilter) {
    EntityStream entityStream = res.getEntityStream();
    entityStream.addObserver(new Observer() {

        private long startTime;

        private long bytesNum = 0;

        @Override
        public void onDataAvailable(ByteString data) {
            if (bytesNum == 0) {
                startTime = System.nanoTime();
            }
            bytesNum += data.length();
        }

        @Override
        public void onDone() {
            long stopTime = System.nanoTime();
            _log.info("Status: success. Total bytes streamed: " + bytesNum + ". Total stream time: " + (stopTime - startTime) + " nano seconds.");
        }

        @Override
        public void onError(Throwable e) {
            long stopTime = System.nanoTime();
            _log.error("Status: failed. Total bytes streamed: " + bytesNum + ". Total stream time before failure: " + (stopTime - startTime) + " nano seconds.");
        }
    });
}

Example 17 with EntityStream

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

the class TestMIMEReaderStateTransitions method testSinglePartMIMEReaderVerifyState.

@Test
public void testSinglePartMIMEReaderVerifyState() {
    // This will cover drainPart() and most of requestPartData().
    // The caveat is that requestPartData() requires a callback to be registered. This
    // will be covered in the next test.
    final EntityStream entityStream = mock(EntityStream.class);
    final StreamRequest streamRequest = mock(StreamRequest.class);
    when(streamRequest.getEntityStream()).thenReturn(entityStream);
    when(streamRequest.getHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER)).thenReturn("multipart/mixed; boundary=\"--123\"");
    MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    final MultiPartMIMEReader.SinglePartMIMEReader singlePartMIMEReader = reader.new SinglePartMIMEReader(Collections.<String, String>emptyMap());
    singlePartMIMEReader.setState(MultiPartMIMEReader.SingleReaderState.FINISHED);
    try {
        singlePartMIMEReader.verifyUsableState();
        Assert.fail();
    } catch (SinglePartFinishedException singlePartFinishedException) {
    }
    singlePartMIMEReader.setState(MultiPartMIMEReader.SingleReaderState.REQUESTED_DATA);
    try {
        singlePartMIMEReader.verifyUsableState();
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
    singlePartMIMEReader.setState(MultiPartMIMEReader.SingleReaderState.REQUESTED_DRAIN);
    try {
        singlePartMIMEReader.verifyUsableState();
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
}

Example 18 with EntityStream

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

the class TestStreamingCompression method testDecompress.

private void testDecompress(StreamingCompressor compressor, byte[] uncompressed, byte[] compressed) throws CompressionException, ExecutionException, InterruptedException {
    ByteWriter writer = new ByteWriter(compressed);
    EntityStream compressedStream = EntityStreams.newEntityStream(writer);
    EntityStream uncompressedStream = compressor.inflate(compressedStream);
    FutureCallback<byte[]> callback = new FutureCallback<>();
    uncompressedStream.setReader(new ByteReader(callback));
    byte[] result = callback.get();
    Assert.assertEquals(result, uncompressed);
}

Example 19 with EntityStream

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

the class TestStreamingCompression method testCompressThenDecompress.

private void testCompressThenDecompress(StreamingCompressor compressor, byte[] origin) throws CompressionException, ExecutionException, InterruptedException {
    ByteWriter writer = new ByteWriter(origin);
    EntityStream uncompressedStream = EntityStreams.newEntityStream(writer);
    EntityStream compressedStream = compressor.deflate(uncompressedStream);
    EntityStream decompressedStream = compressor.inflate(compressedStream);
    FutureCallback<byte[]> callback = new FutureCallback<>();
    decompressedStream.setReader(new ByteReader(callback));
    byte[] result = callback.get();
    Assert.assertEquals(result, origin);
}

Example 20 with EntityStream

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