Extend tracking of parent tasks to master node, replication and broadcast actions

Now MasterNodeOperations, ReplicationAllShards, ReplicationSingleShard,  BroadcastReplication and BroadcastByNode actions keep track of their parent tasks.

Author: imotov

core/src/main/java/org/elasticsearch/action/admin/indices/upgrade/post/TransportUpgradeAction.java

@@ -41,6 +41,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -179,7 +180,7 @@ protected ClusterBlockException checkRequestBlock(ClusterState state, UpgradeReq
     }
 
     @Override
-    protected void doExecute(UpgradeRequest request, final ActionListener<UpgradeResponse> listener) {
+    protected void doExecute(Task task, UpgradeRequest request, final ActionListener<UpgradeResponse> listener) {
         ActionListener<UpgradeResponse> settingsUpdateListener = new ActionListener<UpgradeResponse>() {
             @Override
             public void onResponse(UpgradeResponse upgradeResponse) {
@@ -199,7 +200,7 @@ public void onFailure(Throwable e) {
                 listener.onFailure(e);
             }
         };
-        super.doExecute(request, settingsUpdateListener);
+        super.doExecute(task, request, settingsUpdateListener);
     }
 
     private void updateSettings(final UpgradeResponse upgradeResponse, final ActionListener<UpgradeResponse> listener) {

core/src/main/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryAction.java

@@ -52,6 +52,7 @@
 import org.elasticsearch.search.internal.DefaultSearchContext;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.internal.ShardSearchLocalRequest;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -89,9 +90,9 @@ public TransportValidateQueryAction(Settings settings, ThreadPool threadPool, Cl
     }
 
     @Override
-    protected void doExecute(ValidateQueryRequest request, ActionListener<ValidateQueryResponse> listener) {
+    protected void doExecute(Task task, ValidateQueryRequest request, ActionListener<ValidateQueryResponse> listener) {
         request.nowInMillis = System.currentTimeMillis();
-        super.doExecute(request, listener);
+        super.doExecute(task, request, listener);
     }
 
     @Override

core/src/main/java/org/elasticsearch/action/delete/TransportDeleteAction.java

@@ -44,6 +44,7 @@
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.IndexAlreadyExistsException;
 import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -69,27 +70,27 @@ public TransportDeleteAction(Settings settings, TransportService transportServic
     }
 
     @Override
-    protected void doExecute(final DeleteRequest request, final ActionListener<DeleteResponse> listener) {
+    protected void doExecute(Task task, final DeleteRequest request, final ActionListener<DeleteResponse> listener) {
         ClusterState state = clusterService.state();
         if (autoCreateIndex.shouldAutoCreate(request.index(), state)) {
-            createIndexAction.execute(new CreateIndexRequest().index(request.index()).cause("auto(delete api)").masterNodeTimeout(request.timeout()), new ActionListener<CreateIndexResponse>() {
+            createIndexAction.execute(task, new CreateIndexRequest().index(request.index()).cause("auto(delete api)").masterNodeTimeout(request.timeout()), new ActionListener<CreateIndexResponse>() {
                 @Override
                 public void onResponse(CreateIndexResponse result) {
-                    innerExecute(request, listener);
+                    innerExecute(task, request, listener);
                 }
 
                 @Override
                 public void onFailure(Throwable e) {
                     if (ExceptionsHelper.unwrapCause(e) instanceof IndexAlreadyExistsException) {
                         // we have the index, do it
-                        innerExecute(request, listener);
+                        innerExecute(task, request, listener);
                     } else {
                         listener.onFailure(e);
                     }
                 }
             });
         } else {
-            innerExecute(request, listener);
+            innerExecute(task, request, listener);
         }
     }
 
@@ -114,8 +115,8 @@ protected void resolveRequest(final MetaData metaData, String concreteIndex, Del
         request.setShardId(shardId);
     }
 
-    private void innerExecute(final DeleteRequest request, final ActionListener<DeleteResponse> listener) {
-        super.doExecute(request, listener);
+    private void innerExecute(Task task, final DeleteRequest request, final ActionListener<DeleteResponse> listener) {
+        super.doExecute(task, request, listener);
     }
 
     @Override

core/src/main/java/org/elasticsearch/action/index/TransportIndexAction.java

@@ -48,6 +48,7 @@
 import org.elasticsearch.index.translog.Translog;
 import org.elasticsearch.indices.IndexAlreadyExistsException;
 import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -84,7 +85,7 @@ public TransportIndexAction(Settings settings, TransportService transportService
     }
 
     @Override
-    protected void doExecute(final IndexRequest request, final ActionListener<IndexResponse> listener) {
+    protected void doExecute(Task task, final IndexRequest request, final ActionListener<IndexResponse> listener) {
         // if we don't have a master, we don't have metadata, that's fine, let it find a master using create index API
         ClusterState state = clusterService.state();
         if (autoCreateIndex.shouldAutoCreate(request.index(), state)) {
@@ -93,18 +94,18 @@ protected void doExecute(final IndexRequest request, final ActionListener<IndexR
             createIndexRequest.mapping(request.type());
             createIndexRequest.cause("auto(index api)");
             createIndexRequest.masterNodeTimeout(request.timeout());
-            createIndexAction.execute(createIndexRequest, new ActionListener<CreateIndexResponse>() {
+            createIndexAction.execute(task, createIndexRequest, new ActionListener<CreateIndexResponse>() {
                 @Override
                 public void onResponse(CreateIndexResponse result) {
-                    innerExecute(request, listener);
+                    innerExecute(task, request, listener);
                 }
 
                 @Override
                 public void onFailure(Throwable e) {
                     if (ExceptionsHelper.unwrapCause(e) instanceof IndexAlreadyExistsException) {
                         // we have the index, do it
                         try {
-                            innerExecute(request, listener);
+                            innerExecute(task, request, listener);
                         } catch (Throwable e1) {
                             listener.onFailure(e1);
                         }
@@ -114,7 +115,7 @@ public void onFailure(Throwable e) {
                 }
             });
         } else {
-            innerExecute(request, listener);
+            innerExecute(task, request, listener);
         }
     }
 
@@ -129,8 +130,8 @@ protected void resolveRequest(MetaData metaData, String concreteIndex, IndexRequ
         request.setShardId(shardId);
     }
 
-    private void innerExecute(final IndexRequest request, final ActionListener<IndexResponse> listener) {
-        super.doExecute(request, listener);
+    private void innerExecute(Task task, final IndexRequest request, final ActionListener<IndexResponse> listener) {
+        super.doExecute(task, request, listener);
     }
 
     @Override

core/src/main/java/org/elasticsearch/action/percolate/TransportPercolateAction.java

@@ -41,6 +41,7 @@
 import org.elasticsearch.index.engine.DocumentMissingException;
 import org.elasticsearch.percolator.PercolateException;
 import org.elasticsearch.percolator.PercolatorService;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -70,7 +71,7 @@ public TransportPercolateAction(Settings settings, ThreadPool threadPool, Cluste
     }
 
     @Override
-    protected void doExecute(final PercolateRequest request, final ActionListener<PercolateResponse> listener) {
+    protected void doExecute(Task task, final PercolateRequest request, final ActionListener<PercolateResponse> listener) {
         request.startTime = System.currentTimeMillis();
         if (request.getRequest() != null) {
             //create a new get request to make sure it has the same headers and context as the original percolate request
@@ -84,7 +85,7 @@ public void onResponse(GetResponse getResponse) {
                     }
 
                     BytesReference docSource = getResponse.getSourceAsBytesRef();
-                    TransportPercolateAction.super.doExecute(new PercolateRequest(request, docSource), listener);
+                    TransportPercolateAction.super.doExecute(task, new PercolateRequest(request, docSource), listener);
                 }
 
                 @Override
@@ -93,7 +94,7 @@ public void onFailure(Throwable e) {
                 }
             });
         } else {
-            super.doExecute(request, listener);
+            super.doExecute(task, request, listener);
         }
     }
 

core/src/main/java/org/elasticsearch/action/support/ChildTaskActionRequest.java

@@ -0,0 +1,66 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.action.support;
+
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.Task;
+
+import java.io.IOException;
+
+/**
+ * Base class for action requests that can have associated child tasks
+ */
+public abstract class ChildTaskActionRequest<Request extends ActionRequest<Request>> extends ActionRequest<Request> {
+
+    private String parentTaskNode;
+
+    private long parentTaskId;
+
+    protected ChildTaskActionRequest() {
+
+    }
+
+    public void setParentTask(String parentTaskNode, long parentTaskId) {
+        this.parentTaskNode = parentTaskNode;
+        this.parentTaskId = parentTaskId;
+    }
+
+    @Override
+    public void readFrom(StreamInput in) throws IOException {
+        super.readFrom(in);
+        parentTaskNode = in.readOptionalString();
+        parentTaskId = in.readLong();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeOptionalString(parentTaskNode);
+        out.writeLong(parentTaskId);
+    }
+
+    @Override
+    public Task createTask(long id, String type, String action) {
+        return new Task(id, type, action, this::getDescription, parentTaskNode, parentTaskId);
+    }
+
+}

core/src/main/java/org/elasticsearch/action/support/ChildTaskRequest.java

@@ -19,10 +19,8 @@
 
 package org.elasticsearch.action.support;
 
-import org.elasticsearch.Version;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.tasks.ChildTask;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportRequest;
 
@@ -61,6 +59,6 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     public Task createTask(long id, String type, String action) {
-        return new ChildTask(id, type, action, this::getDescription, parentTaskNode, parentTaskId);
+        return new Task(id, type, action, this::getDescription, parentTaskNode, parentTaskId);
     }
 }

core/src/main/java/org/elasticsearch/action/support/HandledTransportAction.java

@@ -44,13 +44,14 @@ protected HandledTransportAction(Settings settings, String actionName, ThreadPoo
     class TransportHandler implements TransportRequestHandler<Request> {
 
         @Override
-        public final void messageReceived(final Request request, final TransportChannel channel, Task task) throws Exception {
-            messageReceived(request, channel);
+        public final void messageReceived(Request request, TransportChannel channel) throws Exception {
+            throw new UnsupportedOperationException("the task parameter is required for this operation");
         }
 
         @Override
-        public final void messageReceived(Request request, TransportChannel channel) throws Exception {
-            execute(request, new ActionListener<Response>() {
+        public final void messageReceived(final Request request, final TransportChannel channel, Task task) throws Exception {
+            // We already got the task created on the netty layer - no need to create it again on the transport layer
+            execute(task, request, new ActionListener<Response>() {
                 @Override
                 public void onResponse(Response response) {
                     try {

core/src/main/java/org/elasticsearch/action/support/TransportAction.java

@@ -66,6 +66,11 @@ public final ActionFuture<Response> execute(Request request) {
         return future;
     }
 
+    /**
+     * Use this method when the transport action call should result in creation of a new task associated with the call.
+     *
+     * This is a typical behavior.
+     */
     public final Task execute(Request request, ActionListener<Response> listener) {
         Task task = taskManager.register("transport", actionName, request);
         if (task == null) {
@@ -88,7 +93,10 @@ public void onFailure(Throwable e) {
         return task;
     }
 
-    private final void execute(Task task, Request request, ActionListener<Response> listener) {
+    /**
+     * Use this method when the transport action should continue to run in the context of the current task
+     */
+    public final void execute(Task task, Request request, ActionListener<Response> listener) {
 
         ActionRequestValidationException validationException = request.validate();
         if (validationException != null) {

core/src/main/java/org/elasticsearch/action/support/broadcast/BroadcastShardRequest.java

@@ -21,18 +21,18 @@
 
 import org.elasticsearch.action.IndicesRequest;
 import org.elasticsearch.action.OriginalIndices;
+import org.elasticsearch.action.support.ChildTaskRequest;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.index.shard.ShardId;
-import org.elasticsearch.transport.TransportRequest;
 
 import java.io.IOException;
 
 /**
  *
  */
-public abstract class BroadcastShardRequest extends TransportRequest implements IndicesRequest {
+public abstract class BroadcastShardRequest extends ChildTaskRequest implements IndicesRequest {
 
     private ShardId shardId;
 

core/src/main/java/org/elasticsearch/action/support/broadcast/TransportBroadcastAction.java

@@ -35,6 +35,7 @@
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.BaseTransportResponseHandler;
 import org.elasticsearch.transport.TransportChannel;
@@ -69,8 +70,13 @@ protected TransportBroadcastAction(Settings settings, String actionName, ThreadP
     }
 
     @Override
-    protected void doExecute(Request request, ActionListener<Response> listener) {
-        new AsyncBroadcastAction(request, listener).start();
+    protected void doExecute(Task task, Request request, ActionListener<Response> listener) {
+        new AsyncBroadcastAction(task, request, listener).start();
+    }
+
+    @Override
+    protected final void doExecute(Request request, ActionListener<Response> listener) {
+        throw new UnsupportedOperationException("the task parameter is required for this operation");
     }
 
     protected abstract Response newResponse(Request request, AtomicReferenceArray shardsResponses, ClusterState clusterState);
@@ -93,6 +99,7 @@ protected void doExecute(Request request, ActionListener<Response> listener) {
 
     protected class AsyncBroadcastAction {
 
+        private final Task task;
         private final Request request;
         private final ActionListener<Response> listener;
         private final ClusterState clusterState;
@@ -102,7 +109,8 @@ protected class AsyncBroadcastAction {
         private final AtomicInteger counterOps = new AtomicInteger();
         private final AtomicReferenceArray shardsResponses;
 
-        protected AsyncBroadcastAction(Request request, ActionListener<Response> listener) {
+        protected AsyncBroadcastAction(Task task, Request request, ActionListener<Response> listener) {
+            this.task = task;
             this.request = request;
             this.listener = listener;
 
@@ -158,6 +166,7 @@ protected void performOperation(final ShardIterator shardIt, final ShardRouting
             } else {
                 try {
                     final ShardRequest shardRequest = newShardRequest(shardIt.size(), shard, request);
+                    shardRequest.setParentTask(clusterService.localNode().getId(), task.getId());
                     DiscoveryNode node = nodes.get(shard.currentNodeId());
                     if (node == null) {
                         // no node connected, act as failure