From fcb15b0ce30e1a46d42f65874c1e4a91e2c3794c Mon Sep 17 00:00:00 2001
From: Dimitris Athanasiou <dimitris@elastic.co>
Date: Sun, 9 Sep 2018 22:53:03 +0100
Subject: [PATCH] [ML] Get job stats request should filter non-ML job tasks
 (#33516)

When requesting job stats for `_all`, all ES tasks are accepted
resulting to loads of cluster traffic and a memory overhead.
This commit correctly filters out non ML job tasks.

Closes #33515
---
 .../core/ml/action/GetJobsStatsAction.java    |  3 +--
 .../xpack/core/ml/action/OpenJobAction.java   | 11 ++++++++--
 .../action/GetJobStatsActionRequestTests.java |  9 +++++++++
 .../action/TransportOpenJobActionTests.java   | 20 +++++++++++++++++++
 4 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetJobsStatsAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetJobsStatsAction.java
index 807c0936375..d2d5d09090e 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetJobsStatsAction.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetJobsStatsAction.java
@@ -14,7 +14,6 @@ import org.elasticsearch.action.TaskOperationFailure;
 import org.elasticsearch.action.support.tasks.BaseTasksRequest;
 import org.elasticsearch.action.support.tasks.BaseTasksResponse;
 import org.elasticsearch.client.ElasticsearchClient;
-import org.elasticsearch.cluster.metadata.MetaData;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.ParseField;
@@ -95,7 +94,7 @@ public class GetJobsStatsAction extends Action<GetJobsStatsAction.Response> {
 
         @Override
         public boolean match(Task task) {
-            return jobId.equals(MetaData.ALL) || OpenJobAction.JobTaskMatcher.match(task, jobId);
+            return OpenJobAction.JobTaskMatcher.match(task, jobId);
         }
 
         @Override
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/OpenJobAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/OpenJobAction.java
index fc38d974def..bbc39c7d731 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/OpenJobAction.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/OpenJobAction.java
@@ -12,6 +12,7 @@ import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.action.support.master.MasterNodeRequest;
 import org.elasticsearch.client.ElasticsearchClient;
+import org.elasticsearch.cluster.metadata.MetaData;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -239,8 +240,14 @@ public class OpenJobAction extends Action<AcknowledgedResponse> {
     public interface JobTaskMatcher {
 
         static boolean match(Task task, String expectedJobId) {
-            String expectedDescription = "job-" + expectedJobId;
-            return task instanceof JobTaskMatcher && expectedDescription.equals(task.getDescription());
+            if (task instanceof JobTaskMatcher) {
+                if (MetaData.ALL.equals(expectedJobId)) {
+                    return true;
+                }
+                String expectedDescription = "job-" + expectedJobId;
+                return expectedDescription.equals(task.getDescription());
+            }
+            return false;
         }
     }
 
diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetJobStatsActionRequestTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetJobStatsActionRequestTests.java
index 913618de38b..edf3f73a8af 100644
--- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetJobStatsActionRequestTests.java
+++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetJobStatsActionRequestTests.java
@@ -6,9 +6,13 @@
 package org.elasticsearch.xpack.core.ml.action;
 
 import org.elasticsearch.cluster.metadata.MetaData;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.test.AbstractStreamableTestCase;
 import org.elasticsearch.xpack.core.ml.action.GetJobsStatsAction.Request;
 
+import static org.hamcrest.Matchers.is;
+import static org.mockito.Mockito.mock;
+
 public class GetJobStatsActionRequestTests extends AbstractStreamableTestCase<Request> {
 
     @Override
@@ -23,4 +27,9 @@ public class GetJobStatsActionRequestTests extends AbstractStreamableTestCase<Re
         return new Request();
     }
 
+    public void testMatch_GivenAll_FailsForNonJobTasks() {
+        Task nonJobTask = mock(Task.class);
+
+        assertThat(new Request("_all").match(nonJobTask), is(false));
+    }
 }
diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportOpenJobActionTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportOpenJobActionTests.java
index bef7705e835..58b60273b0e 100644
--- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportOpenJobActionTests.java
+++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportOpenJobActionTests.java
@@ -30,6 +30,7 @@ import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.persistent.PersistentTasksCustomMetaData;
 import org.elasticsearch.persistent.PersistentTasksCustomMetaData.Assignment;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.VersionUtils;
 import org.elasticsearch.xpack.core.ml.MlMetaIndex;
@@ -66,6 +67,7 @@ import java.util.function.Function;
 
 import static org.elasticsearch.xpack.core.ml.job.config.JobTests.buildJobBuilder;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.is;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -626,6 +628,24 @@ public class TransportOpenJobActionTests extends ESTestCase {
         assertEquals("{_node_name1}{ml.machine_memory=5}{node.ml=true}", TransportOpenJobAction.nodeNameAndMlAttributes(node));
     }
 
+    public void testJobTaskMatcherMatch() {
+        Task nonJobTask1 = mock(Task.class);
+        Task nonJobTask2 = mock(Task.class);
+        TransportOpenJobAction.JobTask jobTask1 = new TransportOpenJobAction.JobTask("ml-1",
+                0, "persistent", "", null, null);
+        TransportOpenJobAction.JobTask jobTask2 = new TransportOpenJobAction.JobTask("ml-2",
+                1, "persistent", "", null, null);
+
+        assertThat(OpenJobAction.JobTaskMatcher.match(nonJobTask1, "_all"), is(false));
+        assertThat(OpenJobAction.JobTaskMatcher.match(nonJobTask2, "_all"), is(false));
+        assertThat(OpenJobAction.JobTaskMatcher.match(jobTask1, "_all"), is(true));
+        assertThat(OpenJobAction.JobTaskMatcher.match(jobTask2, "_all"), is(true));
+        assertThat(OpenJobAction.JobTaskMatcher.match(jobTask1, "ml-1"), is(true));
+        assertThat(OpenJobAction.JobTaskMatcher.match(jobTask2, "ml-1"), is(false));
+        assertThat(OpenJobAction.JobTaskMatcher.match(jobTask1, "ml-2"), is(false));
+        assertThat(OpenJobAction.JobTaskMatcher.match(jobTask2, "ml-2"), is(true));
+    }
+
     public static void addJobTask(String jobId, String nodeId, JobState jobState, PersistentTasksCustomMetaData.Builder builder) {
         builder.addTask(MlTasks.jobTaskId(jobId), OpenJobAction.TASK_NAME, new OpenJobAction.JobParams(jobId),
                 new Assignment(nodeId, "test assignment"));