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

Example 16 with StreamRequest

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

the class TestResponseCompression method testEncodingNotAcceptable.

public void testEncodingNotAcceptable(String acceptEncoding) throws TimeoutException, InterruptedException {
    for (Client client : clients()) {
        StreamRequestBuilder builder = new StreamRequestBuilder((Bootstrap.createHttpURI(PORT, SMALL_URI)));
        if (acceptEncoding != null) {
            builder.addHeaderValue(HttpConstants.ACCEPT_ENCODING, acceptEncoding);
        }
        StreamRequest request = builder.build(EntityStreams.emptyStream());
        final FutureCallback<StreamResponse> callback = new FutureCallback<StreamResponse>();
        client.streamRequest(request, callback);
        try {
            final StreamResponse response = callback.get(60, TimeUnit.SECONDS);
            Assert.fail("Should have thrown exception when encoding is not acceptable");
        } catch (ExecutionException e) {
            Throwable t = e.getCause();
            Assert.assertTrue(t instanceof StreamException);
            StreamResponse response = ((StreamException) t).getResponse();
            Assert.assertEquals(response.getStatus(), HttpConstants.NOT_ACCEPTABLE);
        }
    }
}

Example 17 with StreamRequest

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

the class TestStreamRequest method testBackPressure.

@Test
public void testBackPressure() throws Exception {
    for (Client client : clients()) {
        final long totalBytes = SMALL_BYTES_NUM;
        TimedBytesWriter writer = new TimedBytesWriter(totalBytes, BYTE);
        EntityStream entityStream = EntityStreams.newEntityStream(writer);
        StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, RATE_LIMITED_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);
        TimedBytesReader reader = _rateLimitedRequestHandler.getReader();
        Assert.assertNotNull(reader);
        Assert.assertEquals(totalBytes, reader.getTotalBytes());
        Assert.assertTrue(reader.allBytesCorrect());
        long clientSendTimespan = writer.getStopTime() - writer.getStartTime();
        long serverReceiveTimespan = reader.getStopTime() - reader.getStartTime();
        Assert.assertTrue(serverReceiveTimespan > 1000);
        double diff = Math.abs(serverReceiveTimespan - clientSendTimespan);
        double diffRatio = diff / clientSendTimespan;
        // make it generous to reduce the chance occasional test failures
        Assert.assertTrue(diffRatio < 0.2);
    }
}

Example 18 with StreamRequest

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

the class RestClient method buildStreamRequest.

private StreamRequest buildStreamRequest(URI uri, ResourceMethod method, DataMap dataMap, Map<String, String> headers, List<String> cookies, ProtocolVersion protocolVersion, ContentType contentType, List<AcceptType> acceptTypes, boolean acceptResponseAttachments, List<Object> streamingAttachments) throws Exception {
    StreamRequestBuilder requestBuilder = new StreamRequestBuilder(uri).setMethod(method.getHttpMethod().toString());
    requestBuilder.setHeaders(headers);
    requestBuilder.setCookies(cookies);
    addAcceptHeaders(requestBuilder, acceptTypes, acceptResponseAttachments);
    addProtocolVersionHeader(requestBuilder, protocolVersion);
    if (method.getHttpMethod() == HttpMethod.POST) {
        requestBuilder.setHeader(RestConstants.HEADER_RESTLI_REQUEST_METHOD, method.toString());
    }
    //This request builders enforce this invariant.
    if (streamingAttachments != null) {
        final ByteStringWriter firstPartWriter;
        final ContentType type = resolveContentType(requestBuilder, dataMap, contentType);
        //with empty action parameters will have an empty JSON ({}) as the body.
        assert (type != null);
        switch(type) {
            case PSON:
                firstPartWriter = new ByteStringWriter(ByteString.copy(PSON_DATA_CODEC.mapToBytes(dataMap)));
                break;
            case JSON:
                firstPartWriter = new ByteStringWriter(ByteString.copy(JACKSON_DATA_CODEC.mapToBytes(dataMap)));
                break;
            default:
                throw new IllegalStateException("Unknown ContentType:" + type);
        }
        //Our protocol does not use an epilogue or a preamble.
        final MultiPartMIMEWriter.Builder attachmentsBuilder = new MultiPartMIMEWriter.Builder();
        for (final Object dataSource : streamingAttachments) {
            assert (dataSource instanceof RestLiAttachmentDataSourceWriter || dataSource instanceof RestLiDataSourceIterator);
            if (dataSource instanceof RestLiAttachmentDataSourceWriter) {
                AttachmentUtils.appendSingleAttachmentToBuilder(attachmentsBuilder, (RestLiAttachmentDataSourceWriter) dataSource);
            } else {
                AttachmentUtils.appendMultipleAttachmentsToBuilder(attachmentsBuilder, (RestLiDataSourceIterator) dataSource);
            }
        }
        final MultiPartMIMEWriter multiPartMIMEWriter = AttachmentUtils.createMultiPartMIMEWriter(firstPartWriter, type.getHeaderKey(), attachmentsBuilder);
        final String contentTypeHeader = MultiPartMIMEUtils.buildMIMEContentTypeHeader(AttachmentUtils.RESTLI_MULTIPART_SUBTYPE, multiPartMIMEWriter.getBoundary(), Collections.<String, String>emptyMap());
        requestBuilder.setHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER, contentTypeHeader);
        return requestBuilder.build(multiPartMIMEWriter.getEntityStream());
    } else {
        //taken the RestRequest code path.
        assert (acceptResponseAttachments == true);
        return Messages.toStreamRequest(buildRestRequest(uri, method, dataMap, headers, cookies, protocolVersion, contentType, acceptTypes, acceptResponseAttachments));
    }
}

