[Persistent Tasks] Assign based on ProjectId (elastic#130391)

Pass the ProjectId from PersistentTaskClusterService through to all
PersistentTasksExecutors when creating node assignments.

These PersistentTasksExecutors require the ProjectId during node
assignment:
- OpenJobPersistentTasksExecutor
- SnapshotUpgradeTaskExecutor
- StartDatafeedPersistentTasksExecutor
- TransformPersistentTasksExecutor

Author: prwhelan

server/src/main/java/org/elasticsearch/health/node/selection/HealthNodeTaskExecutor.java

@@ -16,12 +16,14 @@
 import org.elasticsearch.cluster.ClusterChangedEvent;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ClusterStateListener;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.node.NodeClosedException;
 import org.elasticsearch.persistent.AllocatedPersistentTask;
@@ -133,10 +135,11 @@ protected HealthNode createTask(
      * Returns the node id from the eligible health nodes
      */
     @Override
-    public PersistentTasksCustomMetadata.Assignment getAssignment(
+    protected PersistentTasksCustomMetadata.Assignment doGetAssignment(
         HealthNodeTaskParams params,
         Collection<DiscoveryNode> candidateNodes,
-        ClusterState clusterState
+        ClusterState clusterState,
+        @Nullable ProjectId projectId
     ) {
         DiscoveryNode discoveryNode = selectLeastLoadedNode(clusterState, candidateNodes, DiscoveryNode::canContainData);
         if (discoveryNode == null) {

server/src/main/java/org/elasticsearch/persistent/PersistentTasksClusterService.java

@@ -171,9 +171,9 @@ public ClusterState execute(ClusterState currentState) {
                 assert (projectId == null && taskExecutor.scope() == PersistentTasksExecutor.Scope.CLUSTER)
                     || (projectId != null && taskExecutor.scope() == PersistentTasksExecutor.Scope.PROJECT)
                     : "inconsistent project-id [" + projectId + "] and task scope [" + taskExecutor.scope() + "]";
-                taskExecutor.validate(taskParams, currentState);
+                taskExecutor.validate(taskParams, currentState, projectId);
 
-                Assignment assignment = createAssignment(taskName, taskParams, currentState);
+                Assignment assignment = createAssignment(taskName, taskParams, currentState, projectId);
                 logger.debug("creating {} persistent task [{}] with assignment [{}]", taskTypeString(projectId), taskName, assignment);
                 return builder.addTask(taskId, taskName, taskParams, assignment).buildAndUpdate(currentState, projectId);
             }
@@ -449,7 +449,8 @@ public void clusterStateProcessed(ClusterState oldState, ClusterState newState)
     private <Params extends PersistentTaskParams> Assignment createAssignment(
         final String taskName,
         final Params taskParams,
-        final ClusterState currentState
+        final ClusterState currentState,
+        @Nullable final ProjectId projectId
     ) {
         PersistentTasksExecutor<Params> persistentTasksExecutor = registry.getPersistentTaskExecutorSafe(taskName);
 
@@ -468,7 +469,7 @@ private <Params extends PersistentTaskParams> Assignment createAssignment(
         // Task assignment should not rely on node order
         Randomness.shuffle(candidateNodes);
 
-        final Assignment assignment = persistentTasksExecutor.getAssignment(taskParams, candidateNodes, currentState);
+        final Assignment assignment = persistentTasksExecutor.getAssignment(taskParams, candidateNodes, currentState, projectId);
         assert assignment != null : "getAssignment() should always return an Assignment object, containing a node or a reason why not";
         assert (assignment.getExecutorNode() == null
             || currentState.metadata().nodeShutdowns().contains(assignment.getExecutorNode()) == false)
@@ -540,8 +541,8 @@ public void clusterStateProcessed(ClusterState oldState, ClusterState newState)
      * persistent tasks changed.
      */
     boolean shouldReassignPersistentTasks(final ClusterChangedEvent event) {
-        final List<PersistentTasks> allTasks = PersistentTasks.getAllTasks(event.state()).map(Tuple::v2).toList();
-        if (allTasks.isEmpty()) {
+        var projectIdToTasksIterator = PersistentTasks.getAllTasks(event.state()).iterator();
+        if (projectIdToTasksIterator.hasNext() == false) {
             return false;
         }
 
@@ -553,10 +554,16 @@ boolean shouldReassignPersistentTasks(final ClusterChangedEvent event) {
             || event.metadataChanged()
             || masterChanged) {
 
-            for (PersistentTasks tasks : allTasks) {
-                for (PersistentTask<?> task : tasks.tasks()) {
+            while (projectIdToTasksIterator.hasNext()) {
+                var projectIdToTasks = projectIdToTasksIterator.next();
+                for (PersistentTask<?> task : projectIdToTasks.v2().tasks()) {
                     if (needsReassignment(task.getAssignment(), event.state().nodes())) {
-                        Assignment assignment = createAssignment(task.getTaskName(), task.getParams(), event.state());
+                        Assignment assignment = createAssignment(
+                            task.getTaskName(),
+                            task.getParams(),
+                            event.state(),
+                            projectIdToTasks.v1()
+                        );
                         if (Objects.equals(assignment, task.getAssignment()) == false) {
                             return true;
                         }
@@ -602,7 +609,7 @@ private ClusterState reassignClusterOrSingleProjectTasks(@Nullable final Project
             // We need to check if removed nodes were running any of the tasks and reassign them
             for (PersistentTask<?> task : tasks.tasks()) {
                 if (needsReassignment(task.getAssignment(), nodes)) {
-                    Assignment assignment = createAssignment(task.getTaskName(), task.getParams(), clusterState);
+                    Assignment assignment = createAssignment(task.getTaskName(), task.getParams(), clusterState, projectId);
                     if (Objects.equals(assignment, task.getAssignment()) == false) {
                         logger.trace(
                             "reassigning {} task {} from node {} to node {}",

server/src/main/java/org/elasticsearch/persistent/PersistentTasksExecutor.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.persistent;
 
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Tuple;
@@ -63,7 +64,31 @@ public Scope scope() {
      * <p>
      * The default implementation returns the least loaded data node from amongst the collection of candidate nodes
      */
-    public Assignment getAssignment(Params params, Collection<DiscoveryNode> candidateNodes, ClusterState clusterState) {
+    public final Assignment getAssignment(
+        Params params,
+        Collection<DiscoveryNode> candidateNodes,
+        ClusterState clusterState,
+        @Nullable ProjectId projectId
+    ) {
+        assert (scope() == Scope.PROJECT && projectId != null) || (scope() == Scope.CLUSTER && projectId == null)
+            : "inconsistent project-id [" + projectId + "] and task scope [" + scope() + "]";
+        return doGetAssignment(params, candidateNodes, clusterState, projectId);
+    }
+
+    /**
+     * Returns the node id where the params has to be executed,
+     * <p>
+     * The default implementation returns the least loaded data node from amongst the collection of candidate nodes.
+     * <p>
+     * If {@link #scope()} returns CLUSTER, then {@link ProjectId} will be null.
+     * If {@link #scope()} returns PROJECT, then {@link ProjectId} will not be null.
+     */
+    protected Assignment doGetAssignment(
+        Params params,
+        Collection<DiscoveryNode> candidateNodes,
+        ClusterState clusterState,
+        @Nullable ProjectId projectId
+    ) {
         DiscoveryNode discoveryNode = selectLeastLoadedNode(clusterState, candidateNodes, DiscoveryNode::canContainData);
         if (discoveryNode == null) {
             return NO_NODE_FOUND;
@@ -105,7 +130,7 @@ protected DiscoveryNode selectLeastLoadedNode(
      * <p>
      * Throws an exception if the supplied params cannot be executed on the cluster in the current state.
      */
-    public void validate(Params params, ClusterState clusterState) {}
+    public void validate(Params params, ClusterState clusterState, @Nullable ProjectId projectId) {}
 
     /**
      * Creates a AllocatedPersistentTask for communicating with task manager

server/src/test/java/org/elasticsearch/persistent/PersistentTasksClusterServiceTests.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.metadata.NodesShutdownMetadata;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.metadata.ProjectMetadata;
 import org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -1087,7 +1088,12 @@ public Scope scope() {
                 }
 
                 @Override
-                public Assignment getAssignment(P params, Collection<DiscoveryNode> candidateNodes, ClusterState clusterState) {
+                protected Assignment doGetAssignment(
+                    P params,
+                    Collection<DiscoveryNode> candidateNodes,
+                    ClusterState clusterState,
+                    ProjectId projectId
+                ) {
                     return fn.apply(params, candidateNodes, clusterState);
                 }
 

server/src/test/java/org/elasticsearch/persistent/TestPersistentTasksPlugin.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.client.internal.ElasticsearchClient;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
@@ -326,12 +327,17 @@ public static void setNonClusterStateCondition(boolean nonClusterStateCondition)
         }
 
         @Override
-        public Assignment getAssignment(TestParams params, Collection<DiscoveryNode> candidateNodes, ClusterState clusterState) {
+        protected Assignment doGetAssignment(
+            TestParams params,
+            Collection<DiscoveryNode> candidateNodes,
+            ClusterState clusterState,
+            ProjectId projectId
+        ) {
             if (nonClusterStateCondition == false) {
                 return new Assignment(null, "non cluster state condition prevents assignment");
             }
             if (params == null || params.getExecutorNodeAttr() == null) {
-                return super.getAssignment(params, candidateNodes, clusterState);
+                return super.doGetAssignment(params, candidateNodes, clusterState, projectId);
             } else {
                 DiscoveryNode executorNode = selectLeastLoadedNode(
                     clusterState,

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/ShardFollowTasksExecutor.java

@@ -32,6 +32,7 @@
 import org.elasticsearch.cluster.metadata.AliasMetadata;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.MappingMetadata;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.IndexRoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
@@ -43,6 +44,7 @@
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
 import org.elasticsearch.core.CheckedConsumer;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexNotFoundException;
@@ -118,7 +120,7 @@ public ShardFollowTasksExecutor(Client client, ThreadPool threadPool, ClusterSer
     }
 
     @Override
-    public void validate(ShardFollowTask params, ClusterState clusterState) {
+    public void validate(ShardFollowTask params, ClusterState clusterState, @Nullable ProjectId projectId) {
         final IndexRoutingTable routingTable = clusterState.getRoutingTable().index(params.getFollowShardId().getIndex());
         final ShardRouting primaryShard = routingTable.shard(params.getFollowShardId().id()).primaryShard();
         if (primaryShard.active() == false) {
@@ -129,10 +131,11 @@ public void validate(ShardFollowTask params, ClusterState clusterState) {
     private static final Assignment NO_ASSIGNMENT = new Assignment(null, "no nodes found with data and remote cluster client roles");
 
     @Override
-    public Assignment getAssignment(
+    protected Assignment doGetAssignment(
         final ShardFollowTask params,
-        Collection<DiscoveryNode> candidateNodes,
-        final ClusterState clusterState
+        final Collection<DiscoveryNode> candidateNodes,
+        final ClusterState clusterState,
+        @Nullable final ProjectId projectId
     ) {
         final DiscoveryNode node = selectLeastLoadedNode(
             clusterState,

x-pack/plugin/ccr/src/test/java/org/elasticsearch/xpack/ccr/action/ShardFollowTasksExecutorAssignmentTests.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.cluster.node.DiscoveryNodeUtils;
@@ -93,7 +94,8 @@ private void runAssignmentTest(
         final Assignment assignment = executor.getAssignment(
             mock(ShardFollowTask.class),
             clusterStateBuilder.nodes().getAllNodes(),
-            clusterStateBuilder.build()
+            clusterStateBuilder.build(),
+            ProjectId.DEFAULT
         );
         consumer.accept(theSpecial, assignment);
     }

x-pack/plugin/downsample/src/main/java/org/elasticsearch/xpack/downsample/DownsampleShardPersistentTaskExecutor.java

@@ -22,13 +22,15 @@
 import org.elasticsearch.action.support.TransportAction;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.index.mapper.TimeSeriesIdFieldMapper;
 import org.elasticsearch.index.shard.ShardId;
@@ -116,7 +118,7 @@ protected AllocatedPersistentTask createTask(
     }
 
     @Override
-    public void validate(DownsampleShardTaskParams params, ClusterState clusterState) {
+    public void validate(DownsampleShardTaskParams params, ClusterState clusterState, @Nullable ProjectId projectId) {
         // This is just a pre-check, but doesn't prevent from avoiding from aborting the task when source index disappeared
         // after initial creation of the persistent task.
         var indexShardRouting = findShardRoutingTable(params.shardId(), clusterState);
@@ -126,10 +128,11 @@ public void validate(DownsampleShardTaskParams params, ClusterState clusterState
     }
 
     @Override
-    public PersistentTasksCustomMetadata.Assignment getAssignment(
+    protected PersistentTasksCustomMetadata.Assignment doGetAssignment(
         final DownsampleShardTaskParams params,
         final Collection<DiscoveryNode> candidateNodes,
-        final ClusterState clusterState
+        final ClusterState clusterState,
+        @Nullable final ProjectId projectId
     ) {
         // NOTE: downsampling works by running a task per each shard of the source index.
         // Here we make sure we assign the task to the actual node holding the shard identified by

x-pack/plugin/downsample/src/test/java/org/elasticsearch/xpack/downsample/DownsampleShardPersistentTaskExecutorTests.java

@@ -96,7 +96,7 @@ public void testGetAssignment() {
             Strings.EMPTY_ARRAY,
             Strings.EMPTY_ARRAY
         );
-        var result = executor.getAssignment(params, Set.of(node), clusterState);
+        var result = executor.getAssignment(params, Set.of(node), clusterState, projectId);
         assertThat(result.getExecutorNode(), equalTo(node.getId()));
     }
 
@@ -128,7 +128,7 @@ public void testGetAssignmentMissingIndex() {
             Strings.EMPTY_ARRAY,
             Strings.EMPTY_ARRAY
         );
-        var result = executor.getAssignment(params, Set.of(node), clusterState);
+        var result = executor.getAssignment(params, Set.of(node), clusterState, projectId);
         assertThat(result.getExecutorNode(), equalTo(node.getId()));
         assertThat(result.getExplanation(), equalTo("a node to fail and stop this persistent task"));
     }
@@ -165,7 +165,7 @@ public void testGetStatelessAssignment() {
             Strings.EMPTY_ARRAY,
             Strings.EMPTY_ARRAY
         );
-        var result = executor.getAssignment(params, Set.of(indexNode, searchNode), clusterState);
+        var result = executor.getAssignment(params, Set.of(indexNode, searchNode), clusterState, projectId);
         assertThat(result.getExecutorNode(), nullValue());
 
         // Assign a copy of the shard to a search node
@@ -185,7 +185,7 @@ public void testGetStatelessAssignment() {
                     .build()
             )
             .build();
-        result = executor.getAssignment(params, Set.of(indexNode, searchNode), clusterState);
+        result = executor.getAssignment(params, Set.of(indexNode, searchNode), clusterState, projectId);
         assertThat(result.getExecutorNode(), equalTo(searchNode.getId()));
     }
 

x-pack/plugin/migrate/src/main/java/org/elasticsearch/system_indices/task/SystemIndexMigrationExecutor.java

@@ -9,10 +9,12 @@
 
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.settings.IndexScopedSettings;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.indices.SystemIndices;
 import org.elasticsearch.persistent.AllocatedPersistentTask;
 import org.elasticsearch.persistent.PersistentTaskParams;
@@ -86,10 +88,11 @@ protected AllocatedPersistentTask createTask(
     }
 
     @Override
-    public PersistentTasksCustomMetadata.Assignment getAssignment(
+    protected PersistentTasksCustomMetadata.Assignment doGetAssignment(
         SystemIndexMigrationTaskParams params,
         Collection<DiscoveryNode> candidateNodes,
-        ClusterState clusterState
+        ClusterState clusterState,
+        @Nullable ProjectId projectId
     ) {
         // This should select from master-eligible nodes because we already require all master-eligible nodes to have all plugins installed.
         // However, due to a misunderstanding, this code as-written needs to run on the master node in particular. This is not a fundamental