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

Example 1 with CancellableTask

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

the class TimeoutTaskCancellationUtility method wrapWithCancellationListener.

/**
 * Wraps a listener with a timeout listener {@link TimeoutRunnableListener} to schedule the task cancellation for provided tasks on
 * generic thread pool
 * @param client - {@link NodeClient}
 * @param taskToCancel - task to schedule cancellation for
 * @param clusterSettings - {@link ClusterSettings}
 * @param listener - original listener associated with the task
 * @return wrapped listener
 */
public static <Response> ActionListener<Response> wrapWithCancellationListener(NodeClient client, CancellableTask taskToCancel, ClusterSettings clusterSettings, ActionListener<Response> listener) {
    final TimeValue globalTimeout = clusterSettings.get(SEARCH_CANCEL_AFTER_TIME_INTERVAL_SETTING);
    final TimeValue timeoutInterval = (taskToCancel.getCancellationTimeout() == null) ? globalTimeout : taskToCancel.getCancellationTimeout();
    // Note: -1 (or no timeout) will help to turn off cancellation. The combinations will be request level set at -1 or request level
    // set to null and cluster level set to -1.
    ActionListener<Response> listenerToReturn = listener;
    if (timeoutInterval.equals(SearchService.NO_TIMEOUT)) {
        return listenerToReturn;
    }
    try {
        final TimeoutRunnableListener<Response> wrappedListener = new TimeoutRunnableListener<>(timeoutInterval, listener, () -> {
            final CancelTasksRequest cancelTasksRequest = new CancelTasksRequest();
            cancelTasksRequest.setTaskId(new TaskId(client.getLocalNodeId(), taskToCancel.getId()));
            cancelTasksRequest.setReason("Cancellation timeout of " + timeoutInterval + " is expired");
            // force the origin to execute the cancellation as a system user
            new OriginSettingClient(client, TASKS_ORIGIN).admin().cluster().cancelTasks(cancelTasksRequest, ActionListener.wrap(r -> logger.debug("Scheduled cancel task with timeout: {} for original task: {} is successfully completed", timeoutInterval, cancelTasksRequest.getTaskId()), e -> logger.error(new ParameterizedMessage("Scheduled cancel task with timeout: {} for original task: {} is failed", timeoutInterval, cancelTasksRequest.getTaskId()), e)));
        });
        wrappedListener.cancellable = client.threadPool().schedule(wrappedListener, timeoutInterval, ThreadPool.Names.GENERIC);
        listenerToReturn = wrappedListener;
    } catch (Exception ex) {
        // if there is any exception in scheduling the cancellation task then continue without it
        logger.warn("Failed to schedule the cancellation task for original task: {}, will continue without it", taskToCancel.getId());
    }
    return listenerToReturn;
}

Example 2 with CancellableTask

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

the class CancellableTasksIT method testCancelOrphanedTasks.

public void testCancelOrphanedTasks() throws Exception {
    final String nodeWithRootTask = internalCluster().startDataOnlyNode();
    Set<DiscoveryNode> nodes = StreamSupport.stream(clusterService().state().nodes().spliterator(), false).collect(Collectors.toSet());
    TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 3));
    client(nodeWithRootTask).execute(TransportTestAction.ACTION, rootRequest);
    allowPartialRequest(rootRequest);
    try {
        internalCluster().stopRandomNode(InternalTestCluster.nameFilter(nodeWithRootTask));
        assertBusy(() -> {
            for (TransportService transportService : internalCluster().getInstances(TransportService.class)) {
                for (CancellableTask task : transportService.getTaskManager().getCancellableTasks().values()) {
                    if (task.getAction().equals(TransportTestAction.ACTION.name())) {
                        final TaskInfo taskInfo = task.taskInfo(transportService.getLocalNode().getId(), false);
                        assertTrue(taskInfo.toString(), task.isCancelled());
                        assertNotNull(taskInfo.toString(), task.getReasonCancelled());
                        assertThat(taskInfo.toString(), task.getReasonCancelled(), equalTo("channel was closed"));
                    }
                }
            }
        }, 30, TimeUnit.SECONDS);
    } finally {
        allowEntireRequest(rootRequest);
        ensureAllBansRemoved();
    }
}

Example 3 with CancellableTask

use of org.opensearch.tasks.CancellableTask 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();
}

Example 4 with CancellableTask

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

the class CancellableTasksTests method testCancelConcurrently.

public void testCancelConcurrently() throws Exception {
    setupTestNodes(Settings.EMPTY);
    final TaskManager taskManager = testNodes[0].transportService.getTaskManager();
    int numTasks = randomIntBetween(1, 10);
    List<CancellableTask> tasks = new ArrayList<>(numTasks);
    for (int i = 0; i < numTasks; i++) {
        tasks.add((CancellableTask) taskManager.register("type-" + i, "action-" + i, new CancellableNodeRequest()));
    }
    Thread[] threads = new Thread[randomIntBetween(1, 8)];
    AtomicIntegerArray notified = new AtomicIntegerArray(threads.length);
    Phaser phaser = new Phaser(threads.length + 1);
    final CancellableTask cancellingTask = randomFrom(tasks);
    for (int i = 0; i < threads.length; i++) {
        int idx = i;
        threads[i] = new Thread(() -> {
            phaser.arriveAndAwaitAdvance();
            taskManager.cancel(cancellingTask, "test", () -> assertTrue(notified.compareAndSet(idx, 0, 1)));
        });
        threads[i].start();
    }
    phaser.arriveAndAwaitAdvance();
    taskManager.unregister(cancellingTask);
    for (int i = 0; i < threads.length; i++) {
        threads[i].join();
        assertThat(notified.get(i), equalTo(1));
    }
    AtomicBoolean called = new AtomicBoolean();
    taskManager.cancel(cancellingTask, "test", () -> assertTrue(called.compareAndSet(false, true)));
    assertTrue(called.get());
}

Example 5 with CancellableTask

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

the class SearchRestCancellationIT method ensureSearchTaskIsCancelled.

private static void ensureSearchTaskIsCancelled(String transportAction, Function<String, String> nodeIdToName) throws Exception {
    SetOnce<TaskInfo> searchTask = new SetOnce<>();
    ListTasksResponse listTasksResponse = client().admin().cluster().prepareListTasks().get();
    for (TaskInfo task : listTasksResponse.getTasks()) {
        if (task.getAction().equals(transportAction)) {
            searchTask.set(task);
        }
    }
    assertNotNull(searchTask.get());
    TaskId taskId = searchTask.get().getTaskId();
    String nodeName = nodeIdToName.apply(taskId.getNodeId());
    assertBusy(() -> {
        TaskManager taskManager = internalCluster().getInstance(TransportService.class, nodeName).getTaskManager();
        Task task = taskManager.getTask(taskId.getId());
        assertThat(task, instanceOf(CancellableTask.class));
        assertTrue(((CancellableTask) task).isCancelled());
    });
}