[ML] make more GET and heavier processing APIs cancellable (elastic#88142)

This commit makes the following APIs cancellable:

get datafeeds
get jobs
evaluate data frames
explain data frames
get job model snapshot
preview datafeed
preview data frame
get categories (this can be expensive due to grok parsing)

Author: benwtrent

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/EvaluateDataFrameAction.java

@@ -16,6 +16,9 @@
 import org.elasticsearch.common.xcontent.XContentParserUtils;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
@@ -30,6 +33,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 
@@ -175,6 +179,11 @@ public boolean equals(Object o) {
                 && Objects.equals(queryProvider, that.queryProvider)
                 && Objects.equals(evaluation, that.evaluation);
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, "evaluate_data_frame", parentTaskId, headers);
+        }
     }
 
     public static class Response extends ActionResponse implements ToXContentObject {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/ExplainDataFrameAnalyticsAction.java

@@ -6,21 +6,36 @@
  */
 package org.elasticsearch.xpack.core.ml.action;
 
+import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.action.ActionType;
+import org.elasticsearch.action.ValidateActions;
+import org.elasticsearch.action.support.master.AcknowledgedRequest;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xpack.core.common.validation.SourceDestValidator;
+import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
+import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsSource;
 import org.elasticsearch.xpack.core.ml.dataframe.explain.FieldSelection;
 import org.elasticsearch.xpack.core.ml.dataframe.explain.MemoryEstimation;
+import org.elasticsearch.xpack.core.ml.job.messages.Messages;
+import org.elasticsearch.xpack.core.ml.utils.MlStrings;
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.core.Strings.format;
+
 public class ExplainDataFrameAnalyticsAction extends ActionType<ExplainDataFrameAnalyticsAction.Response> {
 
     public static final ExplainDataFrameAnalyticsAction INSTANCE = new ExplainDataFrameAnalyticsAction();
@@ -30,6 +45,118 @@ private ExplainDataFrameAnalyticsAction() {
         super(NAME, ExplainDataFrameAnalyticsAction.Response::new);
     }
 
+    public static class Request extends AcknowledgedRequest<Request> implements ToXContentObject {
+        public static Request parseRequest(XContentParser parser) {
+            DataFrameAnalyticsConfig.Builder configBuilder = DataFrameAnalyticsConfig.STRICT_PARSER.apply(parser, null);
+            DataFrameAnalyticsConfig config = configBuilder.buildForExplain();
+            return new Request(config);
+        }
+
+        private final DataFrameAnalyticsConfig config;
+
+        public Request(StreamInput in) throws IOException {
+            super(in);
+            config = new DataFrameAnalyticsConfig(in);
+        }
+
+        public Request(DataFrameAnalyticsConfig config) {
+            this.config = config;
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            super.writeTo(out);
+            config.writeTo(out);
+        }
+
+        public DataFrameAnalyticsConfig getConfig() {
+            return config;
+        }
+
+        @Override
+        public ActionRequestValidationException validate() {
+            ActionRequestValidationException error = null;
+            error = checkConfigIdIsValid(config, error);
+            error = SourceDestValidator.validateRequest(error, config.getDest().getIndex());
+            error = checkNoIncludedAnalyzedFieldsAreExcludedBySourceFiltering(config, error);
+            return error;
+        }
+
+        private ActionRequestValidationException checkConfigIdIsValid(
+            DataFrameAnalyticsConfig analyticsConfig,
+            ActionRequestValidationException error
+        ) {
+            if (MlStrings.isValidId(analyticsConfig.getId()) == false) {
+                error = ValidateActions.addValidationError(
+                    Messages.getMessage(Messages.INVALID_ID, DataFrameAnalyticsConfig.ID, analyticsConfig.getId()),
+                    error
+                );
+            }
+            if (MlStrings.hasValidLengthForId(analyticsConfig.getId()) == false) {
+                error = ValidateActions.addValidationError(
+                    Messages.getMessage(
+                        Messages.ID_TOO_LONG,
+                        DataFrameAnalyticsConfig.ID,
+                        analyticsConfig.getId(),
+                        MlStrings.ID_LENGTH_LIMIT
+                    ),
+                    error
+                );
+            }
+            return error;
+        }
+
+        private ActionRequestValidationException checkNoIncludedAnalyzedFieldsAreExcludedBySourceFiltering(
+            DataFrameAnalyticsConfig analyticsConfig,
+            ActionRequestValidationException error
+        ) {
+            if (analyticsConfig.getAnalyzedFields() == null) {
+                return error;
+            }
+            for (String analyzedInclude : analyticsConfig.getAnalyzedFields().includes()) {
+                if (analyticsConfig.getSource().isFieldExcluded(analyzedInclude)) {
+                    return ValidateActions.addValidationError(
+                        "field ["
+                            + analyzedInclude
+                            + "] is included in ["
+                            + DataFrameAnalyticsConfig.ANALYZED_FIELDS.getPreferredName()
+                            + "] but not in ["
+                            + DataFrameAnalyticsConfig.SOURCE.getPreferredName()
+                            + "."
+                            + DataFrameAnalyticsSource._SOURCE.getPreferredName()
+                            + "]",
+                        error
+                    );
+                }
+            }
+            return error;
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            config.toXContent(builder, params);
+            return builder;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            Request request = (Request) o;
+            return Objects.equals(config, request.config);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(config);
+        }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, format("explain_data_frame_analytics[%s]", config.getId()), parentTaskId, headers);
+        }
+    }
+
     public static class Response extends ActionResponse implements ToXContentObject {
 
         public static final ParseField TYPE = new ParseField("explain_data_frame_analytics_response");

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetCategoriesAction.java

@@ -11,6 +11,9 @@
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
@@ -24,9 +27,11 @@
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.Objects;
 
 import static org.elasticsearch.action.ValidateActions.addValidationError;
+import static org.elasticsearch.core.Strings.format;
 
 public class GetCategoriesAction extends ActionType<GetCategoriesAction.Response> {
 
@@ -190,6 +195,11 @@ public boolean equals(Object o) {
         public int hashCode() {
             return Objects.hash(jobId, categoryId, pageParams, partitionFieldValue);
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, format("get_categories[%s:%s]", jobId, categoryId), parentTaskId, headers);
+        }
     }
 
     public static class Response extends AbstractGetResourcesResponse<CategoryDefinition> implements ToXContentObject {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedsAction.java

@@ -11,15 +11,21 @@
 import org.elasticsearch.action.support.master.MasterNodeReadRequest;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xpack.core.action.AbstractGetResourcesResponse;
 import org.elasticsearch.xpack.core.action.util.QueryPage;
 import org.elasticsearch.xpack.core.ml.datafeed.DatafeedConfig;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.core.Strings.format;
+
 public class GetDatafeedsAction extends ActionType<GetDatafeedsAction.Response> {
 
     public static final GetDatafeedsAction INSTANCE = new GetDatafeedsAction();
@@ -93,6 +99,11 @@ public boolean equals(Object obj) {
             Request other = (Request) obj;
             return Objects.equals(datafeedId, other.datafeedId) && Objects.equals(allowNoMatch, other.allowNoMatch);
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, format("get_datafeeds[%s]", datafeedId), parentTaskId, headers);
+        }
     }
 
     public static class Response extends AbstractGetResourcesResponse<DatafeedConfig> implements ToXContentObject {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetJobsAction.java

@@ -11,15 +11,21 @@
 import org.elasticsearch.action.support.master.MasterNodeReadRequest;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xpack.core.action.AbstractGetResourcesResponse;
 import org.elasticsearch.xpack.core.action.util.QueryPage;
 import org.elasticsearch.xpack.core.ml.job.config.Job;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.core.Strings.format;
+
 public class GetJobsAction extends ActionType<GetJobsAction.Response> {
 
     public static final GetJobsAction INSTANCE = new GetJobsAction();
@@ -91,6 +97,11 @@ public boolean equals(Object obj) {
             Request other = (Request) obj;
             return Objects.equals(jobId, other.jobId) && Objects.equals(allowNoMatch, other.allowNoMatch);
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, format("get_anomaly_detection_jobs[%s]", jobId), parentTaskId, headers);
+        }
     }
 
     public static class Response extends AbstractGetResourcesResponse<Job> implements ToXContentObject {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetModelSnapshotsAction.java

@@ -12,6 +12,9 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
@@ -25,8 +28,11 @@
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.core.Strings.format;
+
 public class GetModelSnapshotsAction extends ActionType<GetModelSnapshotsAction.Response> {
 
     public static final GetModelSnapshotsAction INSTANCE = new GetModelSnapshotsAction();
@@ -205,6 +211,11 @@ public boolean equals(Object obj) {
                 && Objects.equals(sort, other.sort)
                 && Objects.equals(desc, other.desc);
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, format("get_job_model_snapshot[%s:%s]", jobId, snapshotId), parentTaskId, headers);
+        }
     }
 
     public static class Response extends AbstractGetResourcesResponse<ModelSnapshot> implements ToXContentObject {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/PreviewDataFrameAnalyticsAction.java

@@ -12,6 +12,9 @@
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
@@ -26,6 +29,8 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.core.Strings.format;
+
 public class PreviewDataFrameAnalyticsAction extends ActionType<PreviewDataFrameAnalyticsAction.Response> {
 
     public static final PreviewDataFrameAnalyticsAction INSTANCE = new PreviewDataFrameAnalyticsAction();
@@ -87,6 +92,11 @@ public int hashCode() {
             return Objects.hash(config);
         }
 
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, format("preview_data_frame_analytics[%s]", config.getId()), parentTaskId, headers);
+        }
+
         public static class Builder {
             private DataFrameAnalyticsConfig config;
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/PreviewDatafeedAction.java

@@ -17,6 +17,9 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
@@ -29,9 +32,11 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Map;
 import java.util.Objects;
 import java.util.OptionalLong;
 
+import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.core.ml.action.StartDatafeedAction.DatafeedParams.parseDateOrThrow;
 import static org.elasticsearch.xpack.core.ml.action.StartDatafeedAction.END_TIME;
 import static org.elasticsearch.xpack.core.ml.action.StartDatafeedAction.START_TIME;
@@ -199,6 +204,11 @@ public boolean equals(Object obj) {
                 && Objects.equals(jobConfig, other.jobConfig);
         }
 
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, format("preview_datafeed[%s]", datafeedId), parentTaskId, headers);
+        }
+
         public static class Builder {
             private String datafeedId;
             private DatafeedConfig.Builder datafeedBuilder;