Examples with TaskManager org.elasticsearch.tasks.TaskManager used on opensource projects

Example 1 with TaskManager

use of org.elasticsearch.tasks.TaskManager in project elasticsearch by elastic.

the class TransportActionFilterChainTests method testTooManyContinueProcessingRequest.

public void testTooManyContinueProcessingRequest() throws ExecutionException, InterruptedException {
    final int additionalContinueCount = randomInt(10);
    RequestTestFilter testFilter = new RequestTestFilter(randomInt(), new RequestCallback() {

        @Override
        public <Request extends ActionRequest, Response extends ActionResponse> void execute(Task task, String action, Request request, ActionListener<Response> listener, ActionFilterChain<Request, Response> actionFilterChain) {
            for (int i = 0; i <= additionalContinueCount; i++) {
                actionFilterChain.proceed(task, action, request, listener);
            }
        }
    });
    Set<ActionFilter> filters = new HashSet<>();
    filters.add(testFilter);
    String actionName = randomAsciiOfLength(randomInt(30));
    ActionFilters actionFilters = new ActionFilters(filters);
    TransportAction<TestRequest, TestResponse> transportAction = new TransportAction<TestRequest, TestResponse>(Settings.EMPTY, actionName, null, actionFilters, null, new TaskManager(Settings.EMPTY)) {

        @Override
        protected void doExecute(TestRequest request, ActionListener<TestResponse> listener) {
            listener.onResponse(new TestResponse());
        }
    };
    final CountDownLatch latch = new CountDownLatch(additionalContinueCount + 1);
    final AtomicInteger responses = new AtomicInteger();
    final List<Throwable> failures = new CopyOnWriteArrayList<>();
    transportAction.execute(new TestRequest(), new ActionListener<TestResponse>() {

        @Override
        public void onResponse(TestResponse testResponse) {
            responses.incrementAndGet();
            latch.countDown();
        }

        @Override
        public void onFailure(Exception e) {
            failures.add(e);
            latch.countDown();
        }
    });
    if (!latch.await(10, TimeUnit.SECONDS)) {
        fail("timeout waiting for the filter to notify the listener as many times as expected");
    }
    assertThat(testFilter.runs.get(), equalTo(1));
    assertThat(testFilter.lastActionName, equalTo(actionName));
    assertThat(responses.get(), equalTo(1));
    assertThat(failures.size(), equalTo(additionalContinueCount));
    for (Throwable failure : failures) {
        assertThat(failure, instanceOf(IllegalStateException.class));
    }
}

Example 2 with TaskManager

use of org.elasticsearch.tasks.TaskManager in project elasticsearch by elastic.

the class AsyncBulkByScrollActionTests method setupForTest.

@Before
public void setupForTest() {
    // Fill the context with something random so we can make sure we inherited it appropriately.
    expectedHeaders.clear();
    expectedHeaders.put(randomSimpleString(random()), randomSimpleString(random()));
    setupClient(new TestThreadPool(getTestName()));
    firstSearchRequest = new SearchRequest();
    testRequest = new DummyAbstractBulkByScrollRequest(firstSearchRequest);
    listener = new PlainActionFuture<>();
    scrollId = null;
    taskManager = new TaskManager(Settings.EMPTY);
    testTask = (WorkingBulkByScrollTask) taskManager.register("don'tcare", "hereeither", testRequest);
    localNode = new DiscoveryNode("thenode", buildNewFakeTransportAddress(), emptyMap(), emptySet(), Version.CURRENT);
    taskId = new TaskId(localNode.getId(), testTask.getId());
}

Example 3 with TaskManager

use of org.elasticsearch.tasks.TaskManager in project elasticsearch by elastic.

the class TransportMultiSearchActionTests method testBatchExecute.

