Scope methods to Project and Cluster

Author: prwhelan

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

@@ -16,14 +16,12 @@
 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;
@@ -135,11 +133,10 @@ protected HealthNode createTask(
      * Returns the node id from the eligible health nodes
      */
     @Override
-    public PersistentTasksCustomMetadata.Assignment getAssignment(
+    public PersistentTasksCustomMetadata.Assignment getClusterScopedAssignment(
         HealthNodeTaskParams params,
         Collection<DiscoveryNode> candidateNodes,
-        ClusterState clusterState,
-        @Nullable ProjectId projectId
+        ClusterState clusterState
     ) {
         DiscoveryNode discoveryNode = selectLeastLoadedNode(clusterState, candidateNodes, DiscoveryNode::canContainData);
         if (discoveryNode == null) {

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

@@ -168,10 +168,15 @@ public ClusterState execute(ClusterState currentState) {
                 }
 
                 PersistentTasksExecutor<Params> taskExecutor = registry.getPersistentTaskExecutorSafe(taskName);
-                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, projectId);
+                if (taskExecutor.scope() == PersistentTasksExecutor.Scope.CLUSTER) {
+                    assert projectId == null : "project id must be null when persistent task is cluster scoped";
+                    taskExecutor.validateCluster(taskParams, currentState);
+                } else if (taskExecutor.scope() == PersistentTasksExecutor.Scope.PROJECT) {
+                    assert projectId != null : "project id must not be null when persistent task is project scoped";
+                    taskExecutor.validateProject(taskParams, currentState.projectState(projectId));
+                } else {
+                    assert false : "unknown scope " + taskExecutor.scope();
+                }
 
                 Assignment assignment = createAssignment(taskName, taskParams, currentState, projectId);
                 logger.debug("creating {} persistent task [{}] with assignment [{}]", taskTypeString(projectId), taskName, assignment);
@@ -469,7 +474,16 @@ 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, projectId);
+        final Assignment assignment = switch (persistentTasksExecutor.scope()) {
+            case PROJECT -> {
+                assert projectId != null : "project id must not be null when persistent task is project scoped";
+                yield persistentTasksExecutor.getProjectScopedAssignment(taskParams, candidateNodes, currentState.projectState(projectId));
+            }
+            case CLUSTER -> {
+                assert projectId == null : "project id must be null when persistent task is cluster scoped";
+                yield persistentTasksExecutor.getClusterScopedAssignment(taskParams, candidateNodes, currentState);
+            }
+        };
         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)
@@ -541,7 +555,7 @@ public void clusterStateProcessed(ClusterState oldState, ClusterState newState)
      * persistent tasks changed.
      */
     boolean shouldReassignPersistentTasks(final ClusterChangedEvent event) {
-        var projectIdToTasksIterator = PersistentTasks.getAllTasks(event.state()).iterator();
+        var projectIdToTasksIterator = getAllTasks(event.state()).iterator();
         if (projectIdToTasksIterator.hasNext() == false) {
             return false;
         }
@@ -814,7 +828,7 @@ public void runInternal() {
                 // TODO just run on the elected master?
                 final ClusterState state = clusterService.state();
                 logger.trace("periodic persistent task assignment check running for cluster state {}", state.getVersion());
-                if (isAnyTaskUnassigned(PersistentTasks.getAllTasks(state))) {
+                if (isAnyTaskUnassigned(getAllTasks(state))) {
                     reassignPersistentTasks();
                 }
             }

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

@@ -10,7 +10,7 @@
 package org.elasticsearch.persistent;
 
 import org.elasticsearch.cluster.ClusterState;
-import org.elasticsearch.cluster.metadata.ProjectId;
+import org.elasticsearch.cluster.ProjectState;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Tuple;
@@ -60,16 +60,26 @@ public Scope scope() {
     public static final Assignment NO_NODE_FOUND = new Assignment(null, "no appropriate nodes found for the assignment");
 
     /**
-     * Returns the node id where the params has to be executed,
+     * Returns the node id where the params has to be executed based on the Project scope,
      * <p>
-     * The default implementation returns the least loaded data node from amongst the collection of candidate nodes
+     * The default implementation returns the least loaded data node from amongst the collection of candidate nodes.
+     * <p>
+     * Implement this method if your {@link #scope()} returns PROJECT, or {@link #getClusterScopedAssignment} if your {@link #scope()}
+     * returns CLUSTER.
      */
-    public Assignment getAssignment(
-        Params params,
-        Collection<DiscoveryNode> candidateNodes,
-        ClusterState clusterState,
-        @Nullable ProjectId projectId
-    ) {
+    public Assignment getProjectScopedAssignment(Params params, Collection<DiscoveryNode> candidateNodes, ProjectState projectState) {
+        return getClusterScopedAssignment(params, candidateNodes, projectState.cluster());
+    }
+
+    /**
+     * Returns the node id where the params has to be executed based on the Cluster scope,
+     * <p>
+     * The default implementation returns the least loaded data node from amongst the collection of candidate nodes.
+     * <p>
+     * Implement this method if your {@link #scope()} returns CLUSTER, or {@link #getProjectScopedAssignment} if your {@link #scope()}
+     * returns PROJECT.
+     */
+    public Assignment getClusterScopedAssignment(Params params, Collection<DiscoveryNode> candidateNodes, ClusterState clusterState) {
         DiscoveryNode discoveryNode = selectLeastLoadedNode(clusterState, candidateNodes, DiscoveryNode::canContainData);
         if (discoveryNode == null) {
             return NO_NODE_FOUND;
@@ -110,8 +120,19 @@ protected DiscoveryNode selectLeastLoadedNode(
      * Checks the current cluster state for compatibility with the params
      * <p>
      * Throws an exception if the supplied params cannot be executed on the cluster in the current state.
+     * <p>
+     * Implement this method if your {@link #scope()} returns CLUSTER, or {@link #validateProject} if your {@link #scope()} returns PROJECT.
+     */
+    public void validateCluster(Params params, ClusterState clusterState) {}
+
+    /**
+     * Checks the current project state for compatibility with the params
+     * <p>
+     * Throws an exception if the supplied params cannot be executed on the cluster in the current state.
+     * <p>
+     * Implement this method if your {@link #scope()} returns PROJECT, or {@link #validateCluster} if your {@link #scope()} returns CLUSTER.
      */
-    public void validate(Params params, ClusterState clusterState, @Nullable ProjectId projectId) {}
+    public void validateProject(Params params, ProjectState projectState) {}
 
     /**
      * Creates a AllocatedPersistentTask for communicating with task manager

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

@@ -20,7 +20,6 @@
 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;
@@ -1088,11 +1087,10 @@ public Scope scope() {
                 }
 
                 @Override
-                public Assignment getAssignment(
+                public Assignment getClusterScopedAssignment(
                     P params,
                     Collection<DiscoveryNode> candidateNodes,
-                    ClusterState clusterState,
-                    ProjectId projectId
+                    ClusterState clusterState
                 ) {
                     return fn.apply(params, candidateNodes, clusterState);
                 }

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

@@ -24,8 +24,8 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.client.internal.ElasticsearchClient;
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.ProjectState;
 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;
@@ -327,17 +327,25 @@ public static void setNonClusterStateCondition(boolean nonClusterStateCondition)
         }
 
         @Override
-        public Assignment getAssignment(
+        public Assignment getProjectScopedAssignment(
             TestParams params,
             Collection<DiscoveryNode> candidateNodes,
-            ClusterState clusterState,
-            ProjectId projectId
+            ProjectState projectState
+        ) {
+            return getClusterScopedAssignment(params, candidateNodes, projectState.cluster());
+        }
+
+        @Override
+        public Assignment getClusterScopedAssignment(
+            TestParams params,
+            Collection<DiscoveryNode> candidateNodes,
+            ClusterState clusterState
         ) {
             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, projectId);
+                return super.getClusterScopedAssignment(params, candidateNodes, clusterState);
             } else {
                 DiscoveryNode executorNode = selectLeastLoadedNode(
                     clusterState,

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

@@ -28,11 +28,10 @@
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.client.internal.RemoteClusterClient;
-import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.ProjectState;
 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;
@@ -44,7 +43,6 @@
 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;
@@ -120,8 +118,8 @@ public ShardFollowTasksExecutor(Client client, ThreadPool threadPool, ClusterSer
     }
 
     @Override
-    public void validate(ShardFollowTask params, ClusterState clusterState, @Nullable ProjectId projectId) {
-        final IndexRoutingTable routingTable = clusterState.getRoutingTable().index(params.getFollowShardId().getIndex());
+    public void validateProject(ShardFollowTask params, ProjectState projectState) {
+        final IndexRoutingTable routingTable = projectState.cluster().getRoutingTable().index(params.getFollowShardId().getIndex());
         final ShardRouting primaryShard = routingTable.shard(params.getFollowShardId().id()).primaryShard();
         if (primaryShard.active() == false) {
             throw new IllegalArgumentException("The primary shard of a follower index " + primaryShard + " is not active");
@@ -131,14 +129,13 @@ public void validate(ShardFollowTask params, ClusterState clusterState, @Nullabl
     private static final Assignment NO_ASSIGNMENT = new Assignment(null, "no nodes found with data and remote cluster client roles");
 
     @Override
-    public Assignment getAssignment(
+    public Assignment getProjectScopedAssignment(
         final ShardFollowTask params,
         final Collection<DiscoveryNode> candidateNodes,
-        final ClusterState clusterState,
-        @Nullable final ProjectId projectId
+        final ProjectState projectState
     ) {
         final DiscoveryNode node = selectLeastLoadedNode(
-            clusterState,
+            projectState.cluster(),
             candidateNodes,
             ((Predicate<DiscoveryNode>) DiscoveryNode::canContainData).and(DiscoveryNode::isRemoteClusterClient)
         );

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;
@@ -90,11 +91,10 @@ private void runAssignmentTest(
             nodesBuilder.add(newNode(otherNodesRolesSupplier.get()));
         }
         clusterStateBuilder.nodes(nodesBuilder);
-        final Assignment assignment = executor.getAssignment(
+        final Assignment assignment = executor.getProjectScopedAssignment(
             mock(ShardFollowTask.class),
             clusterStateBuilder.nodes().getAllNodes(),
-            clusterStateBuilder.build(),
-            null
+            clusterStateBuilder.build().projectState(ProjectId.DEFAULT)
         );
         consumer.accept(theSpecial, assignment);
     }

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

@@ -22,15 +22,14 @@
 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.ProjectState;
 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;
@@ -118,21 +117,20 @@ protected AllocatedPersistentTask createTask(
     }
 
     @Override
-    public void validate(DownsampleShardTaskParams params, ClusterState clusterState, @Nullable ProjectId projectId) {
+    public void validateProject(DownsampleShardTaskParams params, ProjectState projectState) {
         // 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);
+        var indexShardRouting = findShardRoutingTable(params.shardId(), projectState.cluster());
         if (indexShardRouting == null) {
             throw new ShardNotFoundException(params.shardId());
         }
     }
 
     @Override
-    public PersistentTasksCustomMetadata.Assignment getAssignment(
+    public PersistentTasksCustomMetadata.Assignment getProjectScopedAssignment(
         final DownsampleShardTaskParams params,
         final Collection<DiscoveryNode> candidateNodes,
-        final ClusterState clusterState,
-        @Nullable final ProjectId projectId
+        final ProjectState projectState
     ) {
         // 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
@@ -142,9 +140,9 @@ public PersistentTasksCustomMetadata.Assignment getAssignment(
         // If during re-assignment the source index was deleted, then we need to break out.
         // Returning NO_NODE_FOUND just keeps the persistent task until the source index appears again (which would never happen)
         // So let's return a node and then in the node operation we would just fail and stop this persistent task
-        var indexShardRouting = findShardRoutingTable(shardId, clusterState);
+        var indexShardRouting = findShardRoutingTable(shardId, projectState.cluster());
         if (indexShardRouting == null) {
-            var node = selectLeastLoadedNode(clusterState, candidateNodes, DiscoveryNode::canContainData);
+            var node = selectLeastLoadedNode(projectState.cluster(), candidateNodes, DiscoveryNode::canContainData);
             return new PersistentTasksCustomMetadata.Assignment(node.getId(), "a node to fail and stop this persistent task");
         }