[Transform] Allow task canceling of validate API calls (elastic#110951)

Validate will initiate a search request. In the event that the search request needs to be cancelled, rather than manually stopping the task, cancelling the Validate task will now propagate the cancel request to the Search task. Relate elastic#88010 Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
stefnestor · Jul 18, 2024 · 006f36e · 006f36e
1 parent 548aea5
commit 006f36e
Show file tree

Hide file tree

Showing 9 changed files with 67 additions and 16 deletions.
diff --git a/docs/changelog/110951.yaml b/docs/changelog/110951.yaml
@@ -0,0 +1,5 @@
+pr: 110951
+summary: Allow task canceling of validate API calls
+area: Transform
+type: bug
+issues: []
diff --git a/.../core/src/main/java/org/elasticsearch/xpack/core/transform/action/PutTransformAction.java b/.../core/src/main/java/org/elasticsearch/xpack/core/transform/action/PutTransformAction.java
@@ -14,6 +14,9 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.common.validation.SourceDestValidator;
 import org.elasticsearch.xpack.core.transform.TransformField;
@@ -22,6 +25,7 @@
 import org.elasticsearch.xpack.core.transform.utils.TransformStrings;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.Objects;
 
 import static org.elasticsearch.action.ValidateActions.addValidationError;
@@ -154,6 +158,11 @@ public boolean equals(Object obj) {
                 && this.deferValidation == other.deferValidation
                 && ackTimeout().equals(other.ackTimeout());
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, getDescription(), parentTaskId, headers);
+        }
     }
 
 }
diff --git a/...ore/src/main/java/org/elasticsearch/xpack/core/transform/action/StartTransformAction.java b/...ore/src/main/java/org/elasticsearch/xpack/core/transform/action/StartTransformAction.java
@@ -14,6 +14,9 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.transform.TransformField;
@@ -22,6 +25,7 @@
 import java.io.IOException;
 import java.time.Instant;
 import java.util.Collections;