public void testBatchExecute() throws Exception {
    // Initialize dependencies of TransportMultiSearchAction
    Settings settings = Settings.builder().put("node.name", TransportMultiSearchActionTests.class.getSimpleName()).build();
    ActionFilters actionFilters = mock(ActionFilters.class);
    when(actionFilters.filters()).thenReturn(new ActionFilter[0]);
    ThreadPool threadPool = new ThreadPool(settings);
    TaskManager taskManager = mock(TaskManager.class);
    TransportService transportService = new TransportService(Settings.EMPTY, null, null, TransportService.NOOP_TRANSPORT_INTERCEPTOR, boundAddress -> DiscoveryNode.createLocal(settings, boundAddress.publishAddress(), UUIDs.randomBase64UUID()), null) {

        @Override
        public TaskManager getTaskManager() {
            return taskManager;
        }
    };
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.state()).thenReturn(ClusterState.builder(new ClusterName("test")).build());
    IndexNameExpressionResolver resolver = new IndexNameExpressionResolver(Settings.EMPTY);
    // Keep track of the number of concurrent searches started by multi search api,
    // and if there are more searches than is allowed create an error and remember that.
    int maxAllowedConcurrentSearches = scaledRandomIntBetween(1, 16);
    AtomicInteger counter = new AtomicInteger();
    AtomicReference<AssertionError> errorHolder = new AtomicReference<>();
    // randomize whether or not requests are executed asynchronously
    final List<String> threadPoolNames = Arrays.asList(ThreadPool.Names.GENERIC, ThreadPool.Names.SAME);
    Randomness.shuffle(threadPoolNames);
    final ExecutorService commonExecutor = threadPool.executor(threadPoolNames.get(0));
    final ExecutorService rarelyExecutor = threadPool.executor(threadPoolNames.get(1));
    final Set<SearchRequest> requests = Collections.newSetFromMap(Collections.synchronizedMap(new IdentityHashMap<>()));
    TransportAction<SearchRequest, SearchResponse> searchAction = new TransportAction<SearchRequest, SearchResponse>(Settings.EMPTY, "action", threadPool, actionFilters, resolver, taskManager) {

        @Override
        protected void doExecute(SearchRequest request, ActionListener<SearchResponse> listener) {
            requests.add(request);
            int currentConcurrentSearches = counter.incrementAndGet();
            if (currentConcurrentSearches > maxAllowedConcurrentSearches) {
                errorHolder.set(new AssertionError("Current concurrent search [" + currentConcurrentSearches + "] is higher than is allowed [" + maxAllowedConcurrentSearches + "]"));
            }
            final ExecutorService executorService = rarely() ? rarelyExecutor : commonExecutor;
            executorService.execute(() -> {
                counter.decrementAndGet();
                listener.onResponse(new SearchResponse());
            });
        }
    };
    TransportMultiSearchAction action = new TransportMultiSearchAction(threadPool, actionFilters, transportService, clusterService, searchAction, resolver, 10);
    // Execute the multi search api and fail if we find an error after executing:
    try {
        /*
             * Allow for a large number of search requests in a single batch as previous implementations could stack overflow if the number
             * of requests in a single batch was large
             */
        int numSearchRequests = scaledRandomIntBetween(1, 8192);
        MultiSearchRequest multiSearchRequest = new MultiSearchRequest();
        multiSearchRequest.maxConcurrentSearchRequests(maxAllowedConcurrentSearches);
        for (int i = 0; i < numSearchRequests; i++) {
            multiSearchRequest.add(new SearchRequest());
        }
        MultiSearchResponse response = action.execute(multiSearchRequest).actionGet();
        assertThat(response.getResponses().length, equalTo(numSearchRequests));
        assertThat(requests.size(), equalTo(numSearchRequests));
        assertThat(errorHolder.get(), nullValue());
    } finally {
        assertTrue(ESTestCase.terminate(threadPool));
    }
}

Example 4 with TaskManager

use of org.elasticsearch.tasks.TaskManager in project elasticsearch by elastic.

the class TransportActionFilterChainTests method testActionFiltersRequest.

