Examples with Task com.linkedin.r2.transport.http.client.RateLimiter.Task used on opensource projects

Example 6 with Task

use of com.linkedin.r2.transport.http.client.RateLimiter.Task in project rest.li by linkedin.

the class TestMIMEInputStream method abortWhenNoOutstandingReadTask.

//Abort in the middle of a write task.
@Test
public void abortWhenNoOutstandingReadTask() throws Exception {
    final StringBuilder builder = new StringBuilder();
    for (int i = 0; i < (TEST_CHUNK_SIZE * 10) + 2; i++) {
        builder.append('a');
    }
    //The slow byte array input stream will verify that we call an abort before the first read task is finished.
    final byte[] largeInputData = builder.toString().getBytes();
    final SlowByteArrayInputStream inputStream = new SlowByteArrayInputStream(largeInputData, 300, 10);
    final SlowByteArrayInputStream spyInputStream = spy(inputStream);
    //Setup:
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(spyInputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).build();
    //By the time the first onWritePossible() completes, half the data should be transferred
    //Then the abort task will run.
    when(writeHandle.remaining()).thenReturn(5, new Integer[] { 4, 3, 2, 1, 0 });
    final ByteArrayOutputStream byteArrayOutputStream = setupMockWriteHandleToOutputStream(writeHandle);
    //Setup for the close on the input stream.
    //The close must happen for the test to finish.
    final CountDownLatch closeInputStreamLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            closeInputStreamLatch.countDown();
            return null;
        }
    }).when(spyInputStream).close();
    ///////////////////////////////////
    //Start things off
    //Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onWritePossible();
    multiPartMIMEInputStream.onAbort(new IOException());
    //Wait to finish
    try {
        boolean successful = closeInputStreamLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        if (!successful) {
            Assert.fail("Timeout when waiting for abort to happen!");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to be closed!");
    }
    ///////////////////////////////////
    //Assert
    final byte[] expectedBytes = Arrays.copyOf(largeInputData, TEST_CHUNK_SIZE * 5);
    Assert.assertEquals(byteArrayOutputStream.toByteArray(), expectedBytes, "Partial data from the input stream should have successfully been transferred");
    //Mock verifies:
    verify(spyInputStream, times(1)).close();
    verify(spyInputStream, times(5)).read(isA(byte[].class));
    verify(writeHandle, times(5)).write(isA(ByteString.class));
    verify(writeHandle, times(6)).remaining();
    verify(writeHandle, never()).error(isA(Throwable.class));
    verify(writeHandle, never()).done();
    verifyNoMoreInteractions(writeHandle);
}

Example 7 with Task

use of com.linkedin.r2.transport.http.client.RateLimiter.Task in project rest.li by linkedin.

the class TestRateLimiter method testMaxRunningTasks.

@Test
public void testMaxRunningTasks() throws Exception {
    final int total = 20;
    final int maxRunning = 5;
    final int period = 100;
    final Random rand = new Random();
    final CountDownLatch latch = new CountDownLatch(total);
    final AtomicInteger totalStarted = new AtomicInteger();
    final AtomicInteger totalFinished = new AtomicInteger();
    final Task r = new Task() {

        @Override
        public void run(final SimpleCallback callback) {
            totalStarted.incrementAndGet();
            int delay = period + rand.nextInt(period);
            _executor.schedule(new Runnable() {

                @Override
                public void run() {
                    totalFinished.incrementAndGet();
                    callback.onDone();
                }
            }, delay, TimeUnit.MILLISECONDS);
            latch.countDown();
        }
    };
    RateLimiter limiter = new ExponentialBackOffRateLimiter(period, period, period, _executor, maxRunning);
    limiter.setPeriod(period);
    for (int i = 0; i < total; ++i) {
        limiter.submit(r);
    }
    // check the current number of concurrent tasks every 100ms.
    for (int i = 0; i < total * 2; ++i) {
        int currentRunning = totalStarted.get() - totalFinished.get();
        Assert.assertTrue(currentRunning <= maxRunning, "Should have less than " + maxRunning + " concurrent tasks");
        Thread.sleep(period);
    }
    Assert.assertTrue(latch.await(30, TimeUnit.SECONDS));
    Assert.assertEquals(total, totalStarted.get());
}

Example 8 with Task

use of com.linkedin.r2.transport.http.client.RateLimiter.Task in project rest.li by linkedin.

the class TestAsyncSharedPoolImpl method testGetItemCancelled.

