Examples with Task org.opensearch.tasks.Task used on opensource projects

Example 1 with Task

use of org.opensearch.tasks.Task in project OpenSearch by opensearch-project.

the class TransportAction method execute.

/**
 * Use this method when the transport action call should result in creation of a new task associated with the call.
 *
 * This is a typical behavior.
 */
public final Task execute(Request request, ActionListener<Response> listener) {
    /*
         * While this version of execute could delegate to the TaskListener
         * version of execute that'd add yet another layer of wrapping on the
         * listener and prevent us from using the listener bare if there isn't a
         * task. That just seems like too many objects. Thus the two versions of
         * this method.
         */
    final Releasable unregisterChildNode = registerChildNode(request.getParentTask());
    final Task task;
    try {
        task = taskManager.register("transport", actionName, request);
    } catch (TaskCancelledException e) {
        unregisterChildNode.close();
        throw e;
    }
    execute(task, request, new ActionListener<Response>() {

        @Override
        public void onResponse(Response response) {
            try {
                Releasables.close(unregisterChildNode, () -> taskManager.unregister(task));
            } finally {
                listener.onResponse(response);
            }
        }

        @Override
        public void onFailure(Exception e) {
            try {
                Releasables.close(unregisterChildNode, () -> taskManager.unregister(task));
            } finally {
                listener.onFailure(e);
            }
        }
    });
    return task;
}

Example 2 with Task

use of org.opensearch.tasks.Task in project OpenSearch by opensearch-project.

the class TransportBulkAction method doInternalExecute.

