Examples with PersistentTask org.opensearch.persistent.PersistentTasksCustomMetadata.PersistentTask used on opensource projects

Example 6 with PersistentTask

use of org.opensearch.persistent.PersistentTasksCustomMetadata.PersistentTask in project OpenSearch by opensearch-project.

the class PersistentTasksCustomMetadataTests method testSerializationContext.

@SuppressWarnings("unchecked")
public void testSerializationContext() throws Exception {
    PersistentTasksCustomMetadata testInstance = createTestInstance();
    for (int i = 0; i < randomInt(10); i++) {
        testInstance = (PersistentTasksCustomMetadata) makeTestChanges(testInstance);
    }
    ToXContent.MapParams params = new ToXContent.MapParams(Collections.singletonMap(Metadata.CONTEXT_MODE_PARAM, randomFrom(CONTEXT_MODE_SNAPSHOT, CONTEXT_MODE_GATEWAY)));
    XContentType xContentType = randomFrom(XContentType.values());
    BytesReference shuffled = toShuffledXContent(testInstance, xContentType, params, false);
    PersistentTasksCustomMetadata newInstance;
    try (XContentParser parser = createParser(XContentFactory.xContent(xContentType), shuffled)) {
        newInstance = doParseInstance(parser);
    }
    assertNotSame(newInstance, testInstance);
    assertEquals(testInstance.tasks().size(), newInstance.tasks().size());
    for (PersistentTask<?> testTask : testInstance.tasks()) {
        PersistentTask<TestParams> newTask = (PersistentTask<TestParams>) newInstance.getTask(testTask.getId());
        assertNotNull(newTask);
        // Things that should be serialized
        assertEquals(testTask.getTaskName(), newTask.getTaskName());
        assertEquals(testTask.getId(), newTask.getId());
        assertEquals(testTask.getState(), newTask.getState());
        assertEquals(testTask.getParams(), newTask.getParams());
        // Things that shouldn't be serialized
        assertEquals(0, newTask.getAllocationId());
        assertNull(newTask.getExecutorNode());
    }
}

Example 7 with PersistentTask

use of org.opensearch.persistent.PersistentTasksCustomMetadata.PersistentTask in project OpenSearch by opensearch-project.

the class PersistentTasksNodeServiceTests method testRegisterTaskFails.

public void testRegisterTaskFails() throws InterruptedException {
    CountDownLatch latch = new CountDownLatch(1);
    final Client mockClient = mock(Client.class);
    final ThreadPool threadPool = mock(ThreadPool.class);
    when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY));
    when(threadPool.generic()).thenReturn(OpenSearchExecutors.newDirectExecutorService());
    when(mockClient.threadPool()).thenReturn(threadPool);
    when(mockClient.settings()).thenReturn(Settings.EMPTY);
    PersistentTasksService persistentTasksService = new PersistentTasksService(null, null, mockClient) {

        @Override
        public void sendCompletionRequest(String taskId, long taskAllocationId, Exception taskFailure, ActionListener<PersistentTask<?>> listener) {
            assertThat(taskFailure, instanceOf(RuntimeException.class));
            assertThat(taskFailure.getMessage(), equalTo("Something went wrong"));
            listener.onResponse(mock(PersistentTask.class));
            latch.countDown();
        }
    };
    @SuppressWarnings("unchecked") PersistentTasksExecutor<TestParams> action = mock(PersistentTasksExecutor.class);
    when(action.getExecutor()).thenReturn(ThreadPool.Names.SAME);
    when(action.getTaskName()).thenReturn(TestPersistentTasksExecutor.NAME);
    when(action.createTask(anyLong(), anyString(), anyString(), any(), any(), any())).thenThrow(new RuntimeException("Something went wrong"));
    PersistentTasksExecutorRegistry registry = new PersistentTasksExecutorRegistry(Collections.singletonList(action));
    MockExecutor executor = new MockExecutor();
    PersistentTasksNodeService coordinator = new PersistentTasksNodeService(persistentTasksService, registry, new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet()), executor);
    ClusterState state = createInitialClusterState(0, Settings.EMPTY);
    PersistentTasksCustomMetadata.Builder tasks = PersistentTasksCustomMetadata.builder();
    tasks.addTask(UUIDs.base64UUID(), TestPersistentTasksExecutor.NAME, new TestParams("this_param"), new Assignment("this_node", "test assignment on this node"));
    Metadata.Builder metadata = Metadata.builder(state.metadata());
    metadata.putCustom(PersistentTasksCustomMetadata.TYPE, tasks.build());
    ClusterState newClusterState = ClusterState.builder(state).metadata(metadata).build();
    coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
    // Failed to start the task, make sure it wasn't invoked further
    assertThat(executor.executions.size(), equalTo(0));
    assertTrue(latch.await(5, TimeUnit.SECONDS));
}

