[ML] Improve cleanup for DF Analytics HLRC tests (#56101)

Adds the step of stopping all data frame analytics before deleting them to the cleanup of the corresponding HLRC tests. Closes #56097
elastic · May 4, 2020 · 254cdd9 · 254cdd9
1 parent 1ea5eb4
commit 254cdd9
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git a/client/rest-high-level/src/test/java/org/elasticsearch/client/MlTestStateCleaner.java b/client/rest-high-level/src/test/java/org/elasticsearch/client/MlTestStateCleaner.java
@@ -29,6 +29,7 @@
 import org.elasticsearch.client.ml.GetDatafeedResponse;
 import org.elasticsearch.client.ml.GetJobRequest;
 import org.elasticsearch.client.ml.GetJobResponse;
+import org.elasticsearch.client.ml.StopDataFrameAnalyticsRequest;
 import org.elasticsearch.client.ml.StopDatafeedRequest;
 import org.elasticsearch.client.ml.datafeed.DatafeedConfig;
 import org.elasticsearch.client.ml.dataframe.DataFrameAnalyticsConfig;
@@ -106,10 +107,28 @@ private void closeAllJobs() {
     }
 
     private void deleteAllDataFrameAnalytics() throws IOException {
+        stopAllDataFrameAnalytics();
+
         GetDataFrameAnalyticsResponse getDataFrameAnalyticsResponse =
             mlClient.getDataFrameAnalytics(GetDataFrameAnalyticsRequest.getAllDataFrameAnalyticsRequest(), RequestOptions.DEFAULT);
         for (DataFrameAnalyticsConfig config : getDataFrameAnalyticsResponse.getAnalytics()) {
             mlClient.deleteDataFrameAnalytics(new DeleteDataFrameAnalyticsRequest(config.getId()), RequestOptions.DEFAULT);
         }
     }
+
+    private void stopAllDataFrameAnalytics() {
+        StopDataFrameAnalyticsRequest stopAllRequest = new StopDataFrameAnalyticsRequest("*");
+        try {
+            mlClient.stopDataFrameAnalytics(stopAllRequest, RequestOptions.DEFAULT);
+        } catch (Exception e1) {
+            logger.warn("failed to stop all data frame analytics. Will proceed to force-stopping", e1);
+            stopAllRequest.setForce(true);
+            try {
+                mlClient.stopDataFrameAnalytics(stopAllRequest, RequestOptions.DEFAULT);
+            } catch (Exception e2) {
+                logger.warn("force-stopping all data frame analytics failed", e2);
+            }
+            throw new RuntimeException("Had to resort to force-stopping data frame analytics, something went wrong?", e1);
+        }
+    }
 }
diff --git a/...ore/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java b/...ore/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java
@@ -94,6 +94,8 @@ private void deleteAllJobs() throws IOException {
     }
 
     private void deleteAllDataFrameAnalytics() throws IOException {
+        stopAllDataFrameAnalytics();
+
         final Request analyticsRequest = new Request("GET", "/_ml/data_frame/analytics?size=10000");
         analyticsRequest.addParameter("filter_path", "data_frame_analytics");
         final Response analyticsResponse = adminClient.performRequest(analyticsRequest);
@@ -108,4 +110,18 @@ private void deleteAllDataFrameAnalytics() throws IOException {
             adminClient.performRequest(new Request("DELETE", "/_ml/data_frame/analytics/" + id));
         }
     }
+
+    private void stopAllDataFrameAnalytics() {
+        try {
+            adminClient.performRequest(new Request("POST", "_ml/data_frame/analytics/*/_stop"));
+        } catch (Exception e1) {
+            logger.warn("failed to stop all data frame analytics. Will proceed to force-stopping", e1);
+            try {
+                adminClient.performRequest(new Request("POST", "_ml/data_frame/analytics/*/_stop?force=true"));
+            } catch (Exception e2) {
+                logger.warn("Force-stopping all data frame analytics failed", e2);
+            }
+            throw new RuntimeException("Had to resort to force-stopping data frame analytics, something went wrong?", e1);
+        }
+    }
 }