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

Example 1 with TaskCancelledException

use of org.opensearch.tasks.TaskCancelledException 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 TaskCancelledException

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

the class CancelTests method testCancel.

/**
 * Executes the cancellation test
 */
private void testCancel(String action, AbstractBulkByScrollRequestBuilder<?, ?> builder, CancelAssertion assertion, Matcher<String> taskDescriptionMatcher) throws Exception {
    createIndex(INDEX);
    // Total number of documents created for this test (~10 per primary shard per slice)
    int numDocs = getNumShards(INDEX).numPrimaries * 10 * builder.request().getSlices();
    ALLOWED_OPERATIONS.release(numDocs);
    logger.debug("setting up [{}] docs", numDocs);
    indexRandom(true, false, true, IntStream.range(0, numDocs).mapToObj(i -> client().prepareIndex().setIndex(INDEX).setId(String.valueOf(i)).setSource("n", i)).collect(Collectors.toList()));
    // Checks that the all documents have been indexed and correctly counted
    assertHitCount(client().prepareSearch(INDEX).setSize(0).get(), numDocs);
    assertThat(ALLOWED_OPERATIONS.drainPermits(), equalTo(0));
    // Scroll by 1 so that cancellation is easier to control
    builder.source().setSize(1);
    /* Allow a random number of the documents less the number of workers
         * to be modified by the reindex action. That way at least one worker
         * is blocked. */
    int numModifiedDocs = randomIntBetween(builder.request().getSlices() * 2, numDocs);
    logger.debug("chose to modify [{}] out of [{}] docs", numModifiedDocs, numDocs);
    ALLOWED_OPERATIONS.release(numModifiedDocs - builder.request().getSlices());
    // Now execute the reindex action...
    ActionFuture<? extends BulkByScrollResponse> future = builder.execute();
    /* ... and wait for the indexing operation listeners to block. It
         * is important to realize that some of the workers might have
         * exhausted their slice while others might have quite a bit left
         * to work on. We can't control that. */
    logger.debug("waiting for updates to be blocked");
    assertBusy(() -> assertTrue("updates blocked", ALLOWED_OPERATIONS.hasQueuedThreads() && ALLOWED_OPERATIONS.availablePermits() == 0), 1, TimeUnit.MINUTES);
    // 10 seconds is usually fine but on heavily loaded machines this can take a while
    // Status should show the task running
    TaskInfo mainTask = findTaskToCancel(action, builder.request().getSlices());
    BulkByScrollTask.Status status = (BulkByScrollTask.Status) mainTask.getStatus();
    assertNull(status.getReasonCancelled());
    // Description shouldn't be empty
    assertThat(mainTask.getDescription(), taskDescriptionMatcher);
    // Cancel the request while the action is blocked by the indexing operation listeners.
    // This will prevent further requests from being sent.
    ListTasksResponse cancelTasksResponse = client().admin().cluster().prepareCancelTasks().setTaskId(mainTask.getTaskId()).get();
    cancelTasksResponse.rethrowFailures("Cancel");
    assertThat(cancelTasksResponse.getTasks(), hasSize(1));
    /* The status should now show canceled. The request will still be in the
         * list because it is (or its children are) still blocked. */
    mainTask = client().admin().cluster().prepareGetTask(mainTask.getTaskId()).get().getTask().getTask();
    status = (BulkByScrollTask.Status) mainTask.getStatus();
    logger.debug("asserting that parent is marked canceled {}", status);
    assertEquals(CancelTasksRequest.DEFAULT_REASON, status.getReasonCancelled());
    if (builder.request().getSlices() > 1) {
        boolean foundCancelled = false;
        ListTasksResponse sliceList = client().admin().cluster().prepareListTasks().setParentTaskId(mainTask.getTaskId()).setDetailed(true).get();
        sliceList.rethrowFailures("Fetch slice tasks");
        logger.debug("finding at least one canceled child among {}", sliceList.getTasks());
        for (TaskInfo slice : sliceList.getTasks()) {
            BulkByScrollTask.Status sliceStatus = (BulkByScrollTask.Status) slice.getStatus();
            if (sliceStatus.getReasonCancelled() == null)
                continue;
            assertEquals(CancelTasksRequest.DEFAULT_REASON, sliceStatus.getReasonCancelled());
            foundCancelled = true;
        }
        assertTrue("Didn't find at least one sub task that was cancelled", foundCancelled);
    }
    logger.debug("unblocking the blocked update");
    ALLOWED_OPERATIONS.release(builder.request().getSlices());
    // Checks that no more operations are executed
    assertBusy(() -> {
        if (builder.request().getSlices() == 1) {
            /* We can only be sure that we've drained all the permits if we only use a single worker. Otherwise some worker may have
                 * exhausted all of its documents before we blocked. */
            assertEquals(0, ALLOWED_OPERATIONS.availablePermits());
        }
        assertEquals(0, ALLOWED_OPERATIONS.getQueueLength());
    });
    // And check the status of the response
    BulkByScrollResponse response;
    try {
        response = future.get(30, TimeUnit.SECONDS);
    } catch (Exception e) {
        if (ExceptionsHelper.unwrapCausesAndSuppressed(e, t -> t instanceof TaskCancelledException).isPresent()) {
            // the scroll request was cancelled
            return;
        }
        String tasks = client().admin().cluster().prepareListTasks().setParentTaskId(mainTask.getTaskId()).setDetailed(true).get().toString();
        throw new RuntimeException("Exception while waiting for the response. Running tasks: " + tasks, e);
    }
    assertThat(response.getReasonCancelled(), equalTo("by user request"));
    assertThat(response.getBulkFailures(), emptyIterable());
    assertThat(response.getSearchFailures(), emptyIterable());
    if (builder.request().getSlices() >= 1) {
        // If we have more than one worker we might not have made all the modifications
        numModifiedDocs -= ALLOWED_OPERATIONS.availablePermits();
    }
    flushAndRefresh(INDEX);
    assertion.assertThat(response, numDocs, numModifiedDocs);
}

