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

Example 1 with ByteStringWriter

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

the class TestNettyRequestAdapter method testStreamToNettyRequest.

@Test
public void testStreamToNettyRequest() throws Exception {
    StreamRequestBuilder streamRequestBuilder = new StreamRequestBuilder(new URI(ANY_URI));
    streamRequestBuilder.setMethod("POST");
    streamRequestBuilder.setHeader("Content-Length", Integer.toString(ANY_ENTITY.length()));
    streamRequestBuilder.setHeader("Content-Type", "application/json");
    streamRequestBuilder.setCookies(Collections.singletonList("anyCookie"));
    StreamRequest streamRequest = streamRequestBuilder.build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(ANY_ENTITY.getBytes()))));
    HttpRequest nettyRequest = NettyRequestAdapter.toNettyRequest(streamRequest);
    Assert.assertEquals(nettyRequest.uri(), "/foo/bar?q=baz");
    Assert.assertEquals(nettyRequest.method(), HttpMethod.POST);
    Assert.assertEquals(nettyRequest.protocolVersion(), HttpVersion.HTTP_1_1);
    Assert.assertNull(nettyRequest.headers().get("Content-Length"));
    Assert.assertEquals(nettyRequest.headers().get("Content-Type"), "application/json");
    Assert.assertEquals(nettyRequest.headers().get("Cookie"), ANY_COOKIE);
}

Example 2 with ByteStringWriter

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

the class TestNettyRequestAdapter method testStreamToNettyRequestContentLengthIgnoreCase.

@Test
public void testStreamToNettyRequestContentLengthIgnoreCase() throws Exception {
    StreamRequestBuilder streamRequestBuilder = new StreamRequestBuilder(new URI(ANY_URI));
    streamRequestBuilder.setHeader("CONTENT-LENGTH", Integer.toString(ANY_ENTITY.length()));
    StreamRequest streamRequest = streamRequestBuilder.build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(ANY_ENTITY.getBytes()))));
    HttpRequest nettyRequest = NettyRequestAdapter.toNettyRequest(streamRequest);
    Assert.assertNull(nettyRequest.headers().get("Content-Length"));
}

Example 3 with ByteStringWriter

use of com.linkedin.r2.message.stream.entitystream.ByteStringWriter 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 4 with ByteStringWriter

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

the class RetryClientTest method testStreamRetry.

@Test
public void testStreamRetry() throws Exception {
    SimpleLoadBalancer balancer = prepareLoadBalancer(Arrays.asList("http://test.linkedin.com/retry1", "http://test.linkedin.com/good"), HttpClientFactory.UNLIMITED_CLIENT_REQUEST_RETRY_RATIO);
    DynamicClient dynamicClient = new DynamicClient(balancer, null);
    RetryClient client = new RetryClient(dynamicClient, balancer, D2ClientConfig.DEFAULT_RETRY_LIMIT, RetryClient.DEFAULT_UPDATE_INTERVAL_MS, RetryClient.DEFAULT_AGGREGATED_INTERVAL_NUM, SystemClock.instance(), true, true);
    URI uri = URI.create("d2://retryService?arg1arg2");
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.newEntityStream(new ByteStringWriter(CONTENT)));
    DegraderTrackerClientTest.TestCallback<StreamResponse> restCallback = new DegraderTrackerClientTest.TestCallback<>();
    client.streamRequest(streamRequest, restCallback);
    assertNull(restCallback.e);
    assertNotNull(restCallback.t);
}

Example 5 with ByteStringWriter

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

the class TestBackupRequestsClient method testStreamRequestWithNoIsFullRequest.

@Test(invocationCount = 3, dataProvider = "isD2Async")
public void testStreamRequestWithNoIsFullRequest(boolean isD2Async) throws Exception {
    // 1s till response comes back
    int responseDelayNano = 100000000;
    // make backup request after 0.5 second
    int backupDelayNano = 50000000;
    Deque<URI> hostsReceivingRequest = new ConcurrentLinkedDeque<>();
    BackupRequestsClient client = createAlwaysBackupClientWithHosts(Arrays.asList("http://test1.com:123", "http://test2.com:123"), hostsReceivingRequest, responseDelayNano, backupDelayNano, isD2Async);
    URI uri = URI.create("d2://testService");
    // if there is no IS_FULL_REQUEST set, backup requests will not happen
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.newEntityStream(new ByteStringWriter(CONTENT)));
    RequestContext context = new RequestContext();
    context.putLocalAttr(R2Constants.OPERATION, "get");
    RequestContext context1 = context.clone();
    CountDownLatch latch = new CountDownLatch(1);
    AtomicReference<AssertionError> failure = new AtomicReference<>();
    client.streamRequest(streamRequest, context1, new Callback<StreamResponse>() {

        @Override
        public void onError(Throwable e) {
            failure.set(new AssertionError("Callback onError"));
            latch.countDown();
        }

        @Override
        public void onSuccess(StreamResponse result) {
            try {
                assertEquals(result.getStatus(), 200);
                assertEquals(result.getHeader("buffered"), "false");
                assertEquals(hostsReceivingRequest.size(), 1);
                assertEquals(new HashSet<>(hostsReceivingRequest).size(), 1);
                hostsReceivingRequest.clear();
            } catch (AssertionError e) {
                failure.set(e);
            }
            latch.countDown();
        }
    });
    latch.await(2, TimeUnit.SECONDS);
    if (failure.get() != null) {
        throw failure.get();
    }
}