use of com.linkedin.r2.message.stream.StreamRequestBuilder in project rest.li by linkedin.
the class TestStreamEcho method testNormalEcho.
private void testNormalEcho(long bytesNum, URI uri) throws Exception {
for (Client client : clients()) {
BytesWriter writer = new BytesWriter(bytesNum, BYTE);
StreamRequest request = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, uri)).build(EntityStreams.newEntityStream(writer));
final AtomicInteger status = new AtomicInteger(-1);
final CountDownLatch latch = new CountDownLatch(1);
final AtomicReference<Throwable> error = new AtomicReference<Throwable>();
final Callback<None> readerCallback = getReaderCallback(latch, error);
final BytesReader reader = new BytesReader(BYTE, readerCallback);
Callback<StreamResponse> callback = getCallback(status, readerCallback, reader);
client.streamRequest(request, callback);
latch.await(60000, TimeUnit.MILLISECONDS);
Assert.assertNull(error.get());
Assert.assertEquals(status.get(), RestStatus.OK);
Assert.assertEquals(reader.getTotalBytes(), bytesNum);
Assert.assertTrue(reader.allBytesCorrect());
}
}
use of com.linkedin.r2.message.stream.StreamRequestBuilder in project rest.li by linkedin.
the class TestStreamRequest method testErrorWriter.
@Test
public void testErrorWriter() throws Exception {
for (Client client : clients()) {
final long totalBytes = SMALL_BYTES_NUM;
EntityStream entityStream = EntityStreams.newEntityStream(new ErrorWriter(totalBytes, BYTE));
StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, FOOBAR_URI));
StreamRequest request = builder.setMethod("POST").build(entityStream);
final CountDownLatch latch = new CountDownLatch(1);
final AtomicReference<Throwable> error = new AtomicReference<Throwable>();
Callback<StreamResponse> callback = new Callback<StreamResponse>() {
@Override
public void onError(Throwable e) {
error.set(e);
latch.countDown();
}
@Override
public void onSuccess(StreamResponse result) {
latch.countDown();
}
};
client.streamRequest(request, callback);
latch.await();
Assert.assertNotNull(error.get());
}
}
use of com.linkedin.r2.message.stream.StreamRequestBuilder in project rest.li by linkedin.
the class TestRestLiServer method testBadRequest.
private void testBadRequest(RestLiServer restLiServer, final ProtocolVersion clientProtocolVersion, String headerConstant, final RestOrStream restOrStream) throws URISyntaxException {
Callback<RestResponse> restResponseCallback = new Callback<RestResponse>() {
@Override
public void onSuccess(RestResponse restResponse) {
fail("The request should have failed!");
}
@Override
public void onError(Throwable e) {
assertEquals(((RestException) e).getResponse().getStatus(), 400);
String expectedErrorMessage = "Rest.li protocol version " + clientProtocolVersion + " used by the client is not supported!";
assertEquals(e.getCause().getMessage(), expectedErrorMessage);
}
};
if (restOrStream == RestOrStream.REST) {
RestRequest request = new RestRequestBuilder(new URI("/statuses/1")).setHeader(headerConstant, clientProtocolVersion.toString()).build();
restLiServer.handleRequest(request, new RequestContext(), restResponseCallback);
} else {
StreamRequest streamRequest = new StreamRequestBuilder(new URI("/statuses/1")).setHeader(headerConstant, clientProtocolVersion.toString()).build(EntityStreams.emptyStream());
Callback<StreamResponse> callback = new Callback<StreamResponse>() {
@Override
public void onSuccess(StreamResponse streamResponse) {
fail("The request should have failed!");
}
@Override
public void onError(Throwable e) {
Messages.toRestException((StreamException) e, new Callback<RestException>() {
@Override
public void onError(Throwable e) {
Assert.fail();
}
@Override
public void onSuccess(RestException result) {
restResponseCallback.onError(result);
}
});
}
};
restLiServer.handleRequest(streamRequest, new RequestContext(), callback);
}
}
use of com.linkedin.r2.message.stream.StreamRequestBuilder in project rest.li by linkedin.
the class SimpleLoadBalancerStateTest method testClientsShutdownAfterPropertyUpdatesStreamRequest.
@Test(groups = { "small", "back-end" })
public void testClientsShutdownAfterPropertyUpdatesStreamRequest() throws URISyntaxException, InterruptedException {
reset();
URI uri = URI.create("http://cluster-1/test");
List<String> schemes = new ArrayList<>();
Map<Integer, PartitionData> partitionData = new HashMap<>(1);
partitionData.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1d));
Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<>();
uriData.put(uri, partitionData);
schemes.add("http");
// set up state
_state.listenToService("service-1", new NullStateListenerCallback());
_state.listenToCluster("cluster-1", new NullStateListenerCallback());
_state.setDelayedExecution(0);
_serviceRegistry.put("service-1", new ServiceProperties("service-1", "cluster-1", "/test", Arrays.asList("random"), Collections.<String, Object>emptyMap(), Collections.<String, Object>emptyMap(), Collections.<String, String>emptyMap(), schemes, Collections.<URI>emptySet()));
_clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1"));
_uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
URI uri1 = URI.create("http://partition-cluster-1/test1");
URI uri2 = URI.create("http://partition-cluster-1/test2");
_state.listenToCluster("partition-cluster-1", new NullStateListenerCallback());
_clusterRegistry.put("partition-cluster-1", new ClusterProperties("partition-cluster-1", null, new HashMap<>(), new HashSet<>(), new RangeBasedPartitionProperties("id=(\\d+)", 0, 100, 2)));
_state.listenToService("partition-service-1", new NullStateListenerCallback());
_serviceRegistry.put("partition-service-1", new ServiceProperties("partition-service-1", "partition-cluster-1", "/partition-test", Arrays.asList("degraderV3"), Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap(), schemes, Collections.emptySet()));
Map<Integer, PartitionData> partitionWeight = new HashMap<>();
partitionWeight.put(0, new PartitionData(1d));
partitionWeight.put(1, new PartitionData(2d));
Map<URI, Map<Integer, PartitionData>> partitionDesc = new HashMap<>();
partitionDesc.put(uri1, partitionWeight);
partitionWeight.remove(0);
partitionWeight.put(2, new PartitionData(1d));
partitionDesc.put(uri2, partitionWeight);
_uriRegistry.put("partition-cluster-1", new UriProperties("partition-cluster-1", partitionDesc));
TrackerClient client1 = _state.getClient("partition-service-1", uri1);
TrackerClient client2 = _state.getClient("partition-service-1", uri2);
assertEquals(client2.getPartitionWeight(1), 2d);
assertEquals(client2.getPartitionWeight(2), 1d);
assertEquals(client1.getPartitionWeight(1), 2d);
// Get client, then refresh cluster
TrackerClient client = _state.getClient("service-1", uri);
client.streamRequest(new StreamRequestBuilder(URI.create("d2://service-1/foo")).build(EntityStreams.emptyStream()), new RequestContext(), Collections.emptyMap(), new TransportCallbackAdapter<>(Callbacks.empty()));
// now force a refresh by adding cluster
_clusterRegistry.put("cluster-1", new ClusterProperties("cluster-1"));
// Get client, then refresh service
client = _state.getClient("service-1", uri);
client.streamRequest(new StreamRequestBuilder(URI.create("d2://service-1/foo")).build(EntityStreams.emptyStream()), new RequestContext(), Collections.emptyMap(), new TransportCallbackAdapter<>(Callbacks.empty()));
// refresh by adding service
_serviceRegistry.put("service-1", new ServiceProperties("service-1", "cluster-1", "/test", Arrays.asList("random"), Collections.<String, Object>emptyMap(), null, null, schemes, null));
// Get client, then mark server up/down
client = _state.getClient("service-1", uri);
client.streamRequest(new StreamRequestBuilder(URI.create("d2://service-1/foo")).build(EntityStreams.emptyStream()), new RequestContext(), Collections.emptyMap(), new TransportCallbackAdapter<>(Callbacks.empty()));
_uriRegistry.put("cluster-1", new UriProperties("cluster-1", Collections.<URI, Map<Integer, PartitionData>>emptyMap()));
_uriRegistry.put("cluster-1", new UriProperties("cluster-1", uriData));
// Get the client one last time
client = _state.getClient("service-1", uri);
client.streamRequest(new StreamRequestBuilder(URI.create("d2://service-1/foo")).build(EntityStreams.emptyStream()), new RequestContext(), Collections.emptyMap(), new TransportCallbackAdapter<>(Callbacks.empty()));
TestShutdownCallback callback = new TestShutdownCallback();
_state.shutdown(callback);
assertTrue(callback.await(10, TimeUnit.SECONDS), "Failed to shut down state");
for (TransportClientFactory factory : _clientFactories.values()) {
SimpleLoadBalancerTest.DoNothingClientFactory f = (SimpleLoadBalancerTest.DoNothingClientFactory) factory;
assertEquals(f.getRunningClientCount(), 0, "not all clients were shut down");
}
}
use of com.linkedin.r2.message.stream.StreamRequestBuilder in project rest.li by linkedin.
the class TestMIMEIntegrationReaderDrain method executeRequestWithDrainStrategy.
// /////////////////////////////////////////////////////////////////////////////////////
private MimeMultipart executeRequestWithDrainStrategy(final int chunkSize, final List<MimeBodyPart> bodyPartList, final String drainStrategy, final String serverHeaderPrefix) throws Exception {
MimeMultipart multiPartMimeBody = new MimeMultipart();
// Add your body parts
for (final MimeBodyPart bodyPart : bodyPartList) {
multiPartMimeBody.addBodyPart(bodyPart);
}
final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
multiPartMimeBody.writeTo(byteArrayOutputStream);
final ByteString requestPayload = ByteString.copy(byteArrayOutputStream.toByteArray());
final VariableByteStringWriter variableByteStringWriter = new VariableByteStringWriter(requestPayload, chunkSize);
final EntityStream entityStream = EntityStreams.newEntityStream(variableByteStringWriter);
final StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, SERVER_URI));
StreamRequest request = builder.setMethod("POST").setHeader(HEADER_CONTENT_TYPE, multiPartMimeBody.getContentType()).setHeader(DRAIN_HEADER, drainStrategy).build(entityStream);
final AtomicInteger status = new AtomicInteger(-1);
final CountDownLatch latch = new CountDownLatch(1);
final Map<String, String> responseHeaders = new HashMap<>();
Callback<StreamResponse> callback = expectSuccessCallback(latch, status, responseHeaders);
_client.streamRequest(request, callback);
latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
Assert.assertEquals(status.get(), RestStatus.OK);
Assert.assertEquals(responseHeaders.get(DRAIN_HEADER), serverHeaderPrefix + drainStrategy);
return multiPartMimeBody;
}
Aggregations