Example 3 with TaskCancelledException

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

the class SearchCancellationIT method verifyCancellationException.

private void verifyCancellationException(ShardSearchFailure[] failures) {
    for (ShardSearchFailure searchFailure : failures) {
        // failure may happen while executing the search or while sending shard request for next phase.
        // Below assertion is handling both the cases
        final Throwable topFailureCause = searchFailure.getCause();
        assertTrue(searchFailure.toString(), topFailureCause instanceof TransportException || topFailureCause instanceof TaskCancelledException);
        if (topFailureCause instanceof TransportException) {
            assertTrue(topFailureCause.getCause() instanceof TaskCancelledException);
        }
    }
}

Example 4 with TaskCancelledException

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

the class FetchPhase method execute.

public void execute(SearchContext context) {
    if (LOGGER.isTraceEnabled()) {
        LOGGER.trace("{}", new SearchContextSourcePrinter(context));
    }
    if (context.isCancelled()) {
        throw new TaskCancelledException("cancelled task with reason: " + context.getTask().getReasonCancelled());
    }
    if (context.docIdsToLoadSize() == 0) {
        // no individual hits to process, so we shortcut
        context.fetchResult().hits(new SearchHits(new SearchHit[0], context.queryResult().getTotalHits(), context.queryResult().getMaxScore()));
        return;
    }
    DocIdToIndex[] docs = new DocIdToIndex[context.docIdsToLoadSize()];
    for (int index = 0; index < context.docIdsToLoadSize(); index++) {
        docs[index] = new DocIdToIndex(context.docIdsToLoad()[context.docIdsToLoadFrom() + index], index);
    }
    // make sure that we iterate in doc id order
    Arrays.sort(docs);
    Map<String, Set<String>> storedToRequestedFields = new HashMap<>();
    FieldsVisitor fieldsVisitor = createStoredFieldsVisitor(context, storedToRequestedFields);
    FetchContext fetchContext = new FetchContext(context);
    SearchHit[] hits = new SearchHit[context.docIdsToLoadSize()];
    List<FetchSubPhaseProcessor> processors = getProcessors(context.shardTarget(), fetchContext);
    int currentReaderIndex = -1;
    LeafReaderContext currentReaderContext = null;
    CheckedBiConsumer<Integer, FieldsVisitor, IOException> fieldReader = null;
    boolean hasSequentialDocs = hasSequentialDocs(docs);
    for (int index = 0; index < context.docIdsToLoadSize(); index++) {
        if (context.isCancelled()) {
            throw new TaskCancelledException("cancelled task with reason: " + context.getTask().getReasonCancelled());
        }
        int docId = docs[index].docId;
        try {
            int readerIndex = ReaderUtil.subIndex(docId, context.searcher().getIndexReader().leaves());
            if (currentReaderIndex != readerIndex) {
                currentReaderContext = context.searcher().getIndexReader().leaves().get(readerIndex);
                currentReaderIndex = readerIndex;
                if (currentReaderContext.reader() instanceof SequentialStoredFieldsLeafReader && hasSequentialDocs && docs.length >= 10) {
                    // All the docs to fetch are adjacent but Lucene stored fields are optimized
                    // for random access and don't optimize for sequential access - except for merging.
                    // So we do a little hack here and pretend we're going to do merges in order to
                    // get better sequential access.
                    SequentialStoredFieldsLeafReader lf = (SequentialStoredFieldsLeafReader) currentReaderContext.reader();
                    fieldReader = lf.getSequentialStoredFieldsReader()::visitDocument;
                } else {
                    fieldReader = currentReaderContext.reader()::document;
                }
                for (FetchSubPhaseProcessor processor : processors) {
                    processor.setNextReader(currentReaderContext);
                }
            }
            assert currentReaderContext != null;
            HitContext hit = prepareHitContext(context, fetchContext.searchLookup(), fieldsVisitor, docId, storedToRequestedFields, currentReaderContext, fieldReader);
            for (FetchSubPhaseProcessor processor : processors) {
                processor.process(hit);
            }
            hits[docs[index].index] = hit.hit();
        } catch (Exception e) {
            throw new FetchPhaseExecutionException(context.shardTarget(), "Error running fetch phase for doc [" + docId + "]", e);
        }
    }
    if (context.isCancelled()) {
        throw new TaskCancelledException("cancelled task with reason: " + context.getTask().getReasonCancelled());
    }
    TotalHits totalHits = context.queryResult().getTotalHits();
    context.fetchResult().hits(new SearchHits(hits, totalHits, context.queryResult().getMaxScore()));
}

