Fix race condition in Feature Migration Status API (#80572)

Prior to this commit, there is a race condition in the Feature Migration
Status API where the returned status can be `MIGRATION_NEEDED`, even if
a migration is already in progress (and therefore the returned value
should have been `IN_PROGRESS`). This commit adds a test for this case
which reliably fails without the fix, and fixes the bug.

The fix is straightforward: While we already examine the persistent task
metadata to determine progress, the part of that metadata that we
examined did was not updated until the task's been running for a bit.
However, if we check for the *existence* of the task metadata, that is
guaranteed to be in the cluster state by the time the request to start the
migration completes (and is removed immediately after the task finishes
- that's why we have separate metadata for the migration results instead
of just using the task state).

Author: gwbrown

modules/reindex/src/internalClusterTest/java/org/elasticsearch/migration/FeatureMigrationIT.java

@@ -85,6 +85,43 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
         return plugins;
     }
 
+    public void testStartMigrationAndImmediatelyCheckStatus() throws Exception {
+        createSystemIndexForDescriptor(INTERNAL_MANAGED);
+        createSystemIndexForDescriptor(INTERNAL_UNMANAGED);
+        createSystemIndexForDescriptor(EXTERNAL_MANAGED);
+        createSystemIndexForDescriptor(EXTERNAL_UNMANAGED);
+
+        TestPlugin.preMigrationHook.set((state) -> Collections.emptyMap());
+        TestPlugin.postMigrationHook.set((state, metadata) -> {});
+
+        ensureGreen();
+
+        PostFeatureUpgradeRequest migrationRequest = new PostFeatureUpgradeRequest();
+        GetFeatureUpgradeStatusRequest getStatusRequest = new GetFeatureUpgradeStatusRequest();
+
+        // Start the migration and *immediately* request the status. We're trying to detect a race condition with this test, so we need to
+        // do this as fast as possible, but not before the request to start the migration completes.
+        PostFeatureUpgradeResponse migrationResponse = client().execute(PostFeatureUpgradeAction.INSTANCE, migrationRequest).get();
+        GetFeatureUpgradeStatusResponse statusResponse = client().execute(GetFeatureUpgradeStatusAction.INSTANCE, getStatusRequest).get();
+
+        // Make sure we actually started the migration
+        final Set<String> migratingFeatures = migrationResponse.getFeatures()
+            .stream()
+            .map(PostFeatureUpgradeResponse.Feature::getFeatureName)
+            .collect(Collectors.toSet());
+        assertThat(migratingFeatures, hasItem(FEATURE_NAME));
+
+        // We should see that the migration is in progress even though we just started the migration.
+        assertThat(statusResponse.getUpgradeStatus(), equalTo(GetFeatureUpgradeStatusResponse.UpgradeStatus.IN_PROGRESS));
+
+        // Now wait for the migration to finish (otherwise the test infra explodes)
+        assertBusy(() -> {
+            GetFeatureUpgradeStatusResponse statusResp = client().execute(GetFeatureUpgradeStatusAction.INSTANCE, getStatusRequest).get();
+            logger.info(Strings.toString(statusResp));
+            assertThat(statusResp.getUpgradeStatus(), equalTo(GetFeatureUpgradeStatusResponse.UpgradeStatus.NO_MIGRATION_NEEDED));
+        });
+    }
+
     public void testMigrateInternalManagedSystemIndex() throws Exception {
         createSystemIndexForDescriptor(INTERNAL_MANAGED);
         createSystemIndexForDescriptor(INTERNAL_UNMANAGED);

server/src/main/java/org/elasticsearch/action/admin/cluster/migration/TransportGetFeatureUpgradeStatusAction.java

@@ -20,17 +20,20 @@
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.indices.SystemIndices;
 import org.elasticsearch.persistent.PersistentTasksCustomMetadata;
+import org.elasticsearch.persistent.PersistentTasksService;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.upgrades.FeatureMigrationResults;
 import org.elasticsearch.upgrades.SingleFeatureMigrationResult;
+import org.elasticsearch.upgrades.SystemIndexMigrationTaskParams;
 import org.elasticsearch.upgrades.SystemIndexMigrationTaskState;
 
 import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.elasticsearch.action.admin.cluster.migration.GetFeatureUpgradeStatusResponse.UpgradeStatus.ERROR;
 import static org.elasticsearch.action.admin.cluster.migration.GetFeatureUpgradeStatusResponse.UpgradeStatus.IN_PROGRESS;
@@ -51,6 +54,7 @@ public class TransportGetFeatureUpgradeStatusAction extends TransportMasterNodeA
     public static final Version NO_UPGRADE_REQUIRED_VERSION = Version.V_8_0_0;
 
     private final SystemIndices systemIndices;
+    PersistentTasksService persistentTasksService;
 
     @Inject
     public TransportGetFeatureUpgradeStatusAction(
@@ -59,6 +63,7 @@ public TransportGetFeatureUpgradeStatusAction(
         ActionFilters actionFilters,
         ClusterService clusterService,
         IndexNameExpressionResolver indexNameExpressionResolver,
+        PersistentTasksService persistentTasksService,
         SystemIndices systemIndices
     ) {
         super(
@@ -73,6 +78,7 @@ public TransportGetFeatureUpgradeStatusAction(
             ThreadPool.Names.SAME
         );
         this.systemIndices = systemIndices;
+        this.persistentTasksService = persistentTasksService;
     }
 
     @Override
@@ -90,24 +96,26 @@ protected void masterOperation(
             .map(feature -> getFeatureUpgradeStatus(state, feature))
             .collect(Collectors.toList());
 
-        GetFeatureUpgradeStatusResponse.UpgradeStatus status = features.stream()
-            .map(GetFeatureUpgradeStatusResponse.FeatureUpgradeStatus::getUpgradeStatus)
-            .reduce(GetFeatureUpgradeStatusResponse.UpgradeStatus::combine)
-            .orElseGet(() -> {
-                assert false : "get feature statuses API doesn't have any features";
-                return NO_MIGRATION_NEEDED;
-            });
+        boolean migrationTaskExists = PersistentTasksCustomMetadata.getTaskWithId(state, SYSTEM_INDEX_UPGRADE_TASK_NAME) != null;
+        GetFeatureUpgradeStatusResponse.UpgradeStatus initalStatus = migrationTaskExists ? IN_PROGRESS : NO_MIGRATION_NEEDED;
+
+        GetFeatureUpgradeStatusResponse.UpgradeStatus status = Stream.concat(
+            Stream.of(initalStatus),
+            features.stream().map(GetFeatureUpgradeStatusResponse.FeatureUpgradeStatus::getUpgradeStatus)
+        ).reduce(GetFeatureUpgradeStatusResponse.UpgradeStatus::combine).orElseGet(() -> {
+            assert false : "get feature statuses API doesn't have any features";
+            return NO_MIGRATION_NEEDED;
+        });
 
         listener.onResponse(new GetFeatureUpgradeStatusResponse(features, status));
     }
 
     static GetFeatureUpgradeStatusResponse.FeatureUpgradeStatus getFeatureUpgradeStatus(ClusterState state, SystemIndices.Feature feature) {
         String featureName = feature.getName();
 
-        final String currentFeature = Optional.ofNullable(
-            state.metadata().<PersistentTasksCustomMetadata>custom(PersistentTasksCustomMetadata.TYPE)
-        )
-            .map(tasksMetdata -> tasksMetdata.getTask(SYSTEM_INDEX_UPGRADE_TASK_NAME))
+        PersistentTasksCustomMetadata.PersistentTask<SystemIndexMigrationTaskParams> migrationTask = PersistentTasksCustomMetadata
+            .getTaskWithId(state, SYSTEM_INDEX_UPGRADE_TASK_NAME);
+        final String currentFeature = Optional.ofNullable(migrationTask)
             .map(task -> task.getState())
             .map(taskState -> ((SystemIndexMigrationTaskState) taskState).getCurrentFeature())
             .orElse(null);