+import java.util.Map;
 import java.util.Objects;
 
 public class StartTransformAction extends ActionType<StartTransformAction.Response> {
@@ -89,6 +93,11 @@ public boolean equals(Object obj) {
             // the base class does not implement equals, therefore we need to check timeout ourselves
             return Objects.equals(id, other.id) && Objects.equals(from, other.from) && ackTimeout().equals(other.ackTimeout());
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, getDescription(), parentTaskId, headers);
+        }
     }
 
     public static class Response extends BaseTasksResponse implements ToXContentObject {

diff --git a/.../src/main/java/org/elasticsearch/xpack/core/transform/action/ValidateTransformAction.java b/.../src/main/java/org/elasticsearch/xpack/core/transform/action/ValidateTransformAction.java
@@ -14,6 +14,9 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xpack.core.common.validation.SourceDestValidator;
 import org.elasticsearch.xpack.core.transform.transforms.TransformConfig;
 
@@ -94,6 +97,11 @@ public int hashCode() {
             // the base class does not implement hashCode, therefore we need to hash timeout ourselves
             return Objects.hash(ackTimeout(), config, deferValidation);
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, getDescription(), parentTaskId, headers);
+        }
     }
 
     public static class Response extends ActionResponse {

diff --git a/...m/src/main/java/org/elasticsearch/xpack/transform/action/TransportPutTransformAction.java b/...m/src/main/java/org/elasticsearch/xpack/transform/action/TransportPutTransformAction.java
@@ -15,6 +15,7 @@
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.action.support.master.AcknowledgedTransportMasterNodeAction;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.client.internal.ParentTaskAssigningClient;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
@@ -25,6 +26,7 @@
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.persistent.PersistentTasksCustomMetadata;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.ClientHelper;
@@ -110,9 +112,11 @@ protected void masterOperation(Task task, Request request, ClusterState clusterS
         );
 
         // <2> Validate source and destination indices
+
+        var parentTaskId = new TaskId(clusterService.localNode().getId(), task.getId());
         ActionListener<Void> checkPrivilegesListener = validateTransformListener.delegateFailureAndWrap(
             (l, aVoid) -> ClientHelper.executeAsyncWithOrigin(
-                client,
+                new ParentTaskAssigningClient(client, parentTaskId),
                 ClientHelper.TRANSFORM_ORIGIN,
                 ValidateTransformAction.INSTANCE,
                 new ValidateTransformAction.Request(config, request.isDeferValidation(), request.ackTimeout()),

diff --git a/...src/main/java/org/elasticsearch/xpack/transform/action/TransportStartTransformAction.java b/...src/main/java/org/elasticsearch/xpack/transform/action/TransportStartTransformAction.java
@@ -17,6 +17,7 @@
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.master.TransportMasterNodeAction;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.client.internal.ParentTaskAssigningClient;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
@@ -31,6 +32,7 @@
 import org.elasticsearch.persistent.PersistentTasksService;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.ClientHelper;
@@ -126,23 +128,25 @@ protected TransportStartTransformAction(
 
     @Override
     protected void masterOperation(
-        Task ignoredTask,
+        Task task,
         StartTransformAction.Request request,
         ClusterState state,
         ActionListener<StartTransformAction.Response> listener
     ) {
         TransformNodes.warnIfNoTransformNodes(state);
 
-        final SetOnce<TransformTaskParams> transformTaskParamsHolder = new SetOnce<>();
-        final SetOnce<TransformConfig> transformConfigHolder = new SetOnce<>();
+        var transformTaskParamsHolder = new SetOnce<TransformTaskParams>();
+        var transformConfigHolder = new SetOnce<TransformConfig>();
+        var parentTaskId = new TaskId(clusterService.localNode().getId(), task.getId());
+        var parentClient = new ParentTaskAssigningClient(client, parentTaskId);
 
         // <5> Wait for the allocated task's state to STARTED
         ActionListener<PersistentTasksCustomMetadata.PersistentTask<TransformTaskParams>> newPersistentTaskActionListener = ActionListener
-            .wrap(task -> {
+            .wrap(t -> {
                 TransformTaskParams transformTask = transformTaskParamsHolder.get();
                 assert transformTask != null;
                 waitForTransformTaskStarted(
-                    task.getId(),
+                    t.getId(),
                     transformTask,
                     request.ackTimeout(),
                     ActionListener.wrap(taskStarted -> listener.onResponse(new StartTransformAction.Response(true)), listener::onFailure)
@@ -196,7 +200,7 @@ protected void masterOperation(
                 return;
             }
             TransformIndex.createDestinationIndex(
-                client,
+                parentClient,
                 auditor,
                 indexNameExpressionResolver,
                 state,
@@ -257,7 +261,7 @@ protected void masterOperation(
                 )
             );
             ClientHelper.executeAsyncWithOrigin(
-                client,
+                parentClient,
                 ClientHelper.TRANSFORM_ORIGIN,
                 ValidateTransformAction.INSTANCE,
                 new ValidateTransformAction.Request(config, false, request.ackTimeout()),

diff --git a/.../main/java/org/elasticsearch/xpack/transform/action/TransportValidateTransformAction.java b/.../main/java/org/elasticsearch/xpack/transform/action/TransportValidateTransformAction.java
@@ -11,6 +11,7 @@
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.client.internal.ParentTaskAssigningClient;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -23,15 +24,14 @@
 import org.elasticsearch.license.License;
 import org.elasticsearch.license.RemoteClusterLicenseChecker;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.common.validation.SourceDestValidator;
 import org.elasticsearch.xpack.core.transform.TransformDeprecations;
 import org.elasticsearch.xpack.core.transform.TransformMessages;
 import org.elasticsearch.xpack.core.transform.action.ValidateTransformAction;
 import org.elasticsearch.xpack.core.transform.action.ValidateTransformAction.Request;
 import org.elasticsearch.xpack.core.transform.action.ValidateTransformAction.Response;
-import org.elasticsearch.xpack.core.transform.transforms.TransformConfig;
-import org.elasticsearch.xpack.transform.transforms.Function;
 import org.elasticsearch.xpack.transform.transforms.FunctionFactory;
 import org.elasticsearch.xpack.transform.transforms.TransformNodes;
 import org.elasticsearch.xpack.transform.utils.SourceDestValidations;
@@ -99,8 +99,10 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li
 
         TransformNodes.warnIfNoTransformNodes(clusterState);
 
-        final TransformConfig config = request.getConfig();
-        final Function function = FunctionFactory.create(config);
+        var config = request.getConfig();
+        var function = FunctionFactory.create(config);
+        var parentTaskId = new TaskId(clusterService.localNode().getId(), task.getId());
+        var parentClient = new ParentTaskAssigningClient(client, parentTaskId);
 
         if (config.getVersion() == null || config.getVersion().before(TransformDeprecations.MIN_TRANSFORM_VERSION)) {
             listener.onFailure(
@@ -130,7 +132,7 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li
             if (request.isDeferValidation()) {
                 deduceMappingsListener.onResponse(emptyMap());
             } else {
-                function.deduceMappings(client, config.getHeaders(), config.getId(), config.getSource(), deduceMappingsListener);
+                function.deduceMappings(parentClient, config.getHeaders(), config.getId(), config.getSource(), deduceMappingsListener);
             }
         }, listener::onFailure);
 
@@ -139,7 +141,7 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li
             if (request.isDeferValidation()) {
                 validateQueryListener.onResponse(true);
             } else {
-                function.validateQuery(client, config.getHeaders(), config.getSource(), request.ackTimeout(), validateQueryListener);
+                function.validateQuery(parentClient, config.getHeaders(), config.getSource(), request.ackTimeout(), validateQueryListener);
             }
         }, listener::onFailure);
 

diff --git a/...m/src/main/java/org/elasticsearch/xpack/transform/rest/action/RestPutTransformAction.java b/...m/src/main/java/org/elasticsearch/xpack/transform/rest/action/RestPutTransformAction.java
@@ -15,6 +15,7 @@
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
+import org.elasticsearch.rest.action.RestCancellableNodeClient;
 import org.elasticsearch.rest.action.RestToXContentListener;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
@@ -66,6 +67,10 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient
 
         PutTransformAction.Request request = PutTransformAction.Request.fromXContent(parser, id, deferValidation, timeout);
 
-        return channel -> client.execute(PutTransformAction.INSTANCE, request, new RestToXContentListener<>(channel));
+        return channel -> new RestCancellableNodeClient(client, restRequest.getHttpChannel()).execute(
+            PutTransformAction.INSTANCE,
+            request,
+            new RestToXContentListener<>(channel)
+        );
     }
 }
diff --git a/...src/main/java/org/elasticsearch/xpack/transform/rest/action/RestStartTransformAction.java b/...src/main/java/org/elasticsearch/xpack/transform/rest/action/RestStartTransformAction.java
@@ -17,6 +17,7 @@
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
+import org.elasticsearch.rest.action.RestCancellableNodeClient;
 import org.elasticsearch.rest.action.RestToXContentListener;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xpack.core.transform.TransformField;
@@ -45,7 +46,11 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient
         TimeValue timeout = restRequest.paramAsTime(TransformField.TIMEOUT.getPreferredName(), AcknowledgedRequest.DEFAULT_ACK_TIMEOUT);
 
         StartTransformAction.Request request = new StartTransformAction.Request(id, from, timeout);
-        return channel -> client.execute(StartTransformAction.INSTANCE, request, new RestToXContentListener<>(channel));
+        return channel -> new RestCancellableNodeClient(client, restRequest.getHttpChannel()).execute(
+            StartTransformAction.INSTANCE,
+            request,
+            new RestToXContentListener<>(channel)
+        );
     }
 
     private static Instant parseDateOrThrow(String date, ParseField paramName, LongSupplier now) {