EQL: Make EQL search task cancellable (#54598)

First step towards async search execution. At the moment we don't try to cancel
the underlying search requests, and just check if the task is canceled before
performing network operation (such as field caps and search)

Relates to #49638

Author: imotov

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/action/EqlSearchRequest.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.action.IndicesRequest;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ObjectParser;
@@ -19,9 +20,12 @@
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.searchafter.SearchAfterBuilder;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.Objects;
 import java.util.function.Supplier;
 
@@ -287,4 +291,16 @@ public EqlSearchRequest indicesOptions(IndicesOptions indicesOptions) {
     public IndicesOptions indicesOptions() {
         return indicesOptions;
     }
+
+    @Override
+    public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+        return new EqlSearchTask(id, type, action, () -> {
+            StringBuilder sb = new StringBuilder();
+            sb.append("indices[");
+            Strings.arrayToDelimitedString(indices, ",", sb);
+            sb.append("], ");
+            sb.append(query);
+            return sb.toString();
+        }, parentTaskId, headers);
+    }
 }

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/action/EqlSearchTask.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.eql.action;
+
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.TaskId;
+
+import java.util.Map;
+import java.util.function.Supplier;
+
+public class EqlSearchTask extends CancellableTask {
+    private final Supplier<String> descriptionSupplier;
+
+    public EqlSearchTask(long id, String type, String action, Supplier<String> descriptionSupplier, TaskId parentTaskId,
+                         Map<String, String> headers) {
+        super(id, type, action, null, parentTaskId, headers);
+        this.descriptionSupplier = descriptionSupplier;
+    }
+
+    @Override
+    public boolean shouldCancelChildrenOnCancellation() {
+        return false;
+    }
+
+    @Override
+    public String getDescription() {
+        return descriptionSupplier.get();
+    }
+}

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/search/Querier.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.aggregations.Aggregation;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.xpack.eql.querydsl.container.QueryContainer;
 import org.elasticsearch.xpack.eql.session.Configuration;
 import org.elasticsearch.xpack.eql.session.EqlSession;
@@ -56,7 +57,9 @@ public void query(List<Attribute> output, QueryContainer container, String index
         if (log.isTraceEnabled()) {
             log.trace("About to execute query {} on {}", StringUtils.toString(sourceBuilder), index);
         }
-        
+        if (cfg.isCancelled()) {
+            throw new TaskCancelledException("cancelled");
+        }
         SearchRequest search = prepareRequest(client, sourceBuilder, cfg.requestTimeout(), false,
                 Strings.commaDelimitedListToStringArray(index));
 
@@ -93,4 +96,4 @@ protected static void logSearchResponse(SearchResponse response, Logger logger)
                 response.getHits().getTotalHits().value, aggs.size(), aggsNames, response.getFailedShards(), response.getSkippedShards(),
                 response.getSuccessfulShards(), response.getTotalShards(), response.getTook(), response.isTimedOut());
     }
-}
+}

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/plugin/TransportEqlSearchAction.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.xpack.eql.action.EqlSearchAction;
 import org.elasticsearch.xpack.eql.action.EqlSearchRequest;
 import org.elasticsearch.xpack.eql.action.EqlSearchResponse;
+import org.elasticsearch.xpack.eql.action.EqlSearchTask;
 import org.elasticsearch.xpack.eql.execution.PlanExecutor;
 import org.elasticsearch.xpack.eql.parser.ParserParams;
 import org.elasticsearch.xpack.eql.session.Configuration;