Example 5 with TaskCancelledException

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

the class CancellableTasksTests method testRegisterAndExecuteChildTaskWhileParentTaskIsBeingCanceled.

public void testRegisterAndExecuteChildTaskWhileParentTaskIsBeingCanceled() throws Exception {
    setupTestNodes(Settings.EMPTY);
    connectNodes(testNodes);
    final TaskManager taskManager = testNodes[0].transportService.getTaskManager();
    CancellableNodesRequest parentRequest = new CancellableNodesRequest("parent");
    final Task parentTask = taskManager.register("test", "test", parentRequest);
    final TaskId parentTaskId = parentTask.taskInfo(testNodes[0].getNodeId(), false).getTaskId();
    taskManager.setBan(new TaskId(testNodes[0].getNodeId(), parentTask.getId()), "test");
    CancellableNodesRequest childRequest = new CancellableNodesRequest("child");
    childRequest.setParentTask(parentTaskId);
    CancellableTestNodesAction testAction = new CancellableTestNodesAction("internal:testAction", threadPool, testNodes[1].clusterService, testNodes[0].transportService, false, new CountDownLatch(1));
    TaskCancelledException cancelledException = expectThrows(TaskCancelledException.class, () -> testAction.execute(childRequest, ActionListener.wrap(() -> fail("must not execute"))));
    assertThat(cancelledException.getMessage(), startsWith("Task cancelled before it started:"));
    CountDownLatch latch = new CountDownLatch(1);
    taskManager.startBanOnChildrenNodes(parentTaskId.getId(), latch::countDown);
    assertTrue("onChildTasksCompleted() is not invoked", latch.await(1, TimeUnit.SECONDS));
}