protected void doInternalExecute(Task task, BulkRequest bulkRequest, String executorName, ActionListener<BulkResponse> listener) {
    final long startTime = relativeTime();
    final AtomicArray<BulkItemResponse> responses = new AtomicArray<>(bulkRequest.requests.size());
    boolean hasIndexRequestsWithPipelines = false;
    final Metadata metadata = clusterService.state().getMetadata();
    final Version minNodeVersion = clusterService.state().getNodes().getMinNodeVersion();
    for (DocWriteRequest<?> actionRequest : bulkRequest.requests) {
        IndexRequest indexRequest = getIndexWriteRequest(actionRequest);
        if (indexRequest != null) {
            // Each index request needs to be evaluated, because this method also modifies the IndexRequest
            boolean indexRequestHasPipeline = IngestService.resolvePipelines(actionRequest, indexRequest, metadata);
            hasIndexRequestsWithPipelines |= indexRequestHasPipeline;
        }
        if (actionRequest instanceof IndexRequest) {
            IndexRequest ir = (IndexRequest) actionRequest;
            ir.checkAutoIdWithOpTypeCreateSupportedByVersion(minNodeVersion);
            if (ir.getAutoGeneratedTimestamp() != IndexRequest.UNSET_AUTO_GENERATED_TIMESTAMP) {
                throw new IllegalArgumentException("autoGeneratedTimestamp should not be set externally");
            }
        }
    }
    if (hasIndexRequestsWithPipelines) {
        // this path is never taken.
        try {
            if (Assertions.ENABLED) {
                final boolean arePipelinesResolved = bulkRequest.requests().stream().map(TransportBulkAction::getIndexWriteRequest).filter(Objects::nonNull).allMatch(IndexRequest::isPipelineResolved);
                assert arePipelinesResolved : bulkRequest;
            }
            if (clusterService.localNode().isIngestNode()) {
                processBulkIndexIngestRequest(task, bulkRequest, executorName, listener);
            } else {
                ingestForwarder.forwardIngestRequest(BulkAction.INSTANCE, bulkRequest, listener);
            }
        } catch (Exception e) {
            listener.onFailure(e);
        }
        return;
    }
    final boolean includesSystem = includesSystem(bulkRequest, clusterService.state().metadata().getIndicesLookup(), systemIndices);
    if (includesSystem || needToCheck()) {
        // Attempt to create all the indices that we're going to need during the bulk before we start.
        // Step 1: collect all the indices in the request
        final Map<String, Boolean> indices = bulkRequest.requests.stream().filter(request -> request.opType() != DocWriteRequest.OpType.DELETE || request.versionType() == VersionType.EXTERNAL || request.versionType() == VersionType.EXTERNAL_GTE).collect(Collectors.toMap(DocWriteRequest::index, DocWriteRequest::isRequireAlias, (v1, v2) -> v1 || v2));
        /* Step 2: filter that to indices that don't exist and we can create. At the same time build a map of indices we can't create
             * that we'll use when we try to run the requests. */
        final Map<String, IndexNotFoundException> indicesThatCannotBeCreated = new HashMap<>();
        Set<String> autoCreateIndices = new HashSet<>();
        ClusterState state = clusterService.state();
        for (Map.Entry<String, Boolean> indexAndFlag : indices.entrySet()) {
            boolean shouldAutoCreate;
            final String index = indexAndFlag.getKey();
            try {
                shouldAutoCreate = shouldAutoCreate(index, state);
            } catch (IndexNotFoundException e) {
                shouldAutoCreate = false;
                indicesThatCannotBeCreated.put(index, e);
            }
            // We should only auto create if we are not requiring it to be an alias
            if (shouldAutoCreate && (indexAndFlag.getValue() == false)) {
                autoCreateIndices.add(index);
            }
        }
        // Step 3: create all the indices that are missing, if there are any missing. start the bulk after all the creates come back.
        if (autoCreateIndices.isEmpty()) {
            executeBulk(task, bulkRequest, startTime, listener, responses, indicesThatCannotBeCreated);
        } else {
            final AtomicInteger counter = new AtomicInteger(autoCreateIndices.size());
            for (String index : autoCreateIndices) {
                createIndex(index, bulkRequest.timeout(), minNodeVersion, new ActionListener<CreateIndexResponse>() {

                    @Override
                    public void onResponse(CreateIndexResponse result) {
                        if (counter.decrementAndGet() == 0) {
                            threadPool.executor(executorName).execute(new ActionRunnable<BulkResponse>(listener) {

                                @Override
                                protected void doRun() {
                                    executeBulk(task, bulkRequest, startTime, listener, responses, indicesThatCannotBeCreated);
                                }
                            });
                        }
                    }

                    @Override
                    public void onFailure(Exception e) {
                        if (!(ExceptionsHelper.unwrapCause(e) instanceof ResourceAlreadyExistsException)) {
                            // fail all requests involving this index, if create didn't work
                            for (int i = 0; i < bulkRequest.requests.size(); i++) {
                                DocWriteRequest<?> request = bulkRequest.requests.get(i);
                                if (request != null && setResponseFailureIfIndexMatches(responses, i, request, index, e)) {
                                    bulkRequest.requests.set(i, null);
                                }
                            }
                        }
                        if (counter.decrementAndGet() == 0) {
                            final ActionListener<BulkResponse> wrappedListener = ActionListener.wrap(listener::onResponse, inner -> {
                                inner.addSuppressed(e);
                                listener.onFailure(inner);
                            });
                            threadPool.executor(executorName).execute(new ActionRunnable<BulkResponse>(wrappedListener) {

                                @Override
                                protected void doRun() {
                                    executeBulk(task, bulkRequest, startTime, wrappedListener, responses, indicesThatCannotBeCreated);
                                }

                                @Override
                                public void onRejection(Exception rejectedException) {
                                    rejectedException.addSuppressed(e);
                                    super.onRejection(rejectedException);
                                }
                            });
                        }
                    }
                });
            }
        }
    } else {
        executeBulk(task, bulkRequest, startTime, listener, responses, emptyMap());
    }
}

Example 3 with Task

use of org.opensearch.tasks.Task in project OpenSearch by opensearch-project.

the class TasksIT method waitForCompletionTestCase.

/**
 * Test wait for completion.
 * @param storeResult should the task store its results
 * @param wait start waiting for a task. Accepts that id of the task to wait for and returns a future waiting for it.
 * @param validator validate the response and return the task ids that were found
 */
