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

Example 6 with Assignment

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

the class PersistentTasksNodeServiceTests method testStartTask.

public void testStartTask() {
    PersistentTasksService persistentTasksService = mock(PersistentTasksService.class);
    @SuppressWarnings("unchecked") PersistentTasksExecutor<TestParams> action = mock(PersistentTasksExecutor.class);
    when(action.getExecutor()).thenReturn(ThreadPool.Names.SAME);
    when(action.getTaskName()).thenReturn(TestPersistentTasksExecutor.NAME);
    int nonLocalNodesCount = randomInt(10);
    // need to account for 5 original tasks on each node and their relocations
    for (int i = 0; i < (nonLocalNodesCount + 1) * 10; i++) {
        TaskId parentId = new TaskId("cluster", i);
        when(action.createTask(anyLong(), anyString(), anyString(), eq(parentId), any(), any())).thenReturn(new TestPersistentTasksPlugin.TestTask(i, "persistent", "test", "", parentId, Collections.emptyMap()));
    }
    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(nonLocalNodesCount, Settings.EMPTY);
    PersistentTasksCustomMetadata.Builder tasks = PersistentTasksCustomMetadata.builder();
    boolean added = false;
    if (nonLocalNodesCount > 0) {
        for (int i = 0; i < randomInt(5); i++) {
            tasks.addTask(UUIDs.base64UUID(), TestPersistentTasksExecutor.NAME, new TestParams("other_" + i), new Assignment("other_node_" + randomInt(nonLocalNodesCount), "test assignment on other node"));
            if (added == false && randomBoolean()) {
                added = true;
                tasks.addTask(UUIDs.base64UUID(), TestPersistentTasksExecutor.NAME, new TestParams("this_param"), new Assignment("this_node", "test assignment on this node"));
            }
        }
    }
    if (added == false) {
        logger.info("No local node action was added");
    }
    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));
    if (added) {
        // ActionType for this node was added, let's make sure it was invoked
        assertThat(executor.executions.size(), equalTo(1));
        // Add task on some other node
        state = newClusterState;
        newClusterState = addTask(state, TestPersistentTasksExecutor.NAME, null, "some_other_node");
        coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
        // Make sure action wasn't called again
        assertThat(executor.executions.size(), equalTo(1));
        assertThat(executor.get(0).task.isCompleted(), is(false));
        // Start another task on this node
        state = newClusterState;
        newClusterState = addTask(state, TestPersistentTasksExecutor.NAME, new TestParams("this_param"), "this_node");
        coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
        // Make sure action was called this time
        assertThat(executor.size(), equalTo(2));
        assertThat(executor.get(1).task.isCompleted(), is(false));
        // Finish both tasks
        executor.get(0).task.markAsFailed(new RuntimeException());
        executor.get(1).task.markAsCompleted();
        assertThat(executor.get(0).task.isCompleted(), is(true));
        assertThat(executor.get(1).task.isCompleted(), is(true));
        String failedTaskId = executor.get(0).task.getPersistentTaskId();
        String finishedTaskId = executor.get(1).task.getPersistentTaskId();
        executor.clear();
        // Add task on some other node
        state = newClusterState;
        newClusterState = addTask(state, TestPersistentTasksExecutor.NAME, null, "some_other_node");
        coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
        // Make sure action wasn't called again
        assertThat(executor.size(), equalTo(0));
        // Simulate reallocation of the failed task on the same node
        state = newClusterState;
        newClusterState = reallocateTask(state, failedTaskId, "this_node");
        coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
        // Simulate removal of the finished task
        state = newClusterState;
        newClusterState = removeTask(state, finishedTaskId);
        coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
        // Make sure action was only allocated on this node once
        assertThat(executor.size(), equalTo(1));
    }
}

Example 7 with Assignment

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

the class PersistentTasksNodeServiceTests method testParamsStatusAndNodeTaskAreDelegated.

public void testParamsStatusAndNodeTaskAreDelegated() throws Exception {
    PersistentTasksService persistentTasksService = mock(PersistentTasksService.class);
    @SuppressWarnings("unchecked") PersistentTasksExecutor<TestParams> action = mock(PersistentTasksExecutor.class);
    when(action.getExecutor()).thenReturn(ThreadPool.Names.SAME);
    when(action.getTaskName()).thenReturn(TestPersistentTasksExecutor.NAME);
    TaskId parentId = new TaskId("cluster", 1);
    AllocatedPersistentTask nodeTask = new TestPersistentTasksPlugin.TestTask(0, "persistent", "test", "", parentId, Collections.emptyMap());
    when(action.createTask(anyLong(), anyString(), anyString(), eq(parentId), any(), any())).thenReturn(nodeTask);
    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(1, Settings.EMPTY);
    PersistentTaskState taskState = new TestPersistentTasksPlugin.State("_test_phase");
    PersistentTasksCustomMetadata.Builder tasks = PersistentTasksCustomMetadata.builder();
    String taskId = UUIDs.base64UUID();
    TestParams taskParams = new TestParams("other_0");
    tasks.addTask(taskId, TestPersistentTasksExecutor.NAME, taskParams, new Assignment("this_node", "test assignment on other node"));
    tasks.updateTaskState(taskId, taskState);
    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));
    assertThat(executor.size(), equalTo(1));
    assertThat(executor.get(0).params, sameInstance(taskParams));
    assertThat(executor.get(0).state, sameInstance(taskState));
    assertThat(executor.get(0).task, sameInstance(nodeTask));
}

Example 8 with Assignment

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

the class PersistentTasksNodeServiceTests method reallocateTask.

private ClusterState reallocateTask(ClusterState state, String taskId, String node) {
    PersistentTasksCustomMetadata.Builder builder = PersistentTasksCustomMetadata.builder(state.getMetadata().custom(PersistentTasksCustomMetadata.TYPE));
    assertTrue(builder.hasTask(taskId));
    return ClusterState.builder(state).metadata(Metadata.builder(state.metadata()).putCustom(PersistentTasksCustomMetadata.TYPE, builder.reassignTask(taskId, new Assignment(node, "test assignment")).build())).build();
}

Example 9 with Assignment

use of org.opensearch.persistent.PersistentTasksCustomMetadata.Assignment 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 10 with Assignment

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

the class PersistentTasksClusterService method shouldReassignPersistentTasks.

/**
 * Returns true if the cluster state change(s) require to reassign some persistent tasks. It can happen in the following
 * situations: a node left or is added, the routing table changed, the master node changed, the metadata changed or the
 * persistent tasks changed.
 */
boolean shouldReassignPersistentTasks(final ClusterChangedEvent event) {
    final PersistentTasksCustomMetadata tasks = event.state().getMetadata().custom(PersistentTasksCustomMetadata.TYPE);
    if (tasks == null) {
        return false;
    }
    boolean masterChanged = event.previousState().nodes().isLocalNodeElectedMaster() == false;
    if (persistentTasksChanged(event) || event.nodesChanged() || event.routingTableChanged() || event.metadataChanged() || masterChanged) {
        for (PersistentTask<?> task : tasks.tasks()) {
            if (needsReassignment(task.getAssignment(), event.state().nodes())) {
                Assignment assignment = createAssignment(task.getTaskName(), task.getParams(), event.state());
                if (Objects.equals(assignment, task.getAssignment()) == false) {
                    return true;
                }
            }
        }
    }
    return false;
}