Put a fake allocation id on allocate stale primary command; remove it after recovery is done

Relates to elastic#33432

Author: Vladimir Dolzhenko

server/src/main/java/org/elasticsearch/cluster/routing/IndexRoutingTable.java

@@ -141,7 +141,8 @@ boolean validate(MetaData metaData) {
 
                 if (shardRouting.primary() && shardRouting.initializing() &&
                     shardRouting.recoverySource().getType() == RecoverySource.Type.EXISTING_STORE &&
-                    inSyncAllocationIds.contains(shardRouting.allocationId().getId()) == false)
+                    inSyncAllocationIds.contains(shardRouting.allocationId().getId()) == false &&
+                    inSyncAllocationIds.contains(RecoverySource.ExistingStoreRecoverySource.FORCED_ALLOCATION_ID) == false)
                     throw new IllegalStateException("a primary shard routing " + shardRouting + " is a primary that is recovering from " +
                         "a known allocation id but has no corresponding entry in the in-sync " +
                         "allocation set " + inSyncAllocationIds);

server/src/main/java/org/elasticsearch/cluster/routing/RecoverySource.java

@@ -132,6 +132,7 @@ public String toString() {
      * Recovery from an existing on-disk store
      */
     public static final class ExistingStoreRecoverySource extends RecoverySource {
+        public static final String FORCED_ALLOCATION_ID = "_forced_allocation";
         public static final ExistingStoreRecoverySource INSTANCE = new ExistingStoreRecoverySource(false);
         public static final ExistingStoreRecoverySource FORCE_STALE_PRIMARY_INSTANCE = new ExistingStoreRecoverySource(true);
 

server/src/main/java/org/elasticsearch/cluster/routing/allocation/IndexMetaDataUpdater.java

@@ -69,6 +69,13 @@ public void shardInitialized(ShardRouting unassignedShard, ShardRouting initiali
     @Override
     public void shardStarted(ShardRouting initializingShard, ShardRouting startedShard) {
         addAllocationId(startedShard);
+        if (startedShard.primary()
+            // started shard has to have null recoverySource; have to pick up recoverySource from its initializing state
+            && (initializingShard.recoverySource() instanceof RecoverySource.ExistingStoreRecoverySource
+            || initializingShard.recoverySource() instanceof RecoverySource.SnapshotRecoverySource)) {
+            Updates updates = changes(startedShard.shardId());
+            updates.removedAllocationIds.add(RecoverySource.ExistingStoreRecoverySource.FORCED_ALLOCATION_ID);
+        }
     }
 
     @Override
@@ -144,7 +151,8 @@ private IndexMetaData.Builder updateInSyncAllocations(RoutingTable newRoutingTab
             oldInSyncAllocationIds.contains(updates.initializedPrimary.allocationId().getId()) == false) {
             // we're not reusing an existing in-sync allocation id to initialize a primary, which means that we're either force-allocating
             // an empty or a stale primary (see AllocateEmptyPrimaryAllocationCommand or AllocateStalePrimaryAllocationCommand).
-            RecoverySource.Type recoverySourceType = updates.initializedPrimary.recoverySource().getType();
+            RecoverySource recoverySource = updates.initializedPrimary.recoverySource();
+            RecoverySource.Type recoverySourceType = recoverySource.getType();
             boolean emptyPrimary = recoverySourceType == RecoverySource.Type.EMPTY_STORE;
             assert updates.addedAllocationIds.isEmpty() : (emptyPrimary ? "empty" : "stale") +
                 " primary is not force-initialized in same allocation round where shards are started";
@@ -156,9 +164,12 @@ private IndexMetaData.Builder updateInSyncAllocations(RoutingTable newRoutingTab
                 // forcing an empty primary resets the in-sync allocations to the empty set (ShardRouting.allocatedPostIndexCreate)
                 indexMetaDataBuilder.putInSyncAllocationIds(shardId.id(), Collections.emptySet());
             } else {
+                assert recoverySource instanceof RecoverySource.ExistingStoreRecoverySource
+                    || recoverySource instanceof RecoverySource.SnapshotRecoverySource
+                    : recoverySource;
                 // forcing a stale primary resets the in-sync allocations to the singleton set with the stale id
                 indexMetaDataBuilder.putInSyncAllocationIds(shardId.id(),
-                    Collections.singleton(updates.initializedPrimary.allocationId().getId()));
+                    Collections.singleton(RecoverySource.ExistingStoreRecoverySource.FORCED_ALLOCATION_ID));
             }
         } else {
             // standard path for updating in-sync ids
@@ -291,7 +302,8 @@ void removeAllocationId(ShardRouting shardRouting) {
      * Add allocation id of this shard to the set of in-sync shard copies
      */
     private void addAllocationId(ShardRouting shardRouting) {
-        changes(shardRouting.shardId()).addedAllocationIds.add(shardRouting.allocationId().getId());
+        final Updates changes = changes(shardRouting.shardId());
+        changes.addedAllocationIds.add(shardRouting.allocationId().getId());
     }
 
     /**

server/src/test/java/org/elasticsearch/cluster/routing/PrimaryAllocationIT.java

@@ -204,6 +204,13 @@ public void testForceStaleReplicaToBePromotedToPrimary() throws Exception {
         }
         rerouteBuilder.get();
 
+        ClusterState state = client().admin().cluster().prepareState().get().getState();
+
+        Set<String> expectedAllocationIds = useStaleReplica
+            ? Collections.singleton(RecoverySource.ExistingStoreRecoverySource.FORCED_ALLOCATION_ID)
+            : Collections.emptySet();
+        assertEquals(expectedAllocationIds, state.metaData().index(idxName).inSyncAllocationIds(0));
+
         logger.info("--> check that the stale primary shard gets allocated and that documents are available");
         ensureYellow(idxName);
 
@@ -218,7 +225,8 @@ public void testForceStaleReplicaToBePromotedToPrimary() throws Exception {
         assertHitCount(client().prepareSearch(idxName).setSize(0).setQuery(matchAllQuery()).get(), useStaleReplica ? 1L : 0L);
 
         // allocation id of old primary was cleaned from the in-sync set
-        ClusterState state = client().admin().cluster().prepareState().get().getState();
+        state = client().admin().cluster().prepareState().get().getState();
+
         assertEquals(Collections.singleton(state.routingTable().index(idxName).shard(0).primary.allocationId().getId()),
             state.metaData().index(idxName).inSyncAllocationIds(0));