@Test
public void testGetItemCancelled() throws Exception {
    final LifecycleMock lifecycle = new LifecycleMock();
    final CountDownLatch createLatch = new CountDownLatch(1);
    lifecycle.setCreateConsumer(callback -> {
        try {
            createLatch.await();
            callback.onSuccess(ITEM);
        } catch (Exception e) {
            callback.onError(e);
        }
    });
    AsyncSharedPoolImpl<Object> pool = new AsyncSharedPoolImpl<>(POOL_NAME, lifecycle, SCHEDULER, LIMITER, NO_POOL_TIMEOUT, MAX_WAITERS);
    pool.start();
    // Only one thread will perform the actual item creation task and the rest
    // will return immediately. Therefore we wait for GET_COUNT - 1 threads to complete.
    final CountDownLatch getLatch = new CountDownLatch(GET_COUNT - 1);
    final ConcurrentLinkedQueue<Cancellable> cancellables = new ConcurrentLinkedQueue<>();
    for (int i = 0; i < GET_COUNT; i++) {
        SCHEDULER.execute(() -> {
            cancellables.add(pool.get(new FutureCallback<>()));
            getLatch.countDown();
        });
    }
    if (!getLatch.await(GET_TIMEOUT, TIME_UNIT)) {
        Assert.fail("Timed out awaiting for get");
    }
    Assert.assertEquals(cancellables.size(), GET_COUNT - 1);
    // Cancelling waiters should all succeed
    cancellables.stream().forEach(cancellable -> Assert.assertTrue(cancellable.cancel()));
    // Cancel the last waiter blocking item creation
    Assert.assertEquals(pool.cancelWaiters().size(), 1);
    createLatch.countDown();
    FutureCallback<None> shutdownCallback = new FutureCallback<>();
    pool.shutdown(shutdownCallback);
    shutdownCallback.get(SHUTDOWN_TIMEOUT, TIME_UNIT);
}

Example 9 with Task

use of com.linkedin.r2.transport.http.client.RateLimiter.Task in project rest.li by linkedin.

the class MultiplexedRequestHandlerImpl method handleRequest.

@Override
public void handleRequest(RestRequest request, RequestContext requestContext, final Callback<RestResponse> callback) {
    if (HttpMethod.POST != HttpMethod.valueOf(request.getMethod())) {
        _log.error("POST is expected, but " + request.getMethod() + " received");
        callback.onError(RestException.forError(HttpStatus.S_405_METHOD_NOT_ALLOWED.getCode(), "Invalid method"));
        return;
    }
    IndividualRequestMap individualRequests;
    try {
        individualRequests = extractIndividualRequests(request);
    } catch (RestException e) {
        _log.error("Invalid multiplexed request", e);
        callback.onError(e);
        return;
    } catch (Exception e) {
        _log.error("Invalid multiplexed request", e);
        callback.onError(RestException.forError(HttpStatus.S_400_BAD_REQUEST.getCode(), e));
        return;
    }
    // prepare the map of individual responses to be collected
    final IndividualResponseMap individualResponses = new IndividualResponseMap(individualRequests.size());
    final Map<String, HttpCookie> responseCookies = new HashMap<>();
    // all tasks are Void and side effect based, that will be useful when we add streaming
    Task<?> requestProcessingTask = createParallelRequestsTask(request, requestContext, individualRequests, individualResponses, responseCookies);
    Task<Void> responseAggregationTask = Task.action("send aggregated response", () -> {
        RestResponse aggregatedResponse = aggregateResponses(individualResponses, responseCookies);
        callback.onSuccess(aggregatedResponse);
    });
    _engine.run(requestProcessingTask.andThen(responseAggregationTask), MUX_PLAN_CLASS);
}

Example 10 with Task

use of com.linkedin.r2.transport.http.client.RateLimiter.Task in project rest.li by linkedin.

the class TestMultiplexerRunMode method testMultiplexedAsyncGetRequest.

@Test(dataProvider = "multiplexerConfigurations")
public void testMultiplexedAsyncGetRequest(MultiplexerRunMode multiplexerRunMode) throws URISyntaxException, IOException, InterruptedException {
    RestLiConfig config = new RestLiConfig();
    config.addResourcePackageNames("com.linkedin.restli.server.multiplexer.resources");
    config.setMultiplexerRunMode(multiplexerRunMode);
    CountingEngine engine = engine();
    RestLiServer server = new RestLiServer(config, resourceFactory(), engine);
    IndividualRequest r0 = individualRequest("/users/0", null, Collections.<String, IndividualRequest>emptyMap());
    IndividualRequest r1 = individualRequest("/users/1", null, Collections.<String, IndividualRequest>emptyMap());
    IndividualRequest r2 = individualRequest("/users/2", null, ImmutableMap.of("0", r0, "1", r1));
    // request is seq(par(r0, r1), r2)
    RestRequest request = muxRestRequest(ImmutableMap.of("2", r2));
    CountDownLatch latch = new CountDownLatch(1);
    server.handleRequest(request, new RequestContext(), callback(latch));
    assertTrue(latch.await(5, TimeUnit.SECONDS));
    if (multiplexerRunMode == MultiplexerRunMode.SINGLE_PLAN) {
        assertEquals(engine.plansStarted(), 1);
    } else {
        // in MULTIPLE_PLANS mode: 1 task for multiplexed request itself + 3 individual tasks r0, r1, r2
        assertEquals(engine.plansStarted(), 1 + 3);
    }
}