Search in sources :

Example 51 with EntityStream

use of com.linkedin.r2.message.stream.entitystream.EntityStream 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);
}
Also used : StreamResponseBuilder(com.linkedin.r2.message.stream.StreamResponseBuilder) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) CompositeWriter(com.linkedin.r2.message.stream.entitystream.CompositeWriter) PartialReader(com.linkedin.r2.filter.compression.streaming.PartialReader) StreamEncodingType(com.linkedin.r2.filter.compression.streaming.StreamEncodingType) EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) Callback(com.linkedin.common.callback.Callback) StreamingCompressor(com.linkedin.r2.filter.compression.streaming.StreamingCompressor) AcceptEncoding(com.linkedin.r2.filter.compression.streaming.AcceptEncoding)

Example 52 with EntityStream

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

the class ServerStreamCompressionFilter method onStreamRequest.

/**
 * Handles compression tasks for incoming requests
 */
@Override
public void onStreamRequest(StreamRequest req, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<StreamRequest, StreamResponse> nextFilter) {
    try {
        // Check if the request is compressed, if so, decompress
        String requestContentEncoding = req.getHeader(HttpConstants.CONTENT_ENCODING);
        if (requestContentEncoding != null) {
            // This must be a specific compression type other than *
            StreamEncodingType encoding = StreamEncodingType.get(requestContentEncoding.trim().toLowerCase());
            if (encoding == null || encoding == StreamEncodingType.ANY) {
                throw new CompressionException(CompressionConstants.UNSUPPORTED_ENCODING + requestContentEncoding);
            }
            // Process the correct content-encoding types only
            StreamingCompressor compressor = encoding.getCompressor(_executor);
            if (compressor == null) {
                throw new CompressionException(CompressionConstants.UNKNOWN_ENCODING + encoding);
            }
            EntityStream uncompressedStream = compressor.inflate(req.getEntityStream());
            Map<String, String> headers = stripHeaders(req.getHeaders(), HttpConstants.CONTENT_ENCODING, HttpConstants.CONTENT_LENGTH);
            req = req.builder().setHeaders(headers).build(uncompressedStream);
        }
        // Get client support for compression and flag compress if need be
        String responseCompression = req.getHeader(HttpConstants.ACCEPT_ENCODING);
        if (responseCompression == null) {
            // per RFC 2616, section 14.3, if no Accept-Encoding field is present in a request,
            // server SHOULD use "identity" content-encoding if it is available.
            responseCompression = StreamEncodingType.IDENTITY.getHttpName();
        }
        if (!responseCompression.equalsIgnoreCase(StreamEncodingType.IDENTITY.getHttpName())) {
            requestContext.putLocalAttr(HttpConstants.HEADER_RESPONSE_COMPRESSION_THRESHOLD, _serverCompressionHelper.getResponseCompressionThreshold(req));
        }
        requestContext.putLocalAttr(HttpConstants.ACCEPT_ENCODING, responseCompression);
        nextFilter.onRequest(req, requestContext, wireAttrs);
    } catch (CompressionException ex) {
        LOG.error(ex.getMessage(), ex.getCause());
        StreamResponse streamResponse = new StreamResponseBuilder().setStatus(HttpConstants.UNSUPPORTED_MEDIA_TYPE).build(EntityStreams.emptyStream());
        nextFilter.onError(new StreamException(streamResponse, ex), requestContext, wireAttrs);
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) StreamResponseBuilder(com.linkedin.r2.message.stream.StreamResponseBuilder) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) StreamingCompressor(com.linkedin.r2.filter.compression.streaming.StreamingCompressor) StreamEncodingType(com.linkedin.r2.filter.compression.streaming.StreamEncodingType) StreamException(com.linkedin.r2.message.stream.StreamException)

Example 53 with EntityStream

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

the class TestMIMEReaderStateTransitions method testDrainAllPartsMultiPartMIMEReader.

