Wait for relocations during the first mixed round in SnapshotBasedRec…

…overyIT (#81384)

This commits adds additional logging in order to debug the test failures

Closes #80939

qa/rolling-upgrade/src/test/java/org/elasticsearch/upgrades/SnapshotBasedRecoveryIT.java

@@ -10,21 +10,22 @@
 
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpPost;
+import org.apache.http.util.EntityUtils;
 import org.elasticsearch.Version;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.repositories.blobstore.BlobStoreRepository;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -39,7 +40,6 @@
 import static org.hamcrest.Matchers.notNullValue;
 
 public class SnapshotBasedRecoveryIT extends AbstractRollingTestCase {
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/80939")
     public void testSnapshotBasedRecovery() throws Exception {
         final String indexName = "snapshot_based_recovery";
         final String repositoryName = "snapshot_based_recovery_repo";
@@ -70,10 +70,14 @@ public void testSnapshotBasedRecovery() throws Exception {
             }
             case MIXED, UPGRADED -> {
                 if (FIRST_MIXED_ROUND) {
-                    String upgradedNodeId = getUpgradedNodeId();
-
-                    if (upgradedNodeId != null) {
+                    List<String> upgradedNodeIds = getUpgradedNodeIds();
+                    // It's possible that the test simply does a rolling-restart, i.e. it "upgrades" to
+                    // the same version. In that case we proceed without excluding any node
+                    if (upgradedNodeIds.isEmpty() == false) {
+                        assertThat(upgradedNodeIds.size(), is(equalTo(1)));
+                        String upgradedNodeId = upgradedNodeIds.get(0);
                         updateIndexSettings(indexName, Settings.builder().put("index.routing.allocation.exclude._id", upgradedNodeId));
+                        ensureGreen(indexName);
                     }
 
                     String primaryNodeId = getPrimaryNodeIdOfShard(indexName, 0);
@@ -84,6 +88,7 @@ public void testSnapshotBasedRecovery() throws Exception {
                     // That is an issue only for the first mixed round.
                     // In that case we exclude the upgraded node from the shard allocation and cancel the shard to force moving
                     // the primary to a node in the old version, this allows adding replicas in the first mixed round.
+                    logger.info("--> Primary node in first mixed round {} / {}", primaryNodeId, primaryNodeVersion);
                     if (primaryNodeVersion.after(UPGRADE_FROM_VERSION)) {
                         cancelShard(indexName, 0, primaryNodeId);
 
@@ -105,19 +110,19 @@ public void testSnapshotBasedRecovery() throws Exception {
         }
     }
 
-    @Nullable
-    private String getUpgradedNodeId() throws IOException {
+    private List<String> getUpgradedNodeIds() throws IOException {
         Request request = new Request(HttpGet.METHOD_NAME, "_nodes/_all");
         Response response = client().performRequest(request);
         Map<String, Object> responseMap = responseAsMap(response);
         Map<String, Map<String, Object>> nodes = extractValue(responseMap, "nodes");
+        List<String> upgradedNodes = new ArrayList<>();
         for (Map.Entry<String, Map<String, Object>> nodeInfoEntry : nodes.entrySet()) {
             Version nodeVersion = Version.fromString(extractValue(nodeInfoEntry.getValue(), "version"));
             if (nodeVersion.after(UPGRADE_FROM_VERSION)) {
-                return nodeInfoEntry.getKey();
+                upgradedNodes.add(nodeInfoEntry.getKey());
             }
         }
-        return null;
+        return upgradedNodes;
     }
 
     private Version getNodeVersion(String primaryNodeId) throws IOException {
@@ -173,9 +178,10 @@ private void cancelShard(String indexName, int shard, String nodeName) throws IO
             }
             builder.endObject();
 
-            Request request = new Request(HttpPost.METHOD_NAME, "/_cluster/reroute");
+            Request request = new Request(HttpPost.METHOD_NAME, "/_cluster/reroute?pretty");
             request.setJsonEntity(Strings.toString(builder));
             Response response = client().performRequest(request);
+            logger.info("--> Relocated primary to an older version {}", EntityUtils.toString(response.getEntity()));
             assertOK(response);
         }
     }