Make downsampling project-aware (#124000)

Allows downsampling to work on multiple projects.

Author: nielsbauman

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

@@ -122,9 +122,14 @@ public static PersistentTasksCustomMetadata fromXContent(XContentParser parser)
         return PERSISTENT_TASKS_PARSER.apply(parser, null).build();
     }
 
-    @SuppressWarnings("unchecked")
+    @Deprecated(forRemoval = true)
     public static <Params extends PersistentTaskParams> PersistentTask<Params> getTaskWithId(ClusterState clusterState, String taskId) {
-        PersistentTasksCustomMetadata tasks = get(clusterState.metadata().getProject());
+        return getTaskWithId(clusterState.metadata().getProject(), taskId);
+    }
+
+    @SuppressWarnings("unchecked")
+    public static <Params extends PersistentTaskParams> PersistentTask<Params> getTaskWithId(ProjectMetadata project, String taskId) {
+        PersistentTasksCustomMetadata tasks = get(project);
         if (tasks != null) {
             return (PersistentTask<Params>) tasks.getTask(taskId);
         }

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.client.internal.OriginSettingClient;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ClusterStateObserver;
+import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
@@ -171,16 +172,39 @@ private <Req extends ActionRequest, Resp extends PersistentTaskResponse> void ex
      * @param timeout a timeout for waiting
      * @param listener the callback listener
      */
+    @Deprecated(forRemoval = true)
     public void waitForPersistentTaskCondition(
         final String taskId,
         final Predicate<PersistentTask<?>> predicate,
         final @Nullable TimeValue timeout,
         final WaitForPersistentTaskListener<?> listener
+    ) {
+        final var projectId = clusterService.state().metadata().getProject().id();
+        waitForPersistentTaskCondition(projectId, taskId, predicate, timeout, listener);
+    }
+
+    /**
+     * Waits for a given persistent task to comply with a given predicate, then call back the listener accordingly.
+     *
+     * @param projectId the project ID
+     * @param taskId the persistent task id
+     * @param predicate the persistent task predicate to evaluate, must be able to handle {@code null} input which means either the project
+     *                  does not exist or persistent tasks for the project do not exist
+     * @param timeout a timeout for waiting
+     * @param listener the callback listener
+     */
+    public void waitForPersistentTaskCondition(
+        final ProjectId projectId,
+        final String taskId,
+        final Predicate<PersistentTask<?>> predicate,
+        final @Nullable TimeValue timeout,
+        final WaitForPersistentTaskListener<?> listener
     ) {
         ClusterStateObserver.waitForState(clusterService, threadPool.getThreadContext(), new ClusterStateObserver.Listener() {
             @Override
             public void onNewClusterState(ClusterState state) {
-                listener.onResponse(PersistentTasksCustomMetadata.getTaskWithId(state, taskId));
+                final var project = state.metadata().projects().get(projectId);
+                listener.onResponse(project == null ? null : PersistentTasksCustomMetadata.getTaskWithId(project, taskId));
             }
 
             @Override
@@ -192,7 +216,15 @@ public void onClusterServiceClose() {
             public void onTimeout(TimeValue timeout) {
                 listener.onTimeout(timeout);
             }
-        }, clusterState -> predicate.test(PersistentTasksCustomMetadata.getTaskWithId(clusterState, taskId)), timeout, logger);
+        }, clusterState -> {
+            final var project = clusterState.metadata().projects().get(projectId);
+            if (project == null) {
+                logger.debug("project [{}] not found while waiting for persistent task [{}] to pass predicate", projectId, taskId);
+                return predicate.test(null);
+            } else {
+                return predicate.test(PersistentTasksCustomMetadata.getTaskWithId(project, taskId));
+            }
+        }, timeout, logger);
     }
 
     // visible for testing

test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java

@@ -554,6 +554,16 @@ public static ClusterState getClusterStateWithDataStream(String dataStream, List
             .build();
     }
 
+    public static ClusterState getClusterStateWithDataStream(
+        ProjectId projectId,
+        String dataStream,
+        List<Tuple<Instant, Instant>> timeSlices
+    ) {
+        return ClusterState.builder(ClusterName.DEFAULT)
+            .putProjectMetadata(getProjectWithDataStream(projectId, dataStream, timeSlices))
+            .build();
+    }
+
     public static ProjectMetadata getProjectWithDataStream(
         ProjectId projectId,
         String dataStream,

x-pack/plugin/downsample/qa/rest/build.gradle

@@ -34,3 +34,15 @@ if (buildParams.inFipsJvm){
   // This test cluster is using a BASIC license and FIPS 140 mode is not supported in BASIC
   tasks.named("yamlRestTest").configure{enabled = false }
 }
+
+configurations {
+  basicRestSpecs {
+    attributes {
+      attribute(ArtifactTypeDefinition.ARTIFACT_TYPE_ATTRIBUTE, ArtifactTypeDefinition.DIRECTORY_TYPE)
+    }
+  }
+}
+
+artifacts {
+  basicRestSpecs(new File(projectDir, "src/yamlRestTest/resources/rest-api-spec/test"))
+}

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

@@ -116,7 +116,7 @@ protected AllocatedPersistentTask createTask(
     public void validate(DownsampleShardTaskParams params, ClusterState clusterState) {
         // 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 = clusterState.routingTable().shardRoutingTable(params.shardId().getIndexName(), params.shardId().id());
+        var indexShardRouting = findShardRoutingTable(params.shardId(), clusterState);
         if (indexShardRouting == null) {
             throw new ShardNotFoundException(params.shardId());
         }
@@ -178,11 +178,8 @@ private void delegate(final AllocatedPersistentTask task, final DownsampleShardT
     }
 
     private static IndexShardRoutingTable findShardRoutingTable(ShardId shardId, ClusterState clusterState) {
-        var indexRoutingTable = clusterState.routingTable().index(shardId.getIndexName());
-        if (indexRoutingTable != null) {
-            return indexRoutingTable.shard(shardId.getId());
-        }
-        return null;
+        var indexRoutingTable = clusterState.globalRoutingTable().indexRouting(clusterState.metadata(), shardId.getIndex());
+        return indexRoutingTable.map(routingTable -> routingTable.shard(shardId.getId())).orElse(null);
     }
 
     static void realNodeOperation(
@@ -327,6 +324,7 @@ protected void doExecute(Task t, Request request, ActionListener<ActionResponse.
                 realNodeOperation(client, indicesService, downsampleMetrics, request.task, request.params, request.lastDownsampleTsid);
                 listener.onResponse(ActionResponse.Empty.INSTANCE);
             }
+
         }
     }
 }