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

Example 1 with TaskAwareRequest

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

the class PersistentTasksNodeService method startTask.

private <Params extends PersistentTaskParams> void startTask(PersistentTask<Params> taskInProgress) {
    PersistentTasksExecutor<Params> executor = persistentTasksExecutorRegistry.getPersistentTaskExecutorSafe(taskInProgress.getTaskName());
    TaskAwareRequest request = new TaskAwareRequest() {

        TaskId parentTaskId = new TaskId("cluster", taskInProgress.getAllocationId());

        @Override
        public void setParentTask(TaskId taskId) {
            throw new UnsupportedOperationException("parent task if for persistent tasks shouldn't change");
        }

        @Override
        public TaskId getParentTask() {
            return parentTaskId;
        }

        @Override
        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
            return executor.createTask(id, type, action, parentTaskId, taskInProgress, headers);
        }
    };
    AllocatedPersistentTask task;
    try {
        task = (AllocatedPersistentTask) taskManager.register("persistent", taskInProgress.getTaskName() + "[c]", request);
    } catch (Exception e) {
        logger.error("Fatal error registering persistent task [" + taskInProgress.getTaskName() + "] with id [" + taskInProgress.getId() + "] and allocation id [" + taskInProgress.getAllocationId() + "], removing from persistent tasks", e);
        notifyMasterOfFailedTask(taskInProgress, e);
        return;
    }
    boolean processed = false;
    try {
        task.init(persistentTasksService, taskManager, taskInProgress.getId(), taskInProgress.getAllocationId());
        logger.trace("Persistent task [{}] with id [{}] and allocation id [{}] was created", task.getAction(), task.getPersistentTaskId(), task.getAllocationId());
        try {
            runningTasks.put(taskInProgress.getAllocationId(), task);
            nodePersistentTasksExecutor.executeTask(taskInProgress.getParams(), taskInProgress.getState(), task, executor);
        } catch (Exception e) {
            // Submit task failure
            task.markAsFailed(e);
        }
        processed = true;
    } finally {
        if (processed == false) {
            // something went wrong - unregistering task
            logger.warn("Persistent task [{}] with id [{}] and allocation id [{}] failed to create", task.getAction(), task.getPersistentTaskId(), task.getAllocationId());
            taskManager.unregister(task);
        }
    }
}

Example 2 with TaskAwareRequest

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

the class TransportServiceDeserializationFailureTests method testDeserializationFailureLogIdentifiesListener.

public void testDeserializationFailureLogIdentifiesListener() {
    final DiscoveryNode localNode = new DiscoveryNode("local", buildNewFakeTransportAddress(), Version.CURRENT);
    final DiscoveryNode otherNode = new DiscoveryNode("other", buildNewFakeTransportAddress(), Version.CURRENT);
    final Settings settings = Settings.builder().put(NODE_NAME_SETTING.getKey(), "local").build();
    final DeterministicTaskQueue deterministicTaskQueue = new DeterministicTaskQueue(settings, random());
    final String testActionName = "internal:test-action";
    final MockTransport transport = new MockTransport() {

        @Override
        protected void onSendRequest(long requestId, String action, TransportRequest request, DiscoveryNode node) {
            if (action.equals(TransportService.HANDSHAKE_ACTION_NAME)) {
                handleResponse(requestId, new TransportService.HandshakeResponse(otherNode, new ClusterName(""), Version.CURRENT));
            }
        }
    };
    final TransportService transportService = transport.createTransportService(Settings.EMPTY, deterministicTaskQueue.getThreadPool(), TransportService.NOOP_TRANSPORT_INTERCEPTOR, ignored -> localNode, null, Collections.emptySet());
    transportService.registerRequestHandler(testActionName, ThreadPool.Names.SAME, TransportRequest.Empty::new, (request, channel, task) -> channel.sendResponse(TransportResponse.Empty.INSTANCE));
    transportService.start();
    transportService.acceptIncomingRequests();
    final PlainActionFuture<Void> connectionFuture = new PlainActionFuture<>();
    transportService.connectToNode(otherNode, connectionFuture);
    assertTrue(connectionFuture.isDone());
    {
        // requests without a parent task are recorded directly in the response context
        transportService.sendRequest(otherNode, testActionName, TransportRequest.Empty.INSTANCE, TransportRequestOptions.EMPTY, new TransportResponseHandler<TransportResponse.Empty>() {

            @Override
            public void handleResponse(TransportResponse.Empty response) {
                fail("should not be called");
            }

            @Override
            public void handleException(TransportException exp) {
                fail("should not be called");
            }

            @Override
            public String executor() {
                return ThreadPool.Names.SAME;
            }

            @Override
            public TransportResponse.Empty read(StreamInput in) {
                throw new AssertionError("should not be called");
            }

            @Override
            public String toString() {
                return "test handler without parent";
            }
        });
        final List<Transport.ResponseContext<? extends TransportResponse>> responseContexts = transport.getResponseHandlers().prune(ignored -> true);
        assertThat(responseContexts, hasSize(1));
        final TransportResponseHandler<? extends TransportResponse> handler = responseContexts.get(0).handler();
        assertThat(handler, hasToString(containsString("test handler without parent")));
    }
    {
        // requests with a parent task get wrapped up by the transport service, including the action name
        final Task parentTask = transportService.getTaskManager().register("test", "test-action", new TaskAwareRequest() {

            @Override
            public void setParentTask(TaskId taskId) {
                fail("should not be called");
            }

            @Override
            public TaskId getParentTask() {
                return TaskId.EMPTY_TASK_ID;
            }
        });
        transportService.sendChildRequest(otherNode, testActionName, TransportRequest.Empty.INSTANCE, parentTask, TransportRequestOptions.EMPTY, new TransportResponseHandler<TransportResponse.Empty>() {

            @Override
            public void handleResponse(TransportResponse.Empty response) {
                fail("should not be called");
            }

            @Override
            public void handleException(TransportException exp) {
                fail("should not be called");
            }

            @Override
            public String executor() {
                return ThreadPool.Names.SAME;
            }

            @Override
            public TransportResponse.Empty read(StreamInput in) {
                throw new AssertionError("should not be called");
            }

            @Override
            public String toString() {
                return "test handler with parent";
            }
        });
        final List<Transport.ResponseContext<? extends TransportResponse>> responseContexts = transport.getResponseHandlers().prune(ignored -> true);
        assertThat(responseContexts, hasSize(1));
        final TransportResponseHandler<? extends TransportResponse> handler = responseContexts.get(0).handler();
        assertThat(handler, hasToString(allOf(containsString("test handler with parent"), containsString(testActionName))));
    }
}