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

Example 66 with StreamResponse

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

the class TestStreamResponse method testBackpressure.

@Test
public void testBackpressure() throws Exception {
    for (Client client : clients()) {
        StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, SMALL_URI));
        StreamRequest request = builder.build(EntityStreams.emptyStream());
        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 TimedBytesReader reader = new TimedBytesReader(BYTE, readerCallback) {

            int count = 0;

            @Override
            protected void requestMore(final ReadHandle rh) {
                count++;
                if (count % 16 == 0) {
                    _scheduler.schedule(new Runnable() {

                        @Override
                        public void run() {
                            rh.request(1);
                        }
                    }, INTERVAL, TimeUnit.MILLISECONDS);
                } else {
                    rh.request(1);
                }
            }
        };
        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);
        long serverSendTimespan = _smallHandler.getWriter().getStopTime() - _smallHandler.getWriter().getStartTime();
        long clientReceiveTimespan = reader.getStopTime() - reader.getStartTime();
        Assert.assertTrue(clientReceiveTimespan > 1000);
        double diff = Math.abs(clientReceiveTimespan - serverSendTimespan);
        double diffRatio = diff / serverSendTimespan;
        // make it generous to reduce the chance occasional test failures
        Assert.assertTrue(diffRatio < 0.2);
    }
}

Example 67 with StreamResponse

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

the class TestStreamResponse method testResponse.

private void testResponse(URI uri) throws Exception {
    for (Client client : clients()) {
        StreamRequestBuilder builder = new StreamRequestBuilder(uri);
        StreamRequest request = builder.build(EntityStreams.emptyStream());
        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(), LARGE_BYTES_NUM);
        Assert.assertTrue(reader.allBytesCorrect());
    }
}

Example 68 with StreamResponse

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

the class TestStreamResponse method testErrorWhileStreaming.

@Test
public void testErrorWhileStreaming() throws Exception {
    HttpClientFactory clientFactory = new HttpClientFactory();
    Map<String, String> clientProperties = new HashMap<String, String>();
    clientProperties.put(HttpClientFactory.HTTP_REQUEST_TIMEOUT, "1000");
    Client client = new TransportClientAdapter(_clientFactory.getClient(clientProperties), true);
    StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, SERVER_ERROR_URI));
    StreamRequest request = builder.build(EntityStreams.emptyStream());
    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(2000, TimeUnit.MILLISECONDS);
    Assert.assertEquals(status.get(), RestStatus.OK);
    Throwable throwable = error.get();
    Assert.assertNotNull(throwable);
    final FutureCallback<None> clientShutdownCallback = new FutureCallback<None>();
    client.shutdown(clientShutdownCallback);
    clientShutdownCallback.get();
    final FutureCallback<None> factoryShutdownCallback = new FutureCallback<None>();
    clientFactory.shutdown(factoryShutdownCallback);
    factoryShutdownCallback.get();
}

Example 69 with StreamResponse

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

the class TestChannelPoolStreamHandler method getChannel.

private static EmbeddedChannel getChannel() {
    EmbeddedChannel ch = new EmbeddedChannel(new RAPResponseDecoder(1000), new RAPStreamResponseHandler(), new ChannelPoolStreamHandler());
    ch.attr(RAPResponseDecoder.TIMEOUT_ATTR_KEY).set(new Timeout<None>(Executors.newSingleThreadScheduledExecutor(), 1000, TimeUnit.MILLISECONDS, None.none()));
    ch.attr(RAPStreamResponseHandler.CALLBACK_ATTR_KEY).set(response -> {
        StreamResponse streamResponse = response.getResponse();
        streamResponse.getEntityStream().setReader(new DrainReader());
    });
    return ch;
}

Example 70 with StreamResponse

use of com.linkedin.r2.message.stream.StreamResponse 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<String>();
    Map<Integer, PartitionData> partitionData = new HashMap<Integer, PartitionData>(1);
    partitionData.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1d));
    Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<URI, Map<Integer, PartitionData>>();
    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<String, String>(), new HashSet<URI>(), 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.<String, Object>emptyMap(), Collections.<String, Object>emptyMap(), Collections.<String, String>emptyMap(), schemes, Collections.<URI>emptySet()));
    Map<Integer, PartitionData> partitionWeight = new HashMap<Integer, PartitionData>();
    partitionWeight.put(0, new PartitionData(1d));
    partitionWeight.put(1, new PartitionData(2d));
    Map<URI, Map<Integer, PartitionData>> partitionDesc = new HashMap<URI, Map<Integer, PartitionData>>();
    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.<String, String>emptyMap(), new TransportCallbackAdapter<StreamResponse>(Callbacks.<StreamResponse>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.<String, String>emptyMap(), new TransportCallbackAdapter<StreamResponse>(Callbacks.<StreamResponse>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.<String, String>emptyMap(), new TransportCallbackAdapter<StreamResponse>(Callbacks.<StreamResponse>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.<String, String>emptyMap(), new TransportCallbackAdapter<StreamResponse>(Callbacks.<StreamResponse>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");
    }
}