Example 8 with PersistentTask

use of org.opensearch.persistent.PersistentTasksCustomMetadata.PersistentTask in project OpenSearch by opensearch-project.

the class PersistentTasksService method waitForPersistentTaskCondition.

/**
 * Waits for a given persistent task to comply with a given predicate, then call back the listener accordingly.
 *
 * @param taskId the persistent task id
 * @param predicate the persistent task predicate to evaluate
 * @param timeout a timeout for waiting
 * @param listener the callback listener
 */
public void waitForPersistentTaskCondition(final String taskId, final Predicate<PersistentTask<?>> predicate, @Nullable final TimeValue timeout, final WaitForPersistentTaskListener<?> listener) {
    final Predicate<ClusterState> clusterStatePredicate = clusterState -> predicate.test(PersistentTasksCustomMetadata.getTaskWithId(clusterState, taskId));
    final ClusterStateObserver observer = new ClusterStateObserver(clusterService, timeout, logger, threadPool.getThreadContext());
    final ClusterState clusterState = observer.setAndGetObservedState();
    if (clusterStatePredicate.test(clusterState)) {
        listener.onResponse(PersistentTasksCustomMetadata.getTaskWithId(clusterState, taskId));
    } else {
        observer.waitForNextChange(new ClusterStateObserver.Listener() {

            @Override
            public void onNewClusterState(ClusterState state) {
                listener.onResponse(PersistentTasksCustomMetadata.getTaskWithId(state, taskId));
            }

            @Override
            public void onClusterServiceClose() {
                listener.onFailure(new NodeClosedException(clusterService.localNode()));
            }

            @Override
            public void onTimeout(TimeValue timeout) {
                listener.onTimeout(timeout);
            }
        }, clusterStatePredicate);
    }
}

Example 9 with PersistentTask

use of org.opensearch.persistent.PersistentTasksCustomMetadata.PersistentTask in project OpenSearch by opensearch-project.

the class PersistentTasksExecutorIT method testPersistentActionStatusUpdate.