@Test
public void testDrainAllPartsMultiPartMIMEReader() {
    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);
    // Test each possible exception:
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.FINISHED);
    try {
        reader.drainAllParts();
        Assert.fail();
    } catch (MultiPartReaderFinishedException multiPartReaderFinishedException) {
    }
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.READING_EPILOGUE);
    try {
        reader.drainAllParts();
        Assert.fail();
    } catch (MultiPartReaderFinishedException multiPartReaderFinishedException) {
    }
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.CALLBACK_BOUND_AND_READING_PREAMBLE);
    try {
        reader.drainAllParts();
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.DRAINING);
    try {
        reader.drainAllParts();
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
    // This is the desired top level reader state
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.READING_PARTS);
    final MultiPartMIMEReader.SinglePartMIMEReader singlePartMIMEReader = reader.new SinglePartMIMEReader(Collections.<String, String>emptyMap());
    // This is a undesired single part state
    singlePartMIMEReader.setState(MultiPartMIMEReader.SingleReaderState.REQUESTED_DATA);
    reader.setCurrentSinglePartMIMEReader(singlePartMIMEReader);
    try {
        reader.drainAllParts();
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) MultiPartReaderFinishedException(com.linkedin.multipart.exceptions.MultiPartReaderFinishedException) StreamBusyException(com.linkedin.multipart.exceptions.StreamBusyException) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) Test(org.testng.annotations.Test)

Example 54 with EntityStream

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

the class TestMIMEReaderStateTransitions method testRegisterSinglePartMIMEReaderCallbackTwice.

// /////////////////////////////////////////////////////////////////////////////////////
// SinglePartMIMEReader exceptions:
@Test
public void testRegisterSinglePartMIMEReaderCallbackTwice() {
    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());
    // This is a undesired single part state
    singlePartMIMEReader.setState(MultiPartMIMEReader.SingleReaderState.REQUESTED_DATA);
    try {
        singlePartMIMEReader.registerReaderCallback(EMPTY_SINGLE_PART_MIME_READER_CALLBACK);
        Assert.fail();
    } catch (SinglePartBindException singlePartBindException) {
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) SinglePartBindException(com.linkedin.multipart.exceptions.SinglePartBindException) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) Test(org.testng.annotations.Test)

Example 55 with EntityStream

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

the class TestMIMEReaderStateTransitions method testSinglePartMIMEReaderRequestData.

@Test
public void testSinglePartMIMEReaderRequestData() {
    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.CREATED);
    try {
        singlePartMIMEReader.requestPartData();
        Assert.fail();
    } catch (SinglePartNotInitializedException singlePartNotInitializedException) {
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) SinglePartNotInitializedException(com.linkedin.multipart.exceptions.SinglePartNotInitializedException) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) Test(org.testng.annotations.Test)

Aggregations

EntityStream (com.linkedin.r2.message.stream.entitystream.EntityStream)58 StreamResponse (com.linkedin.r2.message.stream.StreamResponse)29 Test (org.testng.annotations.Test)29 StreamRequest (com.linkedin.r2.message.stream.StreamRequest)27 CountDownLatch (java.util.concurrent.CountDownLatch)27 StreamRequestBuilder (com.linkedin.r2.message.stream.StreamRequestBuilder)22 ByteString (com.linkedin.data.ByteString)11 StreamResponseBuilder (com.linkedin.r2.message.stream.StreamResponseBuilder)10 FutureCallback (com.linkedin.common.callback.FutureCallback)9 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)9 AbstractServiceTest (test.r2.integ.clientserver.providers.AbstractServiceTest)9 AtomicReference (java.util.concurrent.atomic.AtomicReference)8 Callback (com.linkedin.common.callback.Callback)7 Client (com.linkedin.r2.transport.common.Client)7 BytesWriter (test.r2.integ.helper.BytesWriter)7 ScheduledExecutorService (java.util.concurrent.ScheduledExecutorService)6 TimeoutException (java.util.concurrent.TimeoutException)6 StreamingCompressor (com.linkedin.r2.filter.compression.streaming.StreamingCompressor)5 None (com.linkedin.common.util.None)4 StreamEncodingType (com.linkedin.r2.filter.compression.streaming.StreamEncodingType)4