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

Example 1 with TaskManager

use of org.opensearch.tasks.TaskManager 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 2 with TaskManager

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

the class PrimaryReplicaSyncerTests method testSyncerSendsOffCorrectDocuments.

public void testSyncerSendsOffCorrectDocuments() throws Exception {
    IndexShard shard = newStartedShard(true);
    TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet());
    AtomicBoolean syncActionCalled = new AtomicBoolean();
    List<ResyncReplicationRequest> resyncRequests = new ArrayList<>();
    PrimaryReplicaSyncer.SyncAction syncAction = (request, parentTask, allocationId, primaryTerm, listener) -> {
        logger.info("Sending off {} operations", request.getOperations().length);
        syncActionCalled.set(true);
        resyncRequests.add(request);
        assertThat(parentTask, instanceOf(PrimaryReplicaSyncer.ResyncTask.class));
        listener.onResponse(new ResyncReplicationResponse());
    };
    PrimaryReplicaSyncer syncer = new PrimaryReplicaSyncer(taskManager, syncAction);
    syncer.setChunkSize(new ByteSizeValue(randomIntBetween(1, 10)));
    int numDocs = randomInt(10);
    for (int i = 0; i < numDocs; i++) {
        // Index doc but not advance local checkpoint.
        shard.applyIndexOperationOnPrimary(Versions.MATCH_ANY, VersionType.INTERNAL, new SourceToParse(shard.shardId().getIndexName(), "_doc", Integer.toString(i), new BytesArray("{}"), XContentType.JSON), SequenceNumbers.UNASSIGNED_SEQ_NO, 0, IndexRequest.UNSET_AUTO_GENERATED_TIMESTAMP, true);
    }
    long globalCheckPoint = numDocs > 0 ? randomIntBetween(0, numDocs - 1) : 0;
    boolean syncNeeded = numDocs > 0;
    String allocationId = shard.routingEntry().allocationId().getId();
    shard.updateShardState(shard.routingEntry(), shard.getPendingPrimaryTerm(), null, 1000L, Collections.singleton(allocationId), new IndexShardRoutingTable.Builder(shard.shardId()).addShard(shard.routingEntry()).build());
    shard.updateLocalCheckpointForShard(allocationId, globalCheckPoint);
    assertEquals(globalCheckPoint, shard.getLastKnownGlobalCheckpoint());
    logger.info("Total ops: {}, global checkpoint: {}", numDocs, globalCheckPoint);
    PlainActionFuture<PrimaryReplicaSyncer.ResyncTask> fut = new PlainActionFuture<>();
    syncer.resync(shard, fut);
    PrimaryReplicaSyncer.ResyncTask resyncTask = fut.get();
    if (syncNeeded) {
        assertTrue("Sync action was not called", syncActionCalled.get());
        ResyncReplicationRequest resyncRequest = resyncRequests.remove(0);
        assertThat(resyncRequest.getTrimAboveSeqNo(), equalTo(numDocs - 1L));
        assertThat("trimAboveSeqNo has to be specified in request #0 only", resyncRequests.stream().mapToLong(ResyncReplicationRequest::getTrimAboveSeqNo).filter(seqNo -> seqNo != SequenceNumbers.UNASSIGNED_SEQ_NO).findFirst().isPresent(), is(false));
        assertThat(resyncRequest.getMaxSeenAutoIdTimestampOnPrimary(), equalTo(shard.getMaxSeenAutoIdTimestamp()));
    }
    if (syncNeeded && globalCheckPoint < numDocs - 1) {
        assertThat(resyncTask.getSkippedOperations(), equalTo(0));
        assertThat(resyncTask.getResyncedOperations(), equalTo(Math.toIntExact(numDocs - 1 - globalCheckPoint)));
        if (shard.indexSettings.isSoftDeleteEnabled()) {
            assertThat(resyncTask.getTotalOperations(), equalTo(Math.toIntExact(numDocs - 1 - globalCheckPoint)));
        } else {
            assertThat(resyncTask.getTotalOperations(), equalTo(numDocs));
        }
    } else {
        assertThat(resyncTask.getSkippedOperations(), equalTo(0));
        assertThat(resyncTask.getResyncedOperations(), equalTo(0));
        assertThat(resyncTask.getTotalOperations(), equalTo(0));
    }
    closeShards(shard);
}

Example 3 with TaskManager

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

the class TransportActionFilterChainTests method testActionFiltersRequest.

public void testActionFiltersRequest() throws 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 = randomAlphaOfLength(randomInt(30));
    ActionFilters actionFilters = new ActionFilters(filters);
    TransportAction<TestRequest, TestResponse> transportAction = new TransportAction<TestRequest, TestResponse>(actionName, actionFilters, new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet())) {

        @Override
        protected void doExecute(Task task, TestRequest request, ActionListener<TestResponse> listener) {
            listener.onResponse(new TestResponse());
        }
    };
    ArrayList<ActionFilter> actionFiltersByOrder = new ArrayList<>(filters);
    actionFiltersByOrder.sort(Comparator.comparingInt(ActionFilter::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;
        }
    }
    PlainActionFuture<TestResponse> future = PlainActionFuture.newFuture();
    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);
    }
    testFiltersByLastExecution.sort(Comparator.comparingInt(o -> o.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 4 with TaskManager

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

Example 5 with TaskManager

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