Trigger after shard listener before post recovery (elastic#134419)

Trigger this listener prior to transitioning to post recovery state so
that the shard is not yet writeable.

Author: Tim-Brooks

server/src/main/java/org/elasticsearch/index/shard/IndexEventListener.java

@@ -202,6 +202,13 @@ default void beforeIndexShardRecovery(IndexShard indexShard, IndexSettings index
         listener.onResponse(null);
     }
 
+    /**
+     * Called after the recover process is completed. The recovery state is DONE at this point. However, this is triggered prior to the
+     * index shard state transitions to either POST_RECOVERY to STARTED
+     *
+     * @param indexShard the shard that was recovered
+     * @param listener listener notified when this step completes
+     */
     default void afterIndexShardRecovery(IndexShard indexShard, ActionListener<Void> listener) {
         listener.onResponse(null);
     }

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -1900,6 +1900,7 @@ public void postRecovery(String reason, ActionListener<Void> listener) throws In
         SubscribableListener<Void> subscribableListener = new SubscribableListener<>();
         postRecoveryComplete = subscribableListener;
         final ActionListener<Void> finalListener = ActionListener.runBefore(listener, () -> subscribableListener.onResponse(null));
+
         try {
             // Some engine implementations try to acquire the engine reset write lock during refresh: in case something else is holding the
             // engine read lock at the same time then the refresh is a no-op for those engines. Here we acquire the engine reset write lock
@@ -1922,9 +1923,21 @@ public void postRecovery(String reason, ActionListener<Void> listener) throws In
                     throw new IndexShardStartedException(shardId);
                 }
                 recoveryState.setStage(RecoveryState.Stage.DONE);
-                changeState(IndexShardState.POST_RECOVERY, reason);
             }
-            indexEventListener.afterIndexShardRecovery(this, finalListener);
+
+            SubscribableListener.newForked(
+                (CheckedConsumer<ActionListener<Void>, Exception>) l -> indexEventListener.afterIndexShardRecovery(IndexShard.this, l)
+            ).andThenAccept(v -> {
+                synchronized (mutex) {
+                    if (state == IndexShardState.CLOSED) {
+                        throw new IndexShardClosedException(shardId);
+                    }
+                    if (state == IndexShardState.STARTED) {
+                        throw new IndexShardStartedException(shardId);
+                    }
+                    changeState(IndexShardState.POST_RECOVERY, reason);
+                }
+            }).addListener(finalListener);
         } catch (Exception e) {
             finalListener.onFailure(e);
         }

server/src/test/java/org/elasticsearch/indices/IndicesLifecycleListenerSingleNodeTests.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.indices;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeUtils;
@@ -24,18 +25,22 @@
 import org.elasticsearch.index.seqno.RetentionLeaseSyncer;
 import org.elasticsearch.index.shard.IndexEventListener;
 import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.index.shard.IndexShardState;
 import org.elasticsearch.index.shard.IndexShardTestCase;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.cluster.IndexRemovalReason;
 import org.elasticsearch.indices.recovery.RecoveryState;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 
 import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.LockSupport;
 
 import static java.util.Collections.emptySet;
 import static org.elasticsearch.indices.cluster.IndexRemovalReason.DELETED;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.hamcrest.Matchers.equalTo;
 
 public class IndicesLifecycleListenerSingleNodeTests extends ESSingleNodeTestCase {
 
@@ -106,7 +111,6 @@ public void afterIndexRemoved(Index index, IndexSettings indexSettings, IndexRem
                 assertEquals(9, counter.get());
                 counter.incrementAndGet();
             }
-
         };
         indicesService.removeIndex(idx, DELETED, "simon says", EsExecutors.DIRECT_EXECUTOR_SERVICE, ActionListener.noop());
         try {
@@ -136,4 +140,48 @@ public void afterIndexRemoved(Index index, IndexSettings indexSettings, IndexRem
         assertEquals(10, counter.get());
     }
 
+    public void testAfterRecoveryCallbackTriggeredWhileStillInRecoveryState() throws Throwable {
+        IndicesService indicesService = getInstanceFromNode(IndicesService.class);
+        assertAcked(client().admin().indices().prepareCreate("test").setSettings(indexSettings(1, 0)));
+        ensureGreen();
+        Index idx = resolveIndex("test");
+        IndexMetadata metadata = indicesService.indexService(idx).getMetadata();
+        ShardRouting shardRouting = indicesService.indexService(idx).getShard(0).routingEntry();
+        PlainActionFuture<Void> recoveryTriggered = new PlainActionFuture<>();
+        IndexEventListener recoveryListener = new IndexEventListener() {
+
+            @Override
+            public void afterIndexShardRecovery(IndexShard indexShard, ActionListener<Void> listener) {
+                // Pause to ensure we do not transition to post recovery until after the recovery is complete
+                LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(30));
+                assertThat(indexShard.recoveryState().getStage(), equalTo(RecoveryState.Stage.DONE));
+                assertThat(indexShard.state(), equalTo(IndexShardState.RECOVERING));
+                recoveryTriggered.onResponse(null);
+                listener.onResponse(null);
+            }
+        };
+        indicesService.removeIndex(idx, DELETED, "delete", EsExecutors.DIRECT_EXECUTOR_SERVICE, ActionListener.noop());
+        try {
+            IndexService index = indicesService.createIndex(metadata, Arrays.asList(recoveryListener), false);
+            idx = index.index();
+            ShardRouting newRouting = shardRouting;
+            String nodeId = newRouting.currentNodeId();
+            UnassignedInfo unassignedInfo = new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, "boom");
+            newRouting = newRouting.moveToUnassigned(unassignedInfo)
+                .updateUnassigned(unassignedInfo, RecoverySource.EmptyStoreRecoverySource.INSTANCE);
+            newRouting = ShardRoutingHelper.initialize(newRouting, nodeId);
+            IndexShard shard = index.createShard(newRouting, IndexShardTestCase.NOOP_GCP_SYNCER, RetentionLeaseSyncer.EMPTY);
+            IndexShardTestCase.updateRoutingEntry(shard, newRouting);
+            final DiscoveryNode localNode = DiscoveryNodeUtils.builder("foo").roles(emptySet()).build();
+            shard.markAsRecovering("store", new RecoveryState(newRouting, localNode, null));
+            IndexShardTestCase.recoverFromStore(shard);
+            assertBusy(() -> assertThat(shard.state(), equalTo(IndexShardState.POST_RECOVERY)));
+            newRouting = ShardRoutingHelper.moveToStarted(newRouting);
+            IndexShardTestCase.updateRoutingEntry(shard, newRouting);
+            recoveryTriggered.actionGet();
+            assertBusy(() -> assertThat(shard.state(), equalTo(IndexShardState.STARTED)));
+        } finally {
+            indicesService.removeIndex(idx, DELETED, "simon says", EsExecutors.DIRECT_EXECUTOR_SERVICE, ActionListener.noop());
+        }
+    }
 }