[ML] Fix forecasts deletion action when there are many forecast documents. (#75127)

Author: przemekwitek

server/src/main/java/org/elasticsearch/index/reindex/DeleteByQueryRequest.java

@@ -151,7 +151,7 @@ public String toString() {
     //delete by query deletes all documents that match a query. The indices and indices options that affect how
     //indices are resolved depend entirely on the inner search request. That's why the following methods delegate to it.
     @Override
-    public IndicesRequest indices(String... indices) {
+    public DeleteByQueryRequest indices(String... indices) {
         assert getSearchRequest() != null;
         getSearchRequest().indices(indices);
         return this;

x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/ForecastIT.java

@@ -26,16 +26,20 @@
 
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 import static org.elasticsearch.xpack.core.ml.job.messages.Messages.JOB_FORECAST_NATIVE_PROCESS_KILLED;
 import static org.hamcrest.Matchers.closeTo;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
 
 public class ForecastIT extends MlNativeAutodetectIntegTestCase {
 
@@ -343,7 +347,61 @@ public void testDeleteWildCard() throws Exception {
             assertNull(getForecastStats(job.getId(), forecastId2Duration1HourNoExpiry));
             assertNull(getForecastStats(job.getId(), forecastId2Duration1HourNoExpiry2));
         }
+    }
+
+    public void testDeleteAll() throws Exception {
+        Detector.Builder detector1 = new Detector.Builder("sum", "value").setPartitionFieldName("category");
+        Detector.Builder detector2 = new Detector.Builder("mean", "value").setPartitionFieldName("category");
+
+        TimeValue bucketSpan = TimeValue.timeValueHours(1);
+        AnalysisConfig.Builder analysisConfig = new AnalysisConfig.Builder(Arrays.asList(detector1.build(), detector2.build()));
+        analysisConfig.setBucketSpan(bucketSpan);
+        DataDescription.Builder dataDescription = new DataDescription.Builder();
+        dataDescription.setTimeFormat("epoch");
+
+        Job.Builder job = new Job.Builder("forecast-it-test-delete-wildcard-2");
+        job.setAnalysisConfig(analysisConfig);
+        job.setDataDescription(dataDescription);
+        putJob(job);
+        openJob(job.getId());
+
+        long now = Instant.now().getEpochSecond();
+        long timestamp = now - 50 * bucketSpan.seconds();
+        List<String> data = new ArrayList<>();
+        String[] partitionFieldValues = IntStream.range(0, 20).mapToObj(i -> "category_" + i).toArray(String[]::new);
+        while (timestamp < now) {
+            for (String partitionFieldValue : partitionFieldValues) {
+                data.add(createJsonRecord(createRecord(timestamp, partitionFieldValue, 10.0)));
+                data.add(createJsonRecord(createRecord(timestamp, partitionFieldValue, 30.0)));
+            }
+            timestamp += bucketSpan.seconds();
+        }
 
+        postData(job.getId(), data.stream().collect(Collectors.joining()));
+        flushJob(job.getId(), false);
+
+        long noForecasts = 11;  // We want to make sure we set the search size instead of relying on the default
+        List<String> forecastIds = new ArrayList<>();
+        for (int i = 0; i < noForecasts; ++i) {
+            String forecastId = forecast(job.getId(), TimeValue.timeValueHours(100), TimeValue.ZERO);
+            forecastIds.add(forecastId);
+            waitForecastToFinish(job.getId(), forecastId);
+        }
+        closeJob(job.getId());
+
+        assertThat(getJobStats(job.getId()).get(0).getForecastStats().getTotal(), is(equalTo(noForecasts)));
+        for (String forecastId : forecastIds) {
+            assertNotNull(getForecastStats(job.getId(), forecastId));
+        }
+
+        DeleteForecastAction.Request request = new DeleteForecastAction.Request(job.getId(), randomBoolean() ? "*" : "_all");
+        AcknowledgedResponse response = client().execute(DeleteForecastAction.INSTANCE, request).actionGet();
+        assertTrue(response.isAcknowledged());
+
+        assertThat(getJobStats(job.getId()).get(0).getForecastStats().getTotal(), is(equalTo(0L)));
+        for (String forecastId : forecastIds) {
+            assertNull(getForecastStats(job.getId(), forecastId));
+        }
     }
 
     public void testDelete() throws Exception {
@@ -561,4 +619,10 @@ private static Map<String, Object> createRecord(long timestamp, double value) {
         record.put("value", value);
         return record;
     }
+
+    private static Map<String, Object> createRecord(long timestamp, String partitionFieldValue, double value) {
+        Map<String, Object> record = createRecord(timestamp, value);
+        record.put("category", partitionFieldValue);
+        return record;
+    }
 }

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

@@ -23,13 +23,8 @@
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.common.inject.Inject;
-import org.elasticsearch.common.xcontent.DeprecationHandler;
-import org.elasticsearch.common.xcontent.NamedXContentRegistry;
-import org.elasticsearch.common.xcontent.XContentFactory;
-import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.query.BoolQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
@@ -57,17 +52,13 @@
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 import org.elasticsearch.xpack.ml.utils.QueryBuilderHelper;
 
-import java.io.IOException;
-import java.io.InputStream;
 import java.util.ArrayList;
-import java.util.Collection;
-import java.util.EnumSet;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.TimeoutException;
-import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
+import static java.util.stream.Collectors.toSet;
 import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
 import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin;
 
@@ -80,8 +71,10 @@ public class TransportDeleteForecastAction extends HandledTransportAction<Delete
     private final ClusterService clusterService;
     private static final int MAX_FORECAST_TO_SEARCH = 10_000;
 
-    private static final Set<ForecastRequestStatus> DELETABLE_STATUSES =
-        EnumSet.of(ForecastRequestStatus.FINISHED, ForecastRequestStatus.FAILED);
+    private static final Set<String> DELETABLE_STATUSES =
+        Stream.of(ForecastRequestStatus.FINISHED, ForecastRequestStatus.FAILED)
+            .map(ForecastRequestStatus::toString)
+            .collect(toSet());
 
     @Inject
     public TransportDeleteForecastAction(TransportService transportService,
@@ -105,47 +98,54 @@ protected void doExecute(Task task, DeleteForecastAction.Request request, Action
             e -> handleFailure(e, request, listener)
         );
 
-        SearchSourceBuilder source = new SearchSourceBuilder();
-
-        BoolQueryBuilder builder = QueryBuilders.boolQuery()
-            .filter(QueryBuilders.termQuery(Result.RESULT_TYPE.getPreferredName(), ForecastRequestStats.RESULT_TYPE_VALUE));
+        BoolQueryBuilder query =
+            QueryBuilders.boolQuery()
+                .filter(QueryBuilders.termQuery(Result.RESULT_TYPE.getPreferredName(), ForecastRequestStats.RESULT_TYPE_VALUE));
         QueryBuilderHelper
             .buildTokenFilterQuery(Forecast.FORECAST_ID.getPreferredName(), forecastIds)
-            .ifPresent(builder::filter);
-        source.query(builder);
-
-        SearchRequest searchRequest = new SearchRequest(AnomalyDetectorsIndex.jobResultsAliasedName(jobId));
-        searchRequest.source(source);
+            .ifPresent(query::filter);
+        SearchSourceBuilder source =
+            new SearchSourceBuilder()
+                .size(MAX_FORECAST_TO_SEARCH)
+                // We only need forecast id and status, there is no need fetching the whole source
+                .fetchSource(false)
+                .docValueField(ForecastRequestStats.FORECAST_ID.getPreferredName())
+                .docValueField(ForecastRequestStats.STATUS.getPreferredName())
+                .query(query);
+        SearchRequest searchRequest =
+            new SearchRequest(AnomalyDetectorsIndex.jobResultsAliasedName(jobId))
+                .source(source);
 
         executeAsyncWithOrigin(client, ML_ORIGIN, SearchAction.INSTANCE, searchRequest, forecastStatsHandler);
     }
 
-    static void validateForecastState(Collection<ForecastRequestStats> forecastsToDelete, JobState jobState, String jobId) {
-        List<String> badStatusForecasts = forecastsToDelete.stream()
-            .filter((f) -> DELETABLE_STATUSES.contains(f.getStatus()) == false)
-            .map(ForecastRequestStats::getForecastId)
-            .collect(Collectors.toList());
-        if (badStatusForecasts.size() > 0 && JobState.OPENED.equals(jobState)) {
+    static List<String> extractForecastIds(SearchHit[] forecastsToDelete, JobState jobState, String jobId) {
+        List<String> forecastIds = new ArrayList<>(forecastsToDelete.length);
+        List<String> badStatusForecastIds = new ArrayList<>();
+        for (SearchHit hit : forecastsToDelete) {
+            String forecastId = hit.field(ForecastRequestStats.FORECAST_ID.getPreferredName()).getValue();
+            String forecastStatus = hit.field(ForecastRequestStats.STATUS.getPreferredName()).getValue();
+            if (DELETABLE_STATUSES.contains(forecastStatus)) {
+                forecastIds.add(forecastId);
+            } else {
+                badStatusForecastIds.add(forecastId);
+            }
+        }
+        if (badStatusForecastIds.size() > 0 && JobState.OPENED.equals(jobState)) {
             throw ExceptionsHelper.conflictStatusException(
-                Messages.getMessage(Messages.REST_CANNOT_DELETE_FORECAST_IN_CURRENT_STATE, badStatusForecasts, jobId));
+                Messages.getMessage(Messages.REST_CANNOT_DELETE_FORECAST_IN_CURRENT_STATE, badStatusForecastIds, jobId));
         }
+        return forecastIds;
     }
 
     private void deleteForecasts(SearchResponse searchResponse,
                                  DeleteForecastAction.Request request,
                                  ActionListener<AcknowledgedResponse> listener) {
         final String jobId = request.getJobId();
-        Set<ForecastRequestStats> forecastsToDelete;
-        try {
-            forecastsToDelete = parseForecastsFromSearch(searchResponse);
-        } catch (IOException e) {
-            listener.onFailure(e);
-            return;
-        }
+        SearchHits forecastsToDelete = searchResponse.getHits();
 
-        if (forecastsToDelete.isEmpty()) {
-            if (Strings.isAllOrWildcard(request.getForecastId()) &&
-                request.isAllowNoForecasts()) {
+        if (forecastsToDelete.getHits().length == 0) {
+            if (Strings.isAllOrWildcard(request.getForecastId()) && request.isAllowNoForecasts()) {
                 listener.onResponse(AcknowledgedResponse.TRUE);
             } else {
                 listener.onFailure(
@@ -156,16 +156,15 @@ private void deleteForecasts(SearchResponse searchResponse,
         final ClusterState state = clusterService.state();
         PersistentTasksCustomMetadata persistentTasks = state.metadata().custom(PersistentTasksCustomMetadata.TYPE);
         JobState jobState = MlTasks.getJobState(jobId, persistentTasks);
+        final List<String> forecastIds;
         try {
-            validateForecastState(forecastsToDelete, jobState, jobId);
+            forecastIds = extractForecastIds(forecastsToDelete.getHits(), jobState, jobId);
         } catch (ElasticsearchException ex) {
             listener.onFailure(ex);
             return;
         }
 
-        final List<String> forecastIds = forecastsToDelete.stream().map(ForecastRequestStats::getForecastId).collect(Collectors.toList());
         DeleteByQueryRequest deleteByQueryRequest = buildDeleteByQuery(jobId, forecastIds);
-
         executeAsyncWithOrigin(client, ML_ORIGIN, DeleteByQueryAction.INSTANCE, deleteByQueryRequest, ActionListener.wrap(
             response -> {
                 if (response.isTimedOut()) {
@@ -208,45 +207,29 @@ private static Tuple<RestStatus, Throwable> getStatusAndReason(final BulkByScrol
         return new Tuple<>(status, reason);
     }
 
-    private static Set<ForecastRequestStats> parseForecastsFromSearch(SearchResponse searchResponse) throws IOException {
-        SearchHits hits = searchResponse.getHits();
-        List<ForecastRequestStats> allStats = new ArrayList<>(hits.getHits().length);
-        for (SearchHit hit : hits) {
-            try (InputStream stream = hit.getSourceRef().streamInput();
-                 XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(
-                     NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, stream)) {
-                allStats.add(ForecastRequestStats.STRICT_PARSER.apply(parser, null));
-            }
-        }
-        return new HashSet<>(allStats);
-    }
-
     private DeleteByQueryRequest buildDeleteByQuery(String jobId, List<String> forecastsToDelete) {
-        DeleteByQueryRequest request = new DeleteByQueryRequest()
-            .setAbortOnVersionConflict(false) //since these documents are not updated, a conflict just means it was deleted previously
-            .setMaxDocs(MAX_FORECAST_TO_SEARCH)
-            .setSlices(AbstractBulkByScrollRequest.AUTO_SLICES);
-
-        request.indices(AnomalyDetectorsIndex.jobResultsAliasedName(jobId));
-        BoolQueryBuilder innerBoolQuery = QueryBuilders.boolQuery();
-        innerBoolQuery
+        BoolQueryBuilder innerBoolQuery = QueryBuilders.boolQuery()
             .must(QueryBuilders.termsQuery(Result.RESULT_TYPE.getPreferredName(),
                 ForecastRequestStats.RESULT_TYPE_VALUE, Forecast.RESULT_TYPE_VALUE))
             .must(QueryBuilders.termsQuery(Forecast.FORECAST_ID.getPreferredName(),
                 forecastsToDelete));
-
         QueryBuilder query = QueryBuilders.boolQuery().filter(innerBoolQuery);
-        request.setQuery(query);
-        request.setRefresh(true);
-        return request;
+
+        // We want *all* of the docs to be deleted. Hence, we rely on the default value of max_docs.
+        return new DeleteByQueryRequest()
+            .setAbortOnVersionConflict(false) // since these documents are not updated, a conflict just means it was deleted previously
+            .setSlices(AbstractBulkByScrollRequest.AUTO_SLICES)
+            .indices(AnomalyDetectorsIndex.jobResultsAliasedName(jobId))
+            .setQuery(query)
+            .setRefresh(true);
     }
 
     private static void handleFailure(Exception e,
                                       DeleteForecastAction.Request request,
                                       ActionListener<AcknowledgedResponse> listener) {
         if (ExceptionsHelper.unwrapCause(e) instanceof ResourceNotFoundException) {
             if (request.isAllowNoForecasts() && Strings.isAllOrWildcard(request.getForecastId())) {
-                listener.onResponse(AcknowledgedResponse.of(true));
+                listener.onResponse(AcknowledgedResponse.TRUE);
             } else {
                 listener.onFailure(new ResourceNotFoundException(
                     Messages.getMessage(Messages.REST_NO_SUCH_FORECAST, request.getForecastId(), request.getJobId())