@@ -49,10 +50,10 @@ public TransportEqlSearchAction(Settings settings, ClusterService clusterService
 
     @Override
     protected void doExecute(Task task, EqlSearchRequest request, ActionListener<EqlSearchResponse> listener) {
-        operation(planExecutor, request, username(securityContext), clusterName(clusterService), listener);
+        operation(planExecutor, (EqlSearchTask) task, request, username(securityContext), clusterName(clusterService), listener);
     }
 
-    public static void operation(PlanExecutor planExecutor, EqlSearchRequest request, String username,
+    public static void operation(PlanExecutor planExecutor, EqlSearchTask task, EqlSearchRequest request, String username,
                                  String clusterName, ActionListener<EqlSearchResponse> listener) {
         // TODO: these should be sent by the client
         ZoneId zoneId = DateUtils.of("Z");
@@ -67,7 +68,7 @@ public static void operation(PlanExecutor planExecutor, EqlSearchRequest request
             .implicitJoinKey(request.implicitJoinKeyField());
 
         Configuration cfg = new Configuration(request.indices(), zoneId, username, clusterName, filter, timeout, request.fetchSize(),
-                includeFrozen, clientId);
+                includeFrozen, clientId, task);
         planExecutor.eql(cfg, request.query(), params, wrap(r -> listener.onResponse(createResponse(r)), listener::onFailure));
     }
 

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/session/Configuration.java

@@ -9,22 +9,24 @@
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.xpack.eql.action.EqlSearchTask;
 
 import java.time.ZoneId;
 
 public class Configuration extends org.elasticsearch.xpack.ql.session.Configuration {
-    
+
     private final String[] indices;
     private final TimeValue requestTimeout;
     private final int size;
     private final String clientId;
     private final boolean includeFrozenIndices;
+    private final EqlSearchTask task;
 
     @Nullable
     private QueryBuilder filter;
 
     public Configuration(String[] indices, ZoneId zi, String username, String clusterName, QueryBuilder filter, TimeValue requestTimeout,
-                         int size, boolean includeFrozen, String clientId) {
+                         int size, boolean includeFrozen, String clientId, EqlSearchTask task) {
 
         super(zi, username, clusterName);
 
@@ -34,6 +36,7 @@ public Configuration(String[] indices, ZoneId zi, String username, String cluste
         this.size = size;
         this.clientId = clientId;
         this.includeFrozenIndices = includeFrozen;
+        this.task = task;
     }
 
     public String[] indices() {
@@ -59,4 +62,8 @@ public String clientId() {
     public boolean includeFrozen() {
         return includeFrozenIndices;
     }
-}
+
+    public boolean isCancelled() {
+        return task.isCancelled();
+    }
+}

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/session/EqlSession.java

@@ -9,6 +9,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.xpack.eql.analysis.Analyzer;
 import org.elasticsearch.xpack.eql.analysis.PreAnalyzer;
 import org.elasticsearch.xpack.eql.execution.PlanExecutor;
@@ -35,7 +36,7 @@ public class EqlSession {
 
     public EqlSession(Client client, Configuration cfg, IndexResolver indexResolver, PreAnalyzer preAnalyzer, Analyzer analyzer,
             Optimizer optimizer, Planner planner, PlanExecutor planExecutor) {
-        
+
         this.client = client;
         this.configuration = cfg;
         this.indexResolver = indexResolver;
@@ -60,7 +61,7 @@ public Configuration configuration() {
     public void eql(String eql, ParserParams params, ActionListener<Results> listener) {
         eqlExecutable(eql, params, wrap(e -> e.execute(this, listener), listener::onFailure));
     }
-    
+
     public void eqlExecutable(String eql, ParserParams params, ActionListener<PhysicalPlan> listener) {
         try {
             physicalPlan(doParse(eql, params), listener);
@@ -88,7 +89,9 @@ public void analyzedPlan(LogicalPlan parsed, ActionListener<LogicalPlan> listene
 
     private <T> void preAnalyze(LogicalPlan parsed, ActionListener<LogicalPlan> listener) {
         String indexWildcard = Strings.arrayToCommaDelimitedString(configuration.indices());
-
+        if(configuration.isCancelled()){
+            throw new TaskCancelledException("cancelled");
+        }
         indexResolver.resolveAsMergedMapping(indexWildcard, null, configuration.includeFrozen(), wrap(r -> {
             listener.onResponse(preAnalyzer.preAnalyze(parsed, r));
         }, listener::onFailure));
@@ -97,4 +100,4 @@ private <T> void preAnalyze(LogicalPlan parsed, ActionListener<LogicalPlan> list
     private LogicalPlan doParse(String eql, ParserParams params) {
         return new EqlParser().createStatement(eql, params);
     }
-}
+}

x-pack/plugin/eql/src/test/java/org/elasticsearch/xpack/eql/EqlTestUtils.java

@@ -7,11 +7,16 @@
 package org.elasticsearch.xpack.eql;
 
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.xpack.eql.action.EqlSearchAction;
+import org.elasticsearch.xpack.eql.action.EqlSearchTask;
 import org.elasticsearch.xpack.eql.session.Configuration;
 
+import java.util.Collections;
+
 import static org.elasticsearch.test.ESTestCase.randomAlphaOfLength;
 import static org.elasticsearch.test.ESTestCase.randomBoolean;
 import static org.elasticsearch.test.ESTestCase.randomIntBetween;
+import static org.elasticsearch.test.ESTestCase.randomLong;
 import static org.elasticsearch.test.ESTestCase.randomNonNegativeLong;
 import static org.elasticsearch.test.ESTestCase.randomZone;
 
@@ -21,7 +26,8 @@ private EqlTestUtils() {
     }
 
     public static final Configuration TEST_CFG = new Configuration(new String[]{"none"}, org.elasticsearch.xpack.ql.util.DateUtils.UTC,
-            "nobody", "cluster", null, TimeValue.timeValueSeconds(30), -1, false, "");
+            "nobody", "cluster", null, TimeValue.timeValueSeconds(30), -1, false, "",
+            new EqlSearchTask(-1, "", EqlSearchAction.NAME, () -> "", null, Collections.emptyMap()));
 
     public static Configuration randomConfiguration() {
         return new Configuration(new String[]{randomAlphaOfLength(16)},
@@ -32,6 +38,11 @@ public static Configuration randomConfiguration() {
             new TimeValue(randomNonNegativeLong()),
             randomIntBetween(5, 100),
             randomBoolean(),
-            randomAlphaOfLength(16));
+            randomAlphaOfLength(16),
+            randomTask());
+    }
+
+    public static EqlSearchTask randomTask() {
+        return new EqlSearchTask(randomLong(), "transport", EqlSearchAction.NAME, () -> "", null, Collections.emptyMap());
     }
 }

x-pack/plugin/eql/src/test/java/org/elasticsearch/xpack/eql/analysis/CancellationTests.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.eql.analysis;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.fieldcaps.FieldCapabilities;
+import org.elasticsearch.action.fieldcaps.FieldCapabilitiesResponse;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.tasks.TaskCancelledException;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.eql.action.EqlSearchRequest;
+import org.elasticsearch.xpack.eql.action.EqlSearchResponse;
+import org.elasticsearch.xpack.eql.action.EqlSearchTask;
+import org.elasticsearch.xpack.eql.execution.PlanExecutor;
+import org.elasticsearch.xpack.eql.plugin.TransportEqlSearchAction;
+import org.elasticsearch.xpack.ql.index.IndexResolver;
+import org.elasticsearch.xpack.ql.type.DefaultDataTypeRegistry;
+import org.mockito.stubbing.Answer;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonMap;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+public class CancellationTests extends ESTestCase {
+
+    public void testCancellationBeforeFieldCaps() throws InterruptedException {
+        Client client = mock(Client.class);
+        EqlSearchTask task = mock(EqlSearchTask.class);
+        when(task.isCancelled()).thenReturn(true);
+
+        IndexResolver indexResolver = new IndexResolver(client, randomAlphaOfLength(10), DefaultDataTypeRegistry.INSTANCE);
+        PlanExecutor planExecutor = new PlanExecutor(client, indexResolver, new NamedWriteableRegistry(Collections.emptyList()));
+        CountDownLatch countDownLatch = new CountDownLatch(1);
+        TransportEqlSearchAction.operation(planExecutor, task, new EqlSearchRequest().query("foo where blah"), "", "",
+            new ActionListener<>() {
+                @Override
+                public void onResponse(EqlSearchResponse eqlSearchResponse) {
+                    fail("Shouldn't be here");
+                    countDownLatch.countDown();
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    assertThat(e, instanceOf(TaskCancelledException.class));
+                    countDownLatch.countDown();
+                }
+            });
+        countDownLatch.await();
+        verify(task, times(1)).isCancelled();
+        verifyNoMoreInteractions(client, task);
+    }
+
+    public void testCancellationBeforeSearch() throws InterruptedException {
+        Client client = mock(Client.class);
+
+        AtomicBoolean cancelled = new AtomicBoolean(false);
+        EqlSearchTask task = mock(EqlSearchTask.class);
+        when(task.isCancelled()).then(invocationOnMock -> cancelled.get());
+
+        String[] indices = new String[]{"endgame"};
+
+        FieldCapabilities fooField =
+            new FieldCapabilities("foo", "integer", true, true, indices, null, null, emptyMap());
+        FieldCapabilities categoryField =
+            new FieldCapabilities("event.category", "keyword", true, true, indices, null, null, emptyMap());
+        FieldCapabilities timestampField =
+            new FieldCapabilities("@timestamp", "date", true, true, indices, null, null, emptyMap());
+        Map<String, Map<String, FieldCapabilities>> fields = new HashMap<>();
+        fields.put(fooField.getName(), singletonMap(fooField.getName(), fooField));
+        fields.put(categoryField.getName(), singletonMap(categoryField.getName(), categoryField));
+        fields.put(timestampField.getName(), singletonMap(timestampField.getName(), timestampField));
+
+        FieldCapabilitiesResponse fieldCapabilitiesResponse = mock(FieldCapabilitiesResponse.class);
+        when(fieldCapabilitiesResponse.getIndices()).thenReturn(indices);
+        when(fieldCapabilitiesResponse.get()).thenReturn(fields);
+        doAnswer((Answer<Void>) invocation -> {
+            @SuppressWarnings("unchecked")
+            ActionListener<FieldCapabilitiesResponse> listener = (ActionListener<FieldCapabilitiesResponse>) invocation.getArguments()[1];
+            assertFalse(cancelled.getAndSet(true));
+            listener.onResponse(fieldCapabilitiesResponse);
+            return null;
+        }).when(client).fieldCaps(any(), any());
+
+
+        IndexResolver indexResolver = new IndexResolver(client, randomAlphaOfLength(10), DefaultDataTypeRegistry.INSTANCE);
+        PlanExecutor planExecutor = new PlanExecutor(client, indexResolver, new NamedWriteableRegistry(Collections.emptyList()));
+        CountDownLatch countDownLatch = new CountDownLatch(1);
+        TransportEqlSearchAction.operation(planExecutor, task, new EqlSearchRequest().indices("endgame")
+            .query("process where foo==3"), "", "", new ActionListener<>() {
+            @Override
+            public void onResponse(EqlSearchResponse eqlSearchResponse) {
+                fail("Shouldn't be here");
+                countDownLatch.countDown();
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                assertThat(e, instanceOf(TaskCancelledException.class));
+                countDownLatch.countDown();
+            }
+        });
+        countDownLatch.await();
+        verify(client).fieldCaps(any(), any());
+        verify(task, times(2)).isCancelled();
+        verifyNoMoreInteractions(client, task);
+    }
+
+}