Fix race condition when storing Feature migration callback metadata (#80049)

This PR fixes a race condition where the callback metadata was stored in
a callback which probably won't run right away, and just immediately proceeds,
meaning we move on before the callback's been called.

The existing tests didn't catch this because all the test callbacks are actually
synchronous. I'd like to add a test that catches this, but doing so may be complex
and I'd like to get the fix in sooner rather than later.

Author: gwbrown

server/src/main/java/org/elasticsearch/upgrades/SystemIndexMigrator.java

@@ -331,29 +331,30 @@ private void prepareNextIndex(ClusterState clusterState, Consumer<ClusterState>
             migrationInfo.getFeatureName(),
             migrationInfo.getNextIndexName()
         );
-        final AtomicReference<Map<String, Object>> updatedTaskStateFeatureMetadata = new AtomicReference<>();
         if (migrationInfo.getFeatureName().equals(lastFeatureName) == false) {
             // And then invoke the pre-migration hook for the next one.
-            migrationInfo.prepareForIndicesMigration(
-                clusterService,
-                baseClient,
-                ActionListener.wrap(updatedTaskStateFeatureMetadata::set, this::markAsFailed)
-            );
+            migrationInfo.prepareForIndicesMigration(clusterService, baseClient, ActionListener.wrap(newMetadata -> {
+                currentFeatureCallbackMetadata.set(newMetadata);
+                updateTaskState(migrationInfo, listener, newMetadata);
+            }, this::markAsFailed));
         } else {
             // Otherwise, just re-use what we already have.
-            updatedTaskStateFeatureMetadata.set(currentFeatureCallbackMetadata.get());
+            updateTaskState(migrationInfo, listener, currentFeatureCallbackMetadata.get());
         }
+    }
+
+    private void updateTaskState(SystemIndexMigrationInfo migrationInfo, Consumer<ClusterState> listener, Map<String, Object> metadata) {
         final SystemIndexMigrationTaskState newTaskState = new SystemIndexMigrationTaskState(
             migrationInfo.getCurrentIndexName(),
             migrationInfo.getFeatureName(),
-            updatedTaskStateFeatureMetadata.get()
+            metadata
         );
         logger.debug("updating task state to [{}]", Strings.toString(newTaskState));
-        currentFeatureCallbackMetadata.set(updatedTaskStateFeatureMetadata.get());
+        currentFeatureCallbackMetadata.set(metadata);
         updatePersistentTaskState(newTaskState, ActionListener.wrap(task -> {
             assert newTaskState.equals(task.getState()) : "task state returned by update method did not match submitted task state";
             logger.debug("new task state [{}] accepted", Strings.toString(newTaskState));
-            listener.accept(clusterState);
+            listener.accept(clusterService.state());
         }, this::markAsFailed));
     }