private <T> void waitForCompletionTestCase(boolean storeResult, Function<TaskId, ActionFuture<T>> wait, Consumer<T> validator) throws Exception {
    // Start blocking test task
    TestTaskPlugin.NodesRequest request = new TestTaskPlugin.NodesRequest("test");
    request.setShouldStoreResult(storeResult);
    ActionFuture<TestTaskPlugin.NodesResponse> future = client().execute(TestTaskPlugin.TestTaskAction.INSTANCE, request);
    ActionFuture<T> waitResponseFuture;
    TaskId taskId;
    try {
        taskId = waitForTestTaskStartOnAllNodes();
        // Wait for the task to start
        assertBusy(() -> client().admin().cluster().prepareGetTask(taskId).get());
        // Register listeners so we can be sure the waiting started
        CountDownLatch waitForWaitingToStart = new CountDownLatch(1);
        for (TransportService transportService : internalCluster().getInstances(TransportService.class)) {
            ((MockTaskManager) transportService.getTaskManager()).addListener(new MockTaskManagerListener() {

                @Override
                public void waitForTaskCompletion(Task task) {
                    waitForWaitingToStart.countDown();
                }

                @Override
                public void onTaskRegistered(Task task) {
                }

                @Override
                public void onTaskUnregistered(Task task) {
                }
            });
        }
        // Spin up a request to wait for the test task to finish
        waitResponseFuture = wait.apply(taskId);
        /* Wait for the wait to start. This should count down just *before* we wait for completion but after the list/get has got a
             * reference to the running task. Because we unblock immediately after this the task may no longer be running for us to wait
             * on which is fine. */
        waitForWaitingToStart.await();
    } finally {
        // Unblock the request so the wait for completion request can finish
        new TestTaskPlugin.UnblockTestTasksRequestBuilder(client(), TestTaskPlugin.UnblockTestTasksAction.INSTANCE).get();
    }
    // Now that the task is unblocked the list response will come back
    T waitResponse = waitResponseFuture.get();
    validator.accept(waitResponse);
    TestTaskPlugin.NodesResponse response = future.get();
    assertEquals(emptyList(), response.failures());
}

Example 4 with Task

use of org.opensearch.tasks.Task in project OpenSearch by opensearch-project.

the class TasksIT method testTransportBroadcastReplicationTasks.

public void testTransportBroadcastReplicationTasks() {
    // main task
    registerTaskManagerListeners(RefreshAction.NAME);
    // shard level tasks
    registerTaskManagerListeners(RefreshAction.NAME + "[s]");
    // primary and replica shard tasks
    registerTaskManagerListeners(RefreshAction.NAME + "[s][*]");
    createIndex("test");
    // Make sure all shards are allocated
    ensureGreen("test");
    client().admin().indices().prepareRefresh("test").get();
    // the refresh operation should produce one main task
    NumShards numberOfShards = getNumShards("test");
    logger.debug("number of shards, total: [{}], primaries: [{}] ", numberOfShards.totalNumShards, numberOfShards.numPrimaries);
    logger.debug("main events {}", numberOfEvents(RefreshAction.NAME, Tuple::v1));
    logger.debug("main event node {}", findEvents(RefreshAction.NAME, Tuple::v1).get(0).getTaskId().getNodeId());
    logger.debug("[s] events {}", numberOfEvents(RefreshAction.NAME + "[s]", Tuple::v1));
    logger.debug("[s][*] events {}", numberOfEvents(RefreshAction.NAME + "[s][*]", Tuple::v1));
    logger.debug("nodes with the index {}", internalCluster().nodesInclude("test"));
    assertEquals(1, numberOfEvents(RefreshAction.NAME, Tuple::v1));
    // Because it's broadcast replication action we will have as many [s] level requests
    // as we have primary shards on the coordinating node plus we will have one task per primary outside of the
    // coordinating node due to replication.
    // If all primaries are on the coordinating node, the number of tasks should be equal to the number of primaries
    // If all primaries are not on the coordinating node, the number of tasks should be equal to the number of primaries times 2
    assertThat(numberOfEvents(RefreshAction.NAME + "[s]", Tuple::v1), greaterThanOrEqualTo(numberOfShards.numPrimaries));
    assertThat(numberOfEvents(RefreshAction.NAME + "[s]", Tuple::v1), lessThanOrEqualTo(numberOfShards.numPrimaries * 2));
    // Verify that all [s] events have the proper parent
    // This is complicated because if the shard task runs on the same node it has main task as a parent
    // but if it runs on non-coordinating node it would have another intermediate [s] task on the coordinating node as a parent
    TaskInfo mainTask = findEvents(RefreshAction.NAME, Tuple::v1).get(0);
    List<TaskInfo> sTasks = findEvents(RefreshAction.NAME + "[s]", Tuple::v1);
    for (TaskInfo taskInfo : sTasks) {
        if (mainTask.getTaskId().getNodeId().equals(taskInfo.getTaskId().getNodeId())) {
            // This shard level task runs on the same node as a parent task - it should have the main task as a direct parent
            assertParentTask(Collections.singletonList(taskInfo), mainTask);
        } else {
            String description = taskInfo.getDescription();
            // This shard level task runs on another node - it should have a corresponding shard level task on the node where main task
            // is running
            List<TaskInfo> sTasksOnRequestingNode = findEvents(RefreshAction.NAME + "[s]", event -> event.v1() && mainTask.getTaskId().getNodeId().equals(event.v2().getTaskId().getNodeId()) && description.equals(event.v2().getDescription()));
            // There should be only one parent task
            assertEquals(1, sTasksOnRequestingNode.size());
            assertParentTask(Collections.singletonList(taskInfo), sTasksOnRequestingNode.get(0));
        }
    }
    // we will have as many [s][p] and [s][r] tasks as we have primary and replica shards
    assertEquals(numberOfShards.totalNumShards, numberOfEvents(RefreshAction.NAME + "[s][*]", Tuple::v1));
    // we the [s][p] and [s][r] tasks should have a corresponding [s] task on the same node as a parent
    List<TaskInfo> spEvents = findEvents(RefreshAction.NAME + "[s][*]", Tuple::v1);
    for (TaskInfo taskInfo : spEvents) {
        List<TaskInfo> sTask;
        if (taskInfo.getAction().endsWith("[s][p]")) {
            // A [s][p] level task should have a corresponding [s] level task on the same node
            sTask = findEvents(RefreshAction.NAME + "[s]", event -> event.v1() && taskInfo.getTaskId().getNodeId().equals(event.v2().getTaskId().getNodeId()) && taskInfo.getDescription().equals(event.v2().getDescription()));
        } else {
            // A [s][r] level task should have a corresponding [s] level task on the a different node (where primary is located)
            sTask = findEvents(RefreshAction.NAME + "[s]", event -> event.v1() && taskInfo.getParentTaskId().getNodeId().equals(event.v2().getTaskId().getNodeId()) && taskInfo.getDescription().equals(event.v2().getDescription()));
        }
        // There should be only one parent task
        assertEquals(1, sTask.size());
        assertParentTask(Collections.singletonList(taskInfo), sTask.get(0));
    }
}

