Make the persistent task status available to PersistentTasksExecutor.nodeOperation(...) method

Author: martijnvg

server/src/main/java/org/elasticsearch/persistent/NodePersistentTasksExecutor.java

@@ -20,6 +20,7 @@
 
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 
 /**
@@ -35,6 +36,7 @@ public NodePersistentTasksExecutor(ThreadPool threadPool) {
     }
 
     public <Params extends PersistentTaskParams> void executeTask(@Nullable Params params,
+                                                                  @Nullable Task.Status status,
                                                                   AllocatedPersistentTask task,
                                                                   PersistentTasksExecutor<Params> executor) {
         threadPool.executor(executor.getExecutor()).execute(new AbstractRunnable() {
@@ -47,7 +49,7 @@ public void onFailure(Exception e) {
             @Override
             protected void doRun() throws Exception {
                 try {
-                    executor.nodeOperation(task, params);
+                    executor.nodeOperation(task, params, status);
                 } catch (Exception ex) {
                     task.markAsFailed(ex);
                 }

server/src/main/java/org/elasticsearch/persistent/PersistentTasksExecutor.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.component.AbstractComponent;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.persistent.PersistentTasksCustomMetaData.Assignment;
 import org.elasticsearch.persistent.PersistentTasksCustomMetaData.PersistentTask;
@@ -115,10 +116,10 @@ protected String getDescription(PersistentTask<Params> taskInProgress) {
     /**
      * This operation will be executed on the executor node.
      * <p>
-     * NOTE: The nodeOperation has to throws an exception, trigger task.markAsCompleted() or task.completeAndNotifyIfNeeded() methods to
+     * NOTE: The nodeOperation has to throw an exception, trigger task.markAsCompleted() or task.completeAndNotifyIfNeeded() methods to
      * indicate that the persistent task has finished.
      */
-    protected abstract void nodeOperation(AllocatedPersistentTask task, @Nullable Params params);
+    protected abstract void nodeOperation(AllocatedPersistentTask task, @Nullable Params params, @Nullable Task.Status status);
 
     public String getExecutor() {
         return executor;

server/src/main/java/org/elasticsearch/persistent/PersistentTasksNodeService.java

@@ -173,7 +173,7 @@ public Task createTask(long id, String type, String action, TaskId parentTaskId)
                     task.getPersistentTaskId(), task.getAllocationId());
             try {
                 runningTasks.put(taskInProgress.getAllocationId(), task);
-                nodePersistentTasksExecutor.executeTask(taskInProgress.getParams(), task, executor);
+                nodePersistentTasksExecutor.executeTask(taskInProgress.getParams(), taskInProgress.getStatus(), task, executor);
             } catch (Exception e) {
                 // Submit task failure
                 task.markAsFailed(e);

server/src/test/java/org/elasticsearch/persistent/PersistentTasksNodeServiceTests.java

@@ -49,12 +49,14 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.sameInstance;
 import static org.hamcrest.core.IsEqual.equalTo;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyLong;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class PersistentTasksNodeServiceTests extends ESTestCase {
@@ -167,6 +169,41 @@ public void testStartTask() throws Exception {
 
     }
 
+    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);
+        when(action.createTask(anyLong(), anyString(), anyString(), eq(parentId), any())).thenReturn(nodeTask);
+        PersistentTasksExecutorRegistry registry = new PersistentTasksExecutorRegistry(Settings.EMPTY, Collections.singletonList(action));
+
+        MockExecutor executor = new MockExecutor();
+        PersistentTasksNodeService coordinator = new PersistentTasksNodeService(Settings.EMPTY, persistentTasksService,
+                registry, new TaskManager(Settings.EMPTY), executor);
+
+        ClusterState state = createInitialClusterState(1, Settings.EMPTY);
+
+        Task.Status status = new TestPersistentTasksPlugin.Status("_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.updateTaskStatus(taskId, status);
+        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).status, sameInstance(status));
+        assertThat(executor.get(0).task, sameInstance(nodeTask));
+    }
+
     public void testTaskCancellation() {
         AtomicLong capturedTaskId = new AtomicLong();
         AtomicReference<ActionListener<CancelTasksResponse>> capturedListener = new AtomicReference<>();
@@ -271,11 +308,13 @@ private ClusterState removeTask(ClusterState state, String taskId) {
     private class Execution {
         private final PersistentTaskParams params;
         private final AllocatedPersistentTask task;
+        private final Task.Status status;
         private final PersistentTasksExecutor<?> holder;
 
-        Execution(PersistentTaskParams params, AllocatedPersistentTask task, PersistentTasksExecutor<?> holder) {
+        Execution(PersistentTaskParams params, AllocatedPersistentTask task, Task.Status status, PersistentTasksExecutor<?> holder) {
             this.params = params;
             this.task = task;
+            this.status = status;
             this.holder = holder;
         }
     }
@@ -288,9 +327,11 @@ private class MockExecutor extends NodePersistentTasksExecutor {
         }
 
         @Override
-        public <Params extends PersistentTaskParams> void executeTask(Params params, AllocatedPersistentTask task,
+        public <Params extends PersistentTaskParams> void executeTask(Params params,
+                                                                      Task.Status status,
+                                                                      AllocatedPersistentTask task,
                                                                       PersistentTasksExecutor<Params> executor) {
-            executions.add(new Execution(params, task, executor));
+            executions.add(new Execution(params, task, status, executor));
         }
 
         public Execution get(int i) {

server/src/test/java/org/elasticsearch/persistent/TestPersistentTasksPlugin.java

@@ -323,7 +323,7 @@ public Assignment getAssignment(TestParams params, ClusterState clusterState) {
         }
 
         @Override
-        protected void nodeOperation(AllocatedPersistentTask task, TestParams params) {
+        protected void nodeOperation(AllocatedPersistentTask task, TestParams params, Task.Status status) {
             logger.info("started node operation for the task {}", task);
             try {
                 TestTask testTask = (TestTask) task;
@@ -346,9 +346,9 @@ protected void nodeOperation(AllocatedPersistentTask task, TestParams params) {
                     } else if ("update_status".equals(testTask.getOperation())) {
                         testTask.setOperation(null);
                         CountDownLatch latch = new CountDownLatch(1);
-                        Status status = new Status("phase " + phase.incrementAndGet());
-                        logger.info("updating the task status to {}", status);
-                        task.updatePersistentStatus(status, new ActionListener<PersistentTask<?>>() {
+                        Status newStatus = new Status("phase " + phase.incrementAndGet());
+                        logger.info("updating the task status to {}", newStatus);
+                        task.updatePersistentStatus(newStatus, new ActionListener<PersistentTask<?>>() {
                             @Override
                             public void onResponse(PersistentTask<?> persistentTask) {
                                 logger.info("updating was successful");