Specify generic types and check project removal

Author: nielsbauman

server/src/internalClusterTest/java/org/elasticsearch/persistent/PersistentTasksExecutorIT.java

@@ -290,7 +290,7 @@ public void testPersistentActionStatusUpdate() throws Exception {
             );
 
             int finalI = i;
-            WaitForPersistentTaskFuture<?> future1 = new WaitForPersistentTaskFuture<>();
+            WaitForPersistentTaskFuture<TestParams> future1 = new WaitForPersistentTaskFuture<>();
             waitForPersistentTaskCondition(
                 persistentTasksService,
                 taskId,
@@ -304,7 +304,7 @@ public void testPersistentActionStatusUpdate() throws Exception {
             assertThat(future1.get().getId(), equalTo(taskId));
         }
 
-        WaitForPersistentTaskFuture<?> future1 = new WaitForPersistentTaskFuture<>();
+        WaitForPersistentTaskFuture<TestParams> future1 = new WaitForPersistentTaskFuture<>();
         waitForPersistentTaskCondition(persistentTasksService, taskId, task -> false, TimeValue.timeValueMillis(10), future1);
 
         assertFutureThrows(future1, IllegalStateException.class, "timed out after 10ms");
@@ -318,7 +318,7 @@ public void testPersistentActionStatusUpdate() throws Exception {
         );
 
         // Wait for the task to disappear
-        WaitForPersistentTaskFuture<?> future2 = new WaitForPersistentTaskFuture<>();
+        WaitForPersistentTaskFuture<TestParams> future2 = new WaitForPersistentTaskFuture<>();
         waitForPersistentTaskCondition(persistentTasksService, taskId, Objects::isNull, TimeValue.timeValueSeconds(10), future2);
 
         logger.info("Completing the running task");
@@ -522,9 +522,9 @@ public void testAbortLocally() throws Exception {
     private void waitForPersistentTaskCondition(
         PersistentTasksService persistentTasksService,
         String taskId,
-        Predicate<PersistentTask<?>> predicate,
+        Predicate<PersistentTask<TestParams>> predicate,
         @Nullable TimeValue timeout,
-        WaitForPersistentTaskListener<?> listener
+        WaitForPersistentTaskListener<TestParams> listener
     ) throws Exception {
         if (scope == PersistentTasksExecutor.Scope.CLUSTER) {
             @FixForMultiProject(description = "can be replaced if PersistentTasksService supports waiting for cluster task conditions")

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

@@ -105,10 +105,10 @@ public long getAllocationId() {
      * @param timeout a timeout for waiting
      * @param listener the callback listener
      */
-    public void waitForPersistentTask(
-        final Predicate<PersistentTasksCustomMetadata.PersistentTask<?>> predicate,
+    public <P extends PersistentTaskParams> void waitForPersistentTask(
+        final Predicate<PersistentTasksCustomMetadata.PersistentTask<P>> predicate,
         final @Nullable TimeValue timeout,
-        final PersistentTasksService.WaitForPersistentTaskListener<?> listener
+        final PersistentTasksService.WaitForPersistentTaskListener<P> listener
     ) {
         persistentTasksService.waitForPersistentTaskCondition(persistentTaskId, predicate, timeout, listener);
     }

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

@@ -173,11 +173,11 @@ private <Req extends ActionRequest, Resp extends PersistentTaskResponse> void ex
      * @param listener the callback listener
      */
     @Deprecated(forRemoval = true)
-    public void waitForPersistentTaskCondition(
+    public <P extends PersistentTaskParams> void waitForPersistentTaskCondition(
         final String taskId,
-        final Predicate<PersistentTask<?>> predicate,
+        final Predicate<PersistentTask<P>> predicate,
         final @Nullable TimeValue timeout,
-        final WaitForPersistentTaskListener<?> listener
+        final WaitForPersistentTaskListener<P> listener
     ) {
         final var projectId = clusterService.state().metadata().getProject().id();
         waitForPersistentTaskCondition(projectId, taskId, predicate, timeout, listener);
@@ -192,17 +192,19 @@ public void waitForPersistentTaskCondition(
      * @param timeout a timeout for waiting
      * @param listener the callback listener
      */
-    public void waitForPersistentTaskCondition(
+    public <P extends PersistentTaskParams> void waitForPersistentTaskCondition(
         final ProjectId projectId,
         final String taskId,
-        final Predicate<PersistentTask<?>> predicate,
+        final Predicate<PersistentTask<P>> predicate,
         final @Nullable TimeValue timeout,
-        final WaitForPersistentTaskListener<?> listener
+        final WaitForPersistentTaskListener<P> listener
     ) {
         ClusterStateObserver.waitForState(clusterService, threadPool.getThreadContext(), new ClusterStateObserver.Listener() {
             @Override
             public void onNewClusterState(ClusterState state) {
-                listener.onResponse(PersistentTasksCustomMetadata.getTaskWithId(state.metadata().getProject(projectId), taskId));
+                final var project = state.metadata().projects().get(projectId);
+                final PersistentTask<P> task = project == null ? null : PersistentTasksCustomMetadata.getTaskWithId(project, taskId);
+                listener.onResponse(task);
             }
 
             @Override
@@ -214,13 +216,11 @@ public void onClusterServiceClose() {
             public void onTimeout(TimeValue timeout) {
                 listener.onTimeout(timeout);
             }
-        },
-            clusterState -> predicate.test(
-                PersistentTasksCustomMetadata.getTaskWithId(clusterState.metadata().getProject(projectId), taskId)
-            ),
-            timeout,
-            logger
-        );
+        }, clusterState -> {
+            final var project = clusterState.metadata().projects().get(projectId);
+            final PersistentTask<P> task = project == null ? null : PersistentTasksCustomMetadata.getTaskWithId(project, taskId);
+            return predicate.test(task);
+        }, timeout, logger);
     }
 
     // visible for testing

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

@@ -66,7 +66,6 @@
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.injection.guice.Inject;
-import org.elasticsearch.persistent.PersistentTaskParams;
 import org.elasticsearch.persistent.PersistentTasksCustomMetadata;
 import org.elasticsearch.persistent.PersistentTasksService;
 import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;
@@ -505,18 +504,18 @@ private void performShardDownsampling(
                 dimensionFields,
                 shardId
             );
-            Predicate<PersistentTasksCustomMetadata.PersistentTask<?>> predicate = runningTask -> {
+            Predicate<PersistentTasksCustomMetadata.PersistentTask<DownsampleShardTaskParams>> predicate = runningTask -> {
                 if (runningTask == null) {
                     // NOTE: don't need to wait if the persistent task completed and was removed
                     return true;
                 }
                 DownsampleShardPersistentTaskState runningPersistentTaskState = (DownsampleShardPersistentTaskState) runningTask.getState();
                 return runningPersistentTaskState != null && runningPersistentTaskState.done();
             };
-            var taskListener = new PersistentTasksService.WaitForPersistentTaskListener<>() {
+            var taskListener = new PersistentTasksService.WaitForPersistentTaskListener<DownsampleShardTaskParams>() {
 
                 @Override
-                public void onResponse(PersistentTasksCustomMetadata.PersistentTask<PersistentTaskParams> persistentTask) {
+                public void onResponse(PersistentTasksCustomMetadata.PersistentTask<DownsampleShardTaskParams> persistentTask) {
                     if (persistentTask != null) {
                         var runningPersistentTaskState = (DownsampleShardPersistentTaskState) persistentTask.getState();
                         if (runningPersistentTaskState != null) {

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportOpenJobAction.java

@@ -357,14 +357,14 @@ public void onFailure(Exception e) {
      * Important: the methods of this class must NOT throw exceptions.  If they did then the callers
      * of endpoints waiting for a condition tested by this predicate would never get a response.
      */
-    private static class JobPredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<?>> {
+    private static class JobPredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<OpenJobAction.JobParams>> {
 
         private volatile Exception exception;
         private volatile String node = "";
         private volatile boolean shouldCancel;
 
         @Override
-        public boolean test(PersistentTasksCustomMetadata.PersistentTask<?> persistentTask) {
+        public boolean test(PersistentTasksCustomMetadata.PersistentTask<OpenJobAction.JobParams> persistentTask) {
             JobState jobState = JobState.CLOSED;
             String reason = null;
             if (persistentTask != null) {
@@ -381,7 +381,7 @@ public boolean test(PersistentTasksCustomMetadata.PersistentTask<?> persistentTa
 
                 // This logic is only appropriate when opening a job, not when reallocating following a failure,
                 // and this is why this class must only be used when opening a job
-                OpenJobAction.JobParams params = (OpenJobAction.JobParams) persistentTask.getParams();
+                OpenJobAction.JobParams params = persistentTask.getParams();
                 Optional<ElasticsearchException> assignmentException = checkAssignmentState(assignment, params.getJobId(), logger);
                 if (assignmentException.isPresent()) {
                     exception = assignmentException.get();

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportStartDataFrameAnalyticsAction.java

@@ -36,7 +36,6 @@
 import org.elasticsearch.license.LicenseUtils;
 import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.persistent.AllocatedPersistentTask;
-import org.elasticsearch.persistent.PersistentTaskParams;
 import org.elasticsearch.persistent.PersistentTaskState;
 import org.elasticsearch.persistent.PersistentTasksCustomMetadata;
 import org.elasticsearch.persistent.PersistentTasksService;
@@ -458,11 +457,10 @@ private void waitForAnalyticsStarted(
             task.getId(),
             predicate,
             timeout,
-
-            new PersistentTasksService.WaitForPersistentTaskListener<PersistentTaskParams>() {
+            new PersistentTasksService.WaitForPersistentTaskListener<>() {
 
                 @Override
-                public void onResponse(PersistentTasksCustomMetadata.PersistentTask<PersistentTaskParams> persistentTask) {
+                public void onResponse(PersistentTasksCustomMetadata.PersistentTask<TaskParams> persistentTask) {
                     if (predicate.exception != null) {
                         // We want to return to the caller without leaving an unassigned persistent task, to match
                         // what would have happened if the error had been detected in the "fast fail" validation
@@ -530,14 +528,14 @@ private StartContext(DataFrameAnalyticsConfig config, List<PhaseProgress> progre
      * Important: the methods of this class must NOT throw exceptions.  If they did then the callers
      * of endpoints waiting for a condition tested by this predicate would never get a response.
      */
-    private static class AnalyticsPredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<?>> {
+    private static class AnalyticsPredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<TaskParams>> {
 
         private volatile Exception exception;
         private volatile String node = "";
         private volatile String assignmentExplanation;
 
         @Override
-        public boolean test(PersistentTasksCustomMetadata.PersistentTask<?> persistentTask) {
+        public boolean test(PersistentTasksCustomMetadata.PersistentTask<TaskParams> persistentTask) {
             if (persistentTask == null) {
                 return false;
             }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportStartDatafeedAction.java

@@ -369,7 +369,7 @@ private void waitForDatafeedStarted(
             taskId,
             predicate,
             params.getTimeout(),
-            new PersistentTasksService.WaitForPersistentTaskListener<StartDatafeedAction.DatafeedParams>() {
+            new PersistentTasksService.WaitForPersistentTaskListener<>() {
                 @Override
                 public void onResponse(PersistentTasksCustomMetadata.PersistentTask<StartDatafeedAction.DatafeedParams> persistentTask) {
                     if (predicate.exception != null) {
@@ -711,13 +711,15 @@ public GetDatafeedRunningStateAction.Response.RunningState getRunningState() {
      * Important: the methods of this class must NOT throw exceptions.  If they did then the callers
      * of endpoints waiting for a condition tested by this predicate would never get a response.
      */
-    private static class DatafeedPredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<?>> {
+    private static class DatafeedPredicate
+        implements
+            Predicate<PersistentTasksCustomMetadata.PersistentTask<StartDatafeedAction.DatafeedParams>> {
 
         private volatile Exception exception;
         private volatile String node = "";
 
         @Override
-        public boolean test(PersistentTasksCustomMetadata.PersistentTask<?> persistentTask) {
+        public boolean test(PersistentTasksCustomMetadata.PersistentTask<StartDatafeedAction.DatafeedParams> persistentTask) {
             if (persistentTask == null) {
                 return false;
             }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportUpgradeJobModelSnapshotAction.java

@@ -247,7 +247,7 @@ private void waitForJobStarted(
             taskId,
             predicate,
             request.getTimeout(),
-            new PersistentTasksService.WaitForPersistentTaskListener<SnapshotUpgradeTaskParams>() {
+            new PersistentTasksService.WaitForPersistentTaskListener<>() {
                 @Override
                 public void onResponse(PersistentTask<SnapshotUpgradeTaskParams> persistentTask) {
                     if (predicate.getException() != null) {

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/snapshot/upgrader/SnapshotUpgradePredicate.java

@@ -20,7 +20,7 @@
 
 import static org.elasticsearch.xpack.ml.job.task.OpenJobPersistentTasksExecutor.checkAssignmentState;
 
-public class SnapshotUpgradePredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<?>> {
+public class SnapshotUpgradePredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<SnapshotUpgradeTaskParams>> {
     private final boolean waitForCompletion;
     private final Logger logger;
     private volatile Exception exception;
@@ -50,7 +50,7 @@ public boolean isCompleted() {
     }
 
     @Override
-    public boolean test(PersistentTasksCustomMetadata.PersistentTask<?> persistentTask) {
+    public boolean test(PersistentTasksCustomMetadata.PersistentTask<SnapshotUpgradeTaskParams> persistentTask) {
         // Persistent task being null means it has been removed from state, and is now complete
         if (persistentTask == null) {
             isCompleted = true;
@@ -64,7 +64,7 @@ public boolean test(PersistentTasksCustomMetadata.PersistentTask<?> persistentTa
         PersistentTasksCustomMetadata.Assignment assignment = persistentTask.getAssignment();
         // This logic is only appropriate when opening a job, not when reallocating following a failure,
         // and this is why this class must only be used when opening a job
-        SnapshotUpgradeTaskParams params = (SnapshotUpgradeTaskParams) persistentTask.getParams();
+        SnapshotUpgradeTaskParams params = persistentTask.getParams();
         Optional<ElasticsearchException> assignmentException = checkAssignmentState(assignment, params.getJobId(), logger);
         if (assignmentException.isPresent()) {
             exception = assignmentException.get();

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportStartTransformAction.java

@@ -339,7 +339,7 @@ private void waitForTransformTaskStarted(
             taskId,
             predicate,
             timeout,
-            new PersistentTasksService.WaitForPersistentTaskListener<TransformTaskParams>() {
+            new PersistentTasksService.WaitForPersistentTaskListener<>() {
                 @Override
                 public void onResponse(PersistentTasksCustomMetadata.PersistentTask<TransformTaskParams> persistentTask) {
                     if (predicate.exception != null) {
@@ -374,12 +374,12 @@ public void onTimeout(TimeValue timeout) {
      * Important: the methods of this class must NOT throw exceptions.  If they did then the callers
      * of endpoints waiting for a condition tested by this predicate would never get a response.
      */
-    private static class TransformPredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<?>> {
+    private static class TransformPredicate implements Predicate<PersistentTasksCustomMetadata.PersistentTask<TransformTaskParams>> {
 
         private volatile Exception exception;
 
         @Override
-        public boolean test(PersistentTasksCustomMetadata.PersistentTask<?> persistentTask) {
+        public boolean test(PersistentTasksCustomMetadata.PersistentTask<TransformTaskParams> persistentTask) {
             if (persistentTask == null) {
                 return false;
             }