public void testPersistentActionStatusUpdate() throws Exception {
    PersistentTasksService persistentTasksService = internalCluster().getInstance(PersistentTasksService.class);
    PlainActionFuture<PersistentTask<TestParams>> future = new PlainActionFuture<>();
    persistentTasksService.sendStartRequest(UUIDs.base64UUID(), TestPersistentTasksExecutor.NAME, new TestParams("Blah"), future);
    String taskId = future.get().getId();
    waitForTaskToStart();
    TaskInfo firstRunningTask = client().admin().cluster().prepareListTasks().setActions(TestPersistentTasksExecutor.NAME + "[c]").get().getTasks().get(0);
    PersistentTasksCustomMetadata tasksInProgress = internalCluster().clusterService().state().getMetadata().custom(PersistentTasksCustomMetadata.TYPE);
    assertThat(tasksInProgress.tasks().size(), equalTo(1));
    assertThat(tasksInProgress.tasks().iterator().next().getState(), nullValue());
    int numberOfUpdates = randomIntBetween(1, 10);
    for (int i = 0; i < numberOfUpdates; i++) {
        logger.info("Updating the task states");
        // Complete the running task and make sure it finishes properly
        assertThat(new TestTasksRequestBuilder(client()).setOperation("update_status").setTaskId(firstRunningTask.getTaskId()).get().getTasks().size(), equalTo(1));
        int finalI = i;
        WaitForPersistentTaskFuture<?> future1 = new WaitForPersistentTaskFuture<>();
        persistentTasksService.waitForPersistentTaskCondition(taskId, task -> task != null && task.getState() != null && task.getState().toString() != null && task.getState().toString().equals("{\"phase\":\"phase " + (finalI + 1) + "\"}"), TimeValue.timeValueSeconds(10), future1);
        assertThat(future1.get().getId(), equalTo(taskId));
    }
    WaitForPersistentTaskFuture<?> future1 = new WaitForPersistentTaskFuture<>();
    persistentTasksService.waitForPersistentTaskCondition(taskId, task -> false, TimeValue.timeValueMillis(10), future1);
    assertFutureThrows(future1, IllegalStateException.class, "timed out after 10ms");
    PlainActionFuture<PersistentTask<?>> failedUpdateFuture = new PlainActionFuture<>();
    persistentTasksService.sendUpdateStateRequest(taskId, -2, new State("should fail"), failedUpdateFuture);
    assertFutureThrows(failedUpdateFuture, ResourceNotFoundException.class, "the task with id " + taskId + " and allocation id -2 doesn't exist");
    // Wait for the task to disappear
    WaitForPersistentTaskFuture<?> future2 = new WaitForPersistentTaskFuture<>();
    persistentTasksService.waitForPersistentTaskCondition(taskId, Objects::isNull, TimeValue.timeValueSeconds(10), future2);
    logger.info("Completing the running task");
    // Complete the running task and make sure it finishes properly
    assertThat(new TestTasksRequestBuilder(client()).setOperation("finish").setTaskId(firstRunningTask.getTaskId()).get().getTasks().size(), equalTo(1));
    assertThat(future2.get(), nullValue());
}

Example 10 with PersistentTask

use of org.opensearch.persistent.PersistentTasksCustomMetadata.PersistentTask in project OpenSearch by opensearch-project.

the class PersistentTasksExecutorIT method testPersistentActionCompletion.

public void testPersistentActionCompletion() throws Exception {
    PersistentTasksService persistentTasksService = internalCluster().getInstance(PersistentTasksService.class);
    PlainActionFuture<PersistentTask<TestParams>> future = new PlainActionFuture<>();
    String taskId = UUIDs.base64UUID();
    persistentTasksService.sendStartRequest(taskId, TestPersistentTasksExecutor.NAME, new TestParams("Blah"), future);
    long allocationId = future.get().getAllocationId();
    assertBusy(() -> {
        // Wait for the task to start
        assertThat(client().admin().cluster().prepareListTasks().setActions(TestPersistentTasksExecutor.NAME + "[c]").get().getTasks().size(), equalTo(1));
    });
    TaskInfo firstRunningTask = client().admin().cluster().prepareListTasks().setActions(TestPersistentTasksExecutor.NAME + "[c]").setDetailed(true).get().getTasks().get(0);
    logger.info("Found running task with id {} and parent {}", firstRunningTask.getId(), firstRunningTask.getParentTaskId());
    // Verifying parent and description
    assertThat(firstRunningTask.getParentTaskId().getId(), equalTo(allocationId));
    assertThat(firstRunningTask.getParentTaskId().getNodeId(), equalTo("cluster"));
    assertThat(firstRunningTask.getDescription(), equalTo("id=" + taskId));
    if (randomBoolean()) {
        logger.info("Simulating errant completion notification");
        // try sending completion request with incorrect allocation id
        PlainActionFuture<PersistentTask<?>> failedCompletionNotificationFuture = new PlainActionFuture<>();
        persistentTasksService.sendCompletionRequest(taskId, Long.MAX_VALUE, null, failedCompletionNotificationFuture);
        assertFutureThrows(failedCompletionNotificationFuture, ResourceNotFoundException.class);
        // Make sure that the task is still running
        assertThat(client().admin().cluster().prepareListTasks().setActions(TestPersistentTasksExecutor.NAME + "[c]").setDetailed(true).get().getTasks().size(), equalTo(1));
    }
    stopOrCancelTask(firstRunningTask.getTaskId());
}