Example 19 with StreamRequest

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

the class TestRestClientRequestBuilder method testUpdate.

@Test(dataProvider = "data")
public void testUpdate(RestClient.ContentType contentType, String expectedContentTypeHeader, String expectedRequestBody, String expectedEntitiesBody, List<RestClient.AcceptType> acceptTypes, String expectedAcceptHeader, boolean acceptContentTypePerClient, boolean streamAttachments, boolean acceptResponseAttachments) throws URISyntaxException {
    //its a RestRequest.
    if (streamAttachments == false && acceptResponseAttachments == false) {
        RestRequest restRequest = clientGeneratedRestRequest(UpdateRequest.class, ResourceMethod.UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient);
        Assert.assertEquals(restRequest.getHeader(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        Assert.assertEquals(restRequest.getEntity().asAvroString(), expectedRequestBody);
        Assert.assertEquals(restRequest.getHeader(ACCEPT_TYPE_HEADER), expectedAcceptHeader);
        RestRequest restRequestBatch = clientGeneratedRestRequest(BatchUpdateRequest.class, ResourceMethod.BATCH_UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient);
        Assert.assertEquals(restRequestBatch.getHeader(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        Assert.assertEquals(restRequestBatch.getEntity().asAvroString(), expectedEntitiesBody);
        Assert.assertEquals(restRequestBatch.getHeader(ACCEPT_TYPE_HEADER), expectedAcceptHeader);
        RestRequest restRequestPartial = clientGeneratedRestRequest(PartialUpdateRequest.class, ResourceMethod.PARTIAL_UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient);
        Assert.assertEquals(restRequestPartial.getHeader(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        Assert.assertEquals(restRequestPartial.getEntity().asAvroString(), expectedRequestBody);
        Assert.assertEquals(restRequestPartial.getHeader(ACCEPT_TYPE_HEADER), expectedAcceptHeader);
        RestRequest restRequestBatchPartial = clientGeneratedRestRequest(BatchPartialUpdateRequest.class, ResourceMethod.BATCH_PARTIAL_UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient);
        Assert.assertEquals(restRequestBatchPartial.getHeader(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        Assert.assertEquals(restRequestBatchPartial.getEntity().asAvroString(), expectedEntitiesBody);
        Assert.assertEquals(restRequestBatchPartial.getHeader(ACCEPT_TYPE_HEADER), expectedAcceptHeader);
    } else {
        StreamRequest streamRequest = clientGeneratedStreamRequest(UpdateRequest.class, ResourceMethod.UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient, streamAttachments ? generateRequestAttachments() : null, acceptResponseAttachments);
        verifyStreamRequest(streamRequest, acceptResponseAttachments, expectedAcceptHeader, streamAttachments, expectedContentTypeHeader, expectedRequestBody);
        StreamRequest streamRequestBatch = clientGeneratedStreamRequest(BatchUpdateRequest.class, ResourceMethod.BATCH_UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient, streamAttachments ? generateRequestAttachments() : null, acceptResponseAttachments);
        verifyStreamRequest(streamRequestBatch, acceptResponseAttachments, expectedAcceptHeader, streamAttachments, expectedContentTypeHeader, expectedEntitiesBody);
        StreamRequest streamRequestPartial = clientGeneratedStreamRequest(PartialUpdateRequest.class, ResourceMethod.PARTIAL_UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient, streamAttachments ? generateRequestAttachments() : null, acceptResponseAttachments);
        verifyStreamRequest(streamRequestPartial, acceptResponseAttachments, expectedAcceptHeader, streamAttachments, expectedContentTypeHeader, expectedRequestBody);
        StreamRequest streamRequestBatchPartial = clientGeneratedStreamRequest(BatchPartialUpdateRequest.class, ResourceMethod.BATCH_PARTIAL_UPDATE, ENTITY_BODY, contentType, acceptTypes, acceptContentTypePerClient, streamAttachments ? generateRequestAttachments() : null, acceptResponseAttachments);
        verifyStreamRequest(streamRequestBatchPartial, acceptResponseAttachments, expectedAcceptHeader, streamAttachments, expectedContentTypeHeader, expectedEntitiesBody);
    }
}

Example 20 with StreamRequest

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

the class TestRestClientRequestBuilder method clientGeneratedStreamRequest.

//This is similar to clientGeneratedRestRequest above except that it will generate a StreamRequest instead
//of a RestRequest. Note that this will ONLY happen if either acceptResponseAttachments below is 'true' OR
//streamingAttachmentDataSources below is non-null with a size greater then 0. If both of these do not hold,
//then a StreamRequest will not be generated by the RestClient.
@SuppressWarnings({ "unchecked", "rawtypes", "deprecation" })
private <T extends Request> StreamRequest clientGeneratedStreamRequest(Class<T> requestClass, ResourceMethod method, DataMap entityBody, RestClient.ContentType contentType, List<RestClient.AcceptType> acceptTypes, boolean acceptContentTypePerClient, List<Object> streamingAttachmentDataSources, boolean acceptResponseAttachments) throws URISyntaxException {
    // massive setup...
    Client mockClient = EasyMock.createMock(Client.class);
    @SuppressWarnings({ "rawtypes" }) Request<?> mockRequest = EasyMock.createMock(requestClass);
    RecordTemplate mockRecordTemplate = EasyMock.createMock(RecordTemplate.class);
    @SuppressWarnings({ "rawtypes" }) RestResponseDecoder mockResponseDecoder = EasyMock.createMock(RestResponseDecoder.class);
    RestliRequestOptions requestOptions = RestliRequestOptions.DEFAULT_OPTIONS;
    //If there is a desire to receive response attachments, then we must use request options.
    if (!acceptContentTypePerClient || acceptResponseAttachments) {
        requestOptions = new RestliRequestOptions(ProtocolVersionOption.USE_LATEST_IF_AVAILABLE, null, null, contentType, acceptTypes, acceptResponseAttachments);
    }
    setCommonExpectations(mockRequest, method, mockResponseDecoder, requestOptions);
    if (streamingAttachmentDataSources != null && streamingAttachmentDataSources.size() > 0) {
        EasyMock.expect(mockRequest.getStreamingAttachments()).andReturn(streamingAttachmentDataSources).times(2);
    } else {
        EasyMock.expect(mockRequest.getStreamingAttachments()).andReturn(null).times(2);
    }
    setResourceMethodExpectations(method, mockRequest, mockRecordTemplate, entityBody);
    Capture<StreamRequest> streamRequestCapture = new Capture<StreamRequest>();
    EasyMock.expect(mockClient.getMetadata(new URI(HOST + SERVICE_NAME))).andReturn(Collections.<String, Object>emptyMap()).once();
    mockClient.streamRequest(EasyMock.capture(streamRequestCapture), (RequestContext) EasyMock.anyObject(), (Callback<StreamResponse>) EasyMock.anyObject());
    EasyMock.expectLastCall().once();
    EasyMock.replay(mockClient, mockRequest, mockRecordTemplate);
    // do work!
    RestClient restClient;
    if (acceptContentTypePerClient) {
        // configuration per client
        restClient = new RestClient(mockClient, HOST, contentType, acceptTypes);
    } else {
        // configuration per request
        restClient = new RestClient(mockClient, HOST);
    }
    restClient.sendRequest(mockRequest);
    return streamRequestCapture.getValue();
}