public void testActionFiltersRequest() throws ExecutionException, InterruptedException {
    int numFilters = randomInt(10);
    Set<Integer> orders = new HashSet<>(numFilters);
    while (orders.size() < numFilters) {
        orders.add(randomInt(10));
    }
    Set<ActionFilter> filters = new HashSet<>();
    for (Integer order : orders) {
        filters.add(new RequestTestFilter(order, randomFrom(RequestOperation.values())));
    }
    String actionName = randomAsciiOfLength(randomInt(30));
    ActionFilters actionFilters = new ActionFilters(filters);
    TransportAction<TestRequest, TestResponse> transportAction = new TransportAction<TestRequest, TestResponse>(Settings.EMPTY, actionName, null, actionFilters, null, new TaskManager(Settings.EMPTY)) {

        @Override
        protected void doExecute(TestRequest request, ActionListener<TestResponse> listener) {
            listener.onResponse(new TestResponse());
        }
    };
    ArrayList<ActionFilter> actionFiltersByOrder = new ArrayList<>(filters);
    Collections.sort(actionFiltersByOrder, new Comparator<ActionFilter>() {

        @Override
        public int compare(ActionFilter o1, ActionFilter o2) {
            return Integer.compare(o1.order(), o2.order());
        }
    });
    List<ActionFilter> expectedActionFilters = new ArrayList<>();
    boolean errorExpected = false;
    for (ActionFilter filter : actionFiltersByOrder) {
        RequestTestFilter testFilter = (RequestTestFilter) filter;
        expectedActionFilters.add(testFilter);
        if (testFilter.callback == RequestOperation.LISTENER_FAILURE) {
            errorExpected = true;
        }
        if (!(testFilter.callback == RequestOperation.CONTINUE_PROCESSING)) {
            break;
        }
    }
    PlainListenableActionFuture<TestResponse> future = new PlainListenableActionFuture<>(null);
    transportAction.execute(new TestRequest(), future);
    try {
        assertThat(future.get(), notNullValue());
        assertThat("shouldn't get here if an error is expected", errorExpected, equalTo(false));
    } catch (ExecutionException e) {
        assertThat("shouldn't get here if an error is not expected " + e.getMessage(), errorExpected, equalTo(true));
    }
    List<RequestTestFilter> testFiltersByLastExecution = new ArrayList<>();
    for (ActionFilter actionFilter : actionFilters.filters()) {
        testFiltersByLastExecution.add((RequestTestFilter) actionFilter);
    }
    Collections.sort(testFiltersByLastExecution, new Comparator<RequestTestFilter>() {

        @Override
        public int compare(RequestTestFilter o1, RequestTestFilter o2) {
            return Integer.compare(o1.executionToken, o2.executionToken);
        }
    });
    ArrayList<RequestTestFilter> finalTestFilters = new ArrayList<>();
    for (ActionFilter filter : testFiltersByLastExecution) {
        RequestTestFilter testFilter = (RequestTestFilter) filter;
        finalTestFilters.add(testFilter);
        if (!(testFilter.callback == RequestOperation.CONTINUE_PROCESSING)) {
            break;
        }
    }
    assertThat(finalTestFilters.size(), equalTo(expectedActionFilters.size()));
    for (int i = 0; i < finalTestFilters.size(); i++) {
        RequestTestFilter testFilter = finalTestFilters.get(i);
        assertThat(testFilter, equalTo(expectedActionFilters.get(i)));
        assertThat(testFilter.runs.get(), equalTo(1));
        assertThat(testFilter.lastActionName, equalTo(actionName));
    }
}

Example 5 with TaskManager

use of org.elasticsearch.tasks.TaskManager in project elasticsearch by elastic.

the class InternalTestCluster method assertShardIndexCounter.

private void assertShardIndexCounter() throws Exception {
    assertBusy(() -> {
        final Collection<NodeAndClient> nodesAndClients = nodes.values();
        for (NodeAndClient nodeAndClient : nodesAndClients) {
            IndicesService indexServices = getInstance(IndicesService.class, nodeAndClient.name);
            for (IndexService indexService : indexServices) {
                for (IndexShard indexShard : indexService) {
                    int activeOperationsCount = indexShard.getActiveOperationsCount();
                    if (activeOperationsCount > 0) {
                        TaskManager taskManager = getInstance(TransportService.class, nodeAndClient.name).getTaskManager();
                        DiscoveryNode localNode = getInstance(ClusterService.class, nodeAndClient.name).localNode();
                        List<TaskInfo> taskInfos = taskManager.getTasks().values().stream().filter(task -> task instanceof ReplicationTask).map(task -> task.taskInfo(localNode.getId(), true)).collect(Collectors.toList());
                        ListTasksResponse response = new ListTasksResponse(taskInfos, Collections.emptyList(), Collections.emptyList());
                        try {
                            XContentBuilder builder = XContentFactory.jsonBuilder().prettyPrint().value(response);
                            throw new AssertionError("expected index shard counter on shard " + indexShard.shardId() + " on node " + nodeAndClient.name + " to be 0 but was " + activeOperationsCount + ". Current replication tasks on node:\n" + builder.string());
                        } catch (IOException e) {
                            throw new RuntimeException("caught exception while building response [" + response + "]", e);
                        }
                    }
                }
            }
        }
    });
}