Preserve Task Id for ML Datafeed (#54943)

This change preserves the task id for internal requests for the `StartDatafeedPersistentTask`.

Task ids are a way to express a relationship between related internal requests.
In this particular case, the task ids are used for debugging and (soon) security auditing,
but not for task cancellation, because there is already a graceful-shutdown of child
internal requests (given a task id) in place.

Author: albertzaharovits

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

@@ -460,6 +460,12 @@ protected void onCancelled() {
             stop(getReasonCancelled(), TimeValue.ZERO);
         }
 
+        @Override
+        public boolean shouldCancelChildrenOnCancellation() {
+            // onCancelled implements graceful shutdown of children
+            return false;
+        }
+
         public void stop(String reason, TimeValue timeout) {
             if (datafeedManager != null) {
                 datafeedManager.stopDatafeed(this, reason, timeout);

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/datafeed/DatafeedJobBuilder.java

@@ -8,11 +8,13 @@
 import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.Client;
+import org.elasticsearch.client.ParentTaskAssigningClient;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.license.RemoteClusterLicenseChecker;
 import org.elasticsearch.node.Node;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xpack.core.action.util.QueryPage;
 import org.elasticsearch.xpack.core.ml.annotations.AnnotationPersister;
 import org.elasticsearch.xpack.core.ml.datafeed.DatafeedConfig;
@@ -74,16 +76,17 @@ public DatafeedJobBuilder(Client client, NamedXContentRegistry xContentRegistry,
         this.nodeName = nodeName;
     }
 
-    void build(String datafeedId, ActionListener<DatafeedJob> listener) {
+    void build(String datafeedId, TaskId parentTaskId, ActionListener<DatafeedJob> listener) {
         AtomicReference<Job> jobHolder = new AtomicReference<>();
         AtomicReference<DatafeedConfig> datafeedConfigHolder = new AtomicReference<>();
+        final ParentTaskAssigningClient parentTaskAssigningClient = new ParentTaskAssigningClient(client, parentTaskId);
 
         // Step 5. Build datafeed job object
         Consumer<Context> contextHanlder = context -> {
             TimeValue frequency = getFrequencyOrDefault(datafeedConfigHolder.get(), jobHolder.get(), xContentRegistry);
             TimeValue queryDelay = datafeedConfigHolder.get().getQueryDelay();
-            DelayedDataDetector delayedDataDetector =
-                    DelayedDataDetectorFactory.buildDetector(jobHolder.get(), datafeedConfigHolder.get(), client, xContentRegistry);
+            DelayedDataDetector delayedDataDetector = DelayedDataDetectorFactory.buildDetector(jobHolder.get(),
+                    datafeedConfigHolder.get(), parentTaskAssigningClient, xContentRegistry);
             DatafeedJob datafeedJob =
                 new DatafeedJob(
                     jobHolder.get().getId(),
@@ -92,7 +95,7 @@ void build(String datafeedId, ActionListener<DatafeedJob> listener) {
                     queryDelay.millis(),
                     context.dataExtractorFactory,
                     context.timingStatsReporter,
-                    client,
+                    parentTaskAssigningClient,
                     auditor,
                     annotationPersister,
                     currentTimeSupplier,
@@ -123,7 +126,7 @@ void build(String datafeedId, ActionListener<DatafeedJob> listener) {
             context.timingStatsReporter =
                 new DatafeedTimingStatsReporter(initialTimingStats, jobResultsPersister::persistDatafeedTimingStats);
             DataExtractorFactory.create(
-                client,
+                parentTaskAssigningClient,
                 datafeedConfigHolder.get(),
                 jobHolder.get(),
                 xContentRegistry,

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/datafeed/DatafeedManager.java

@@ -108,7 +108,7 @@ public void onFailure(Exception e) {
                 }, finishHandler::accept
         );
 
-        datafeedJobBuilder.build(datafeedId, datafeedJobHandler);
+        datafeedJobBuilder.build(datafeedId, task.getParentTaskId(), datafeedJobHandler);
     }
 
     public void stopDatafeed(TransportStartDatafeedAction.DatafeedTask task, String reason, TimeValue timeout) {

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/datafeed/DatafeedJobBuilderTests.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.mock.orig.Mockito;
 import org.elasticsearch.node.Node;
+import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.action.util.QueryPage;
@@ -124,7 +125,7 @@ public void testBuild_GivenScrollDatafeedAndNewJob() throws Exception {
         givenJob(jobBuilder);
         givenDatafeed(datafeed);
 
-        datafeedJobBuilder.build("datafeed1", datafeedJobHandler);
+        datafeedJobBuilder.build("datafeed1", new TaskId(""), datafeedJobHandler);
 
         assertBusy(() -> wasHandlerCalled.get());
     }
@@ -152,7 +153,7 @@ public void testBuild_GivenScrollDatafeedAndOldJobWithLatestRecordTimestampAfter
         givenJob(jobBuilder);
         givenDatafeed(datafeed);
 
-        datafeedJobBuilder.build("datafeed1", datafeedJobHandler);
+        datafeedJobBuilder.build("datafeed1", new TaskId(""), datafeedJobHandler);
 
         assertBusy(() -> wasHandlerCalled.get());
     }
@@ -180,7 +181,7 @@ public void testBuild_GivenScrollDatafeedAndOldJobWithLatestBucketAfterLatestRec
         givenJob(jobBuilder);
         givenDatafeed(datafeed);
 
-        datafeedJobBuilder.build("datafeed1", datafeedJobHandler);
+        datafeedJobBuilder.build("datafeed1", new TaskId(""), datafeedJobHandler);
 
         assertBusy(() -> wasHandlerCalled.get());
     }
@@ -205,7 +206,7 @@ public void testBuild_GivenBucketsRequestFails() {
         givenJob(jobBuilder);
         givenDatafeed(datafeed);
 
-        datafeedJobBuilder.build("datafeed1", ActionListener.wrap(datafeedJob -> fail(), taskHandler));
+        datafeedJobBuilder.build("datafeed1", new TaskId(""), ActionListener.wrap(datafeedJob -> fail(), taskHandler));
 
         verify(taskHandler).accept(error);
     }
@@ -247,7 +248,7 @@ public void testBuildGivenRemoteIndicesButNoRemoteSearching() throws Exception {
 
         givenJob(jobBuilder);
         givenDatafeed(datafeed);
-        datafeedJobBuilder.build("datafeed1", datafeedJobHandler);
+        datafeedJobBuilder.build("datafeed1", new TaskId(""), datafeedJobHandler);
         assertBusy(() -> wasHandlerCalled.get());
     }
 

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/datafeed/DatafeedManagerTests.java

@@ -118,10 +118,10 @@ public void setUpTests() {
         DatafeedJobBuilder datafeedJobBuilder = mock(DatafeedJobBuilder.class);
         doAnswer(invocationOnMock -> {
             @SuppressWarnings("rawtypes")
-            ActionListener listener = (ActionListener) invocationOnMock.getArguments()[1];
+            ActionListener listener = (ActionListener) invocationOnMock.getArguments()[2];
             listener.onResponse(datafeedJob);
             return null;
-        }).when(datafeedJobBuilder).build(any(), any());
+        }).when(datafeedJobBuilder).build(any(), any(), any());
 
         hasOpenAutodetectCommunicator = new AtomicBoolean(true);
         AutodetectProcessManager autodetectProcessManager = mock(AutodetectProcessManager.class);