Refactor tasks to improve APM support (#87917)

Part of #84369. Split out from #87696. Rework how some work is executed
by creating child tasks for them, so that when traced by APM, it results
in more meaningful parent and child tasks in the UI. It also improves
how Elasticsearch is modelling the work.

server/src/internalClusterTest/java/org/elasticsearch/indices/recovery/TaskRecoveryIT.java

@@ -0,0 +1,137 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.indices.recovery;
+
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.CollectionUtils;
+import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.engine.EngineFactory;
+import org.elasticsearch.index.engine.InternalEngine;
+import org.elasticsearch.plugins.EnginePlugin;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.plugins.PluginsService;
+import org.elasticsearch.tasks.TaskInfo;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.stream.StreamSupport;
+
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.hamcrest.Matchers.equalTo;
+
+@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST, numDataNodes = 0)
+public class TaskRecoveryIT extends ESIntegTestCase {
+
+    @Override
+    protected boolean addMockInternalEngine() {
+        return false;
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return CollectionUtils.appendToCopy(super.nodePlugins(), TaskRecoveryIT.EngineTestPlugin.class);
+    }
+
+    /**
+     * Checks that the parent / child task hierarchy is correct for tasks that are initiated by a recovery task.
+     * We use an engine plugin that stalls translog recovery, which gives us the opportunity to inspect the
+     * task hierarchy.
+     */
+    public void testTaskForOngoingRecovery() throws Exception {
+        String indexName = "test";
+        internalCluster().startMasterOnlyNode();
+        String nodeWithPrimary = internalCluster().startDataOnlyNode();
+        assertAcked(
+            client().admin()
+                .indices()
+                .prepareCreate(indexName)
+                .setSettings(
+                    Settings.builder()
+                        .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                        .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+                        .put("index.routing.allocation.include._name", nodeWithPrimary)
+                )
+        );
+        try {
+            String nodeWithReplica = internalCluster().startDataOnlyNode();
+
+            // Create an index so that there is something to recover
+            assertAcked(
+                client().admin()
+                    .indices()
+                    .prepareUpdateSettings(indexName)
+                    .setSettings(
+                        Settings.builder()
+                            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+                            .put("index.routing.allocation.include._name", nodeWithPrimary + "," + nodeWithReplica)
+                    )
+            );
+            // Translog recovery is stalled, so we can inspect the running tasks.
+            assertBusy(() -> {
+                List<TaskInfo> primaryTasks = client().admin()
+                    .cluster()
+                    .prepareListTasks(nodeWithPrimary)
+                    .setActions(PeerRecoverySourceService.Actions.START_RECOVERY)
+                    .get()
+                    .getTasks();
+                assertThat("Expected a single primary task", primaryTasks.size(), equalTo(1));
+                List<TaskInfo> replicaTasks = client().admin()
+                    .cluster()
+                    .prepareListTasks(nodeWithReplica)
+                    .setActions(PeerRecoveryTargetService.Actions.PREPARE_TRANSLOG)
+                    .get()
+                    .getTasks();
+                assertThat("Expected a single replica task", replicaTasks.size(), equalTo(1));
+                assertThat(
+                    "Replica task's parent task ID was incorrect",
+                    replicaTasks.get(0).parentTaskId(),
+                    equalTo(primaryTasks.get(0).taskId())
+                );
+            });
+        } finally {
+            // Release the EngineTestPlugin, which will allow translog recovery to complete
+            StreamSupport.stream(internalCluster().getInstances(PluginsService.class).spliterator(), false)
+                .flatMap(ps -> ps.filterPlugins(EnginePlugin.class).stream())
+                .map(EngineTestPlugin.class::cast)
+                .forEach(EngineTestPlugin::release);
+        }
+        ensureGreen(indexName);
+    }
+
+    /**
+     * An engine plugin that defers translog recovery until the engine is released via {@link #release()}.
+     */
+    public static class EngineTestPlugin extends Plugin implements EnginePlugin {
+        private final CountDownLatch latch = new CountDownLatch(1);
+
+        public void release() {
+            latch.countDown();
+        }
+
+        @Override
+        public Optional<EngineFactory> getEngineFactory(IndexSettings indexSettings) {
+            return Optional.of(config -> new InternalEngine(config) {
+
+                @Override
+                public void skipTranslogRecovery() {
+                    try {
+                        latch.await();
+                    } catch (InterruptedException e) {
+                        throw new AssertionError(e);
+                    }
+                    super.skipTranslogRecovery();
+                }
+            });
+        }
+    }
+}

