[Transform] Assign based on ProjectId

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

Implemented TransformPersistentTasksExecutor's getAssignment using the
ProjectId by reading the TransformMetadata and RoutingTable from the
ProjectMetadata.

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;
@@ -136,7 +138,8 @@ protected HealthNode createTask(
     public PersistentTasksCustomMetadata.Assignment getAssignment(
         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;
@@ -58,12 +59,29 @@ public Scope scope() {
 
     public static final Assignment NO_NODE_FOUND = new Assignment(null, "no appropriate nodes found for the assignment");
 
+    /**
+     * Use {@link #getAssignment(PersistentTaskParams, Collection, ClusterState, ProjectId)}.
+     */
+    @Deprecated
+    public PersistentTasksCustomMetadata.Assignment getAssignment(
+        Params params,
+        Collection<DiscoveryNode> candidateNodes,
+        ClusterState clusterState
+    ) {
+        return getAssignment(params, candidateNodes, clusterState, null);
+    }
+
     /**
      * 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
      */
-    public Assignment getAssignment(Params params, Collection<DiscoveryNode> candidateNodes, ClusterState clusterState) {
+    public Assignment getAssignment(
+        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 +123,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) {
+                public Assignment getAssignment(
+                    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) {
+        public Assignment getAssignment(
+            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.getAssignment(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) {
@@ -131,8 +133,9 @@ public void validate(ShardFollowTask params, ClusterState clusterState) {
     @Override
     public Assignment getAssignment(
         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

@@ -93,7 +93,8 @@ private void runAssignmentTest(
         final Assignment assignment = executor.getAssignment(
             mock(ShardFollowTask.class),
             clusterStateBuilder.nodes().getAllNodes(),
-            clusterStateBuilder.build()
+            clusterStateBuilder.build(),
+            null
         );
         consumer.accept(theSpecial, assignment);
     }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/TransformMetadata.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.cluster.Diff;
 import org.elasticsearch.cluster.NamedDiff;
 import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -209,6 +210,9 @@ public TransformMetadata build() {
         }
     }
 
+    /**
+     * @deprecated use {@link #transformMetadata(ClusterState, ProjectId)}
+     */
     @Deprecated(forRemoval = true)
     public static TransformMetadata getTransformMetadata(ClusterState state) {
         TransformMetadata TransformMetadata = (state == null) ? null : state.metadata().getSingleProjectCustom(TYPE);
@@ -218,6 +222,17 @@ public static TransformMetadata getTransformMetadata(ClusterState state) {
         return TransformMetadata;
     }
 
+    public static TransformMetadata transformMetadata(@Nullable ClusterState state, @Nullable ProjectId projectId) {
+        if (state == null || projectId == null) {
+            return EMPTY_METADATA;
+        }
+        TransformMetadata transformMetadata = state.metadata().getProject(projectId).custom(TYPE);
+        if (transformMetadata == null) {
+            return EMPTY_METADATA;
+        }
+        return transformMetadata;
+    }
+
     public static boolean upgradeMode(ClusterState state) {
         return getTransformMetadata(state).upgradeMode();
     }

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);
@@ -129,7 +131,8 @@ public void validate(DownsampleShardTaskParams params, ClusterState clusterState
     public PersistentTasksCustomMetadata.Assignment getAssignment(
         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, null);
         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, null);
         assertThat(result.getExecutorNode(), equalTo(node.getId()));
         assertThat(result.getExplanation(), equalTo("a node to fail and stop this persistent task"));
     }