Example 5 with Task

use of org.opensearch.tasks.Task in project OpenSearch by opensearch-project.

the class CancellableTasksIT method testBanOnlyNodesWithOutstandingDescendantTasks.

public void testBanOnlyNodesWithOutstandingDescendantTasks() throws Exception {
    if (randomBoolean()) {
        internalCluster().startNodes(randomIntBetween(1, 3));
    }
    Set<DiscoveryNode> nodes = StreamSupport.stream(clusterService().state().nodes().spliterator(), false).collect(Collectors.toSet());
    final TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 4));
    ActionFuture<TestResponse> rootTaskFuture = client().execute(TransportTestAction.ACTION, rootRequest);
    Set<TestRequest> pendingRequests = allowPartialRequest(rootRequest);
    TaskId rootTaskId = getRootTaskId(rootRequest);
    ActionFuture<CancelTasksResponse> cancelFuture = client().admin().cluster().prepareCancelTasks().setTaskId(rootTaskId).waitForCompletion(true).execute();
    if (randomBoolean()) {
        List<TaskInfo> runningTasks = client().admin().cluster().prepareListTasks().setActions(TransportTestAction.ACTION.name()).setDetailed(true).get().getTasks();
        for (TaskInfo subTask : randomSubsetOf(runningTasks)) {
            client().admin().cluster().prepareCancelTasks().setTaskId(subTask.getTaskId()).waitForCompletion(false).get();
        }
    }
    assertBusy(() -> {
        for (DiscoveryNode node : nodes) {
            TaskManager taskManager = internalCluster().getInstance(TransportService.class, node.getName()).getTaskManager();
            Set<TaskId> expectedBans = new HashSet<>();
            for (TestRequest req : pendingRequests) {
                if (req.node.equals(node)) {
                    List<Task> childTasks = taskManager.getTasks().values().stream().filter(t -> t.getParentTaskId() != null && t.getDescription().equals(req.taskDescription())).collect(Collectors.toList());
                    assertThat(childTasks, hasSize(1));
                    CancellableTask childTask = (CancellableTask) childTasks.get(0);
                    assertTrue(childTask.isCancelled());
                    expectedBans.add(childTask.getParentTaskId());
                }
            }
            assertThat(taskManager.getBannedTaskIds(), equalTo(expectedBans));
        }
    }, 30, TimeUnit.SECONDS);
    allowEntireRequest(rootRequest);
    cancelFuture.actionGet();
    waitForRootTask(rootTaskFuture);
    ensureAllBansRemoved();
}