server/src/main/java/org/elasticsearch/cluster/ClusterStatePublicationEvent.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.cluster;
 
 import org.elasticsearch.cluster.service.BatchSummary;
+import org.elasticsearch.tasks.Task;
 
 /**
  * Represents a cluster state update computed by the {@link org.elasticsearch.cluster.service.MasterService} for publication to the cluster.
@@ -24,6 +25,7 @@ public class ClusterStatePublicationEvent {
     private final BatchSummary summary;
     private final ClusterState oldState;
     private final ClusterState newState;
+    private final Task task;
     private final long computationTimeMillis;
     private final long publicationStartTimeMillis;
     private volatile long publicationContextConstructionElapsedMillis = NOT_SET;
@@ -35,12 +37,14 @@ public ClusterStatePublicationEvent(
         BatchSummary summary,
         ClusterState oldState,
         ClusterState newState,
+        Task task,
         long computationTimeMillis,
         long publicationStartTimeMillis
     ) {
         this.summary = summary;
         this.oldState = oldState;
         this.newState = newState;
+        this.task = task;
         this.computationTimeMillis = computationTimeMillis;
         this.publicationStartTimeMillis = publicationStartTimeMillis;
     }
@@ -57,6 +61,10 @@ public ClusterState getNewState() {
         return newState;
     }
 
+    public Task getTask() {
+        return task;
+    }
+
     public long getComputationTimeMillis() {
         return computationTimeMillis;
     }

server/src/main/java/org/elasticsearch/cluster/coordination/PublicationTransportHandler.java

@@ -36,6 +36,7 @@
 import org.elasticsearch.core.AbstractRefCounted;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.Releasables;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.BytesTransportRequest;
 import org.elasticsearch.transport.TransportException;
@@ -284,6 +285,7 @@ public class PublicationContext extends AbstractRefCounted {
         private final DiscoveryNodes discoveryNodes;
         private final ClusterState newState;
         private final ClusterState previousState;
+        private final Task task;
         private final boolean sendFullVersion;
 
         // All the values of these maps have one ref for the context (while it's open) and one for each in-flight message.
@@ -294,6 +296,7 @@ public class PublicationContext extends AbstractRefCounted {
             discoveryNodes = clusterStatePublicationEvent.getNewState().nodes();
             newState = clusterStatePublicationEvent.getNewState();
             previousState = clusterStatePublicationEvent.getOldState();
+            task = clusterStatePublicationEvent.getTask();
             sendFullVersion = previousState.getBlocks().disableStatePersistence();
         }
 
@@ -421,10 +424,11 @@ private void sendClusterState(
                 return;
             }
             try {
-                transportService.sendRequest(
+                transportService.sendChildRequest(
                     destination,
                     PUBLISH_STATE_ACTION_NAME,
                     new BytesTransportRequest(bytes, destination.getVersion()),
+                    task,
                     STATE_REQUEST_OPTIONS,
                     new ActionListenerResponseHandler<>(
                         ActionListener.runAfter(listener, bytes::decRef),

server/src/main/java/org/elasticsearch/cluster/service/ClusterService.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.node.Node;
+import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.threadpool.ThreadPool;
 
 public class ClusterService extends AbstractLifecycleComponent {
@@ -55,11 +56,11 @@ public class ClusterService extends AbstractLifecycleComponent {
 
     private RerouteService rerouteService;
 
-    public ClusterService(Settings settings, ClusterSettings clusterSettings, ThreadPool threadPool) {
+    public ClusterService(Settings settings, ClusterSettings clusterSettings, ThreadPool threadPool, TaskManager taskManager) {
         this(
             settings,
             clusterSettings,
-            new MasterService(settings, clusterSettings, threadPool),
+            new MasterService(settings, clusterSettings, threadPool, taskManager),
             new ClusterApplierService(Node.NODE_NAME_SETTING.get(settings), settings, clusterSettings, threadPool)
         );
     }
@@ -262,5 +263,4 @@ public <T extends ClusterStateTaskListener> void submitStateUpdateTask(
     ) {
         masterService.submitStateUpdateTask(source, task, config, executor);
     }
-
 }