Add pre recovery hook

Author: tlrx

server/src/main/java/org/elasticsearch/index/CompositeIndexEventListener.java

@@ -257,4 +257,17 @@ public void onStoreClosed(ShardId shardId) {
             }
         }
     }
+
+    @Override
+    public void beforeIndexShardRecovery(final IndexShard indexShard, final IndexSettings indexSettings) {
+        for (IndexEventListener listener  : listeners) {
+            try {
+                listener.beforeIndexShardRecovery(indexShard, indexSettings);
+            } catch (Exception e) {
+                logger.warn(() -> new ParameterizedMessage("failed to invoke the listener before the shard recovery starts for {}",
+                    indexShard.shardId()), e);
+                throw e;
+            }
+        }
+    }
 }

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

@@ -167,4 +167,15 @@ default void onStoreCreated(ShardId shardId) {}
      * @param shardId the shard ID the store belongs to
      */
     default void onStoreClosed(ShardId shardId) {}
+
+    /**
+     * Called before the index shard starts to recover.
+     * Note: unlike all other methods in this class, this method is not called using the cluster state update thread. When this method is
+     * called the shard already transitioned to the RECOVERING state.
+     *
+     * @param indexShard    the shard that is about to recover
+     * @param indexSettings the shard's index settings
+     */
+    default void beforeIndexShardRecovery(IndexShard indexShard, IndexSettings indexSettings) {
+    }
 }

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

@@ -1313,6 +1313,13 @@ public void close(String reason, boolean flushEngine) throws IOException {
         }
     }
 
+    public void preRecovery() {
+        if (state != IndexShardState.RECOVERING) {
+            throw new IndexShardNotRecoveringException(shardId, state);
+        }
+        indexEventListener.beforeIndexShardRecovery(this, indexSettings);
+    }
+
     public void postRecovery(String reason) throws IndexShardStartedException, IndexShardRelocatedException, IndexShardClosedException {
         synchronized (postRecoveryMutex) {
             // we need to refresh again to expose all operations that were index until now. Otherwise

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

@@ -359,6 +359,7 @@ private ActionListener<Boolean> recoveryListener(IndexShard indexShard, ActionLi
      * Recovers the state of the shard from the store.
      */
     private void internalRecoverFromStore(IndexShard indexShard) throws IndexShardRecoveryException {
+        indexShard.preRecovery();
         final RecoveryState recoveryState = indexShard.recoveryState();
         final boolean indexShouldExists = recoveryState.getRecoverySource().getType() != RecoverySource.Type.EMPTY_STORE;
         indexShard.prepareForIndexRecovery();
@@ -449,6 +450,7 @@ private void addRecoveredFileDetails(SegmentInfos si, Store store, RecoveryState
     private void restore(IndexShard indexShard, Repository repository, SnapshotRecoverySource restoreSource,
                          ActionListener<Boolean> listener) {
         logger.debug("restoring from {} ...", indexShard.recoveryState().getRecoverySource());
+        indexShard.preRecovery();
         final RecoveryState.Translog translogState = indexShard.recoveryState().getTranslog();
         if (restoreSource == null) {
             listener.onFailure(new IndexShardRestoreFailedException(shardId, "empty restore source"));

server/src/main/java/org/elasticsearch/indices/recovery/PeerRecoveryTargetService.java

@@ -171,11 +171,13 @@ private void doRecovery(final long recoveryId) {
             final RecoveryTarget recoveryTarget = recoveryRef.target();
             timer = recoveryTarget.state().getTimer();
             cancellableThreads = recoveryTarget.cancellableThreads();
+            final IndexShard indexShard = recoveryTarget.indexShard();
             try {
+                indexShard.preRecovery();
                 assert recoveryTarget.sourceNode() != null : "can not do a recovery without a source node";
                 logger.trace("{} preparing shard for peer recovery", recoveryTarget.shardId());
-                recoveryTarget.indexShard().prepareForIndexRecovery();
-                final long startingSeqNo = recoveryTarget.indexShard().recoverLocallyUpToGlobalCheckpoint();
+                indexShard.prepareForIndexRecovery();
+                final long startingSeqNo = indexShard.recoverLocallyUpToGlobalCheckpoint();
                 assert startingSeqNo == UNASSIGNED_SEQ_NO || recoveryTarget.state().getStage() == RecoveryState.Stage.TRANSLOG :
                     "unexpected recovery stage [" + recoveryTarget.state().getStage() + "] starting seqno [ " + startingSeqNo + "]";
                 request = getStartRecoveryRequest(logger, clusterService.localNode(), recoveryTarget, startingSeqNo);

x-pack/plugin/searchable-snapshots/src/main/java/org/elasticsearch/xpack/searchablesnapshots/SearchableSnapshotIndexEventListener.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.searchablesnapshots;
+
+import org.apache.lucene.index.SegmentInfos;
+import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.seqno.SequenceNumbers;
+import org.elasticsearch.index.shard.IndexEventListener;
+import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.index.translog.Translog;
+import org.elasticsearch.index.translog.TranslogException;
+import org.elasticsearch.threadpool.ThreadPool;
+
+import java.nio.file.Path;
+
+import static org.elasticsearch.xpack.searchablesnapshots.SearchableSnapshots.isSearchableSnapshotStore;
+
+public class SearchableSnapshotIndexEventListener implements IndexEventListener {
+
+    @Override
+    public void beforeIndexShardRecovery(IndexShard indexShard, IndexSettings indexSettings) {
+        assert Thread.currentThread().getName().contains(ThreadPool.Names.GENERIC);
+        associateNewEmptyTranslogWithIndex(indexShard);
+    }
+
+    private static void associateNewEmptyTranslogWithIndex(IndexShard indexShard) {
+        final ShardId shardId = indexShard.shardId();
+        assert isSearchableSnapshotStore(indexShard.indexSettings().getSettings()) : "Expected a searchable snapshot shard " + shardId;
+        try {
+            final SegmentInfos segmentInfos = indexShard.store().readLastCommittedSegmentsInfo();
+            final long localCheckpoint = Long.parseLong(segmentInfos.userData.get(SequenceNumbers.LOCAL_CHECKPOINT_KEY));
+            final long primaryTerm = indexShard.getPendingPrimaryTerm();
+            final String translogUUID = segmentInfos.userData.get(Translog.TRANSLOG_UUID_KEY);
+            final Path translogLocation = indexShard.shardPath().resolveTranslog();
+            Translog.createEmptyTranslog(translogLocation, shardId, localCheckpoint, primaryTerm, translogUUID, null);
+        } catch (Exception e) {
+            throw new TranslogException(shardId, "failed to associate a new translog", e);
+        }
+    }
+}

x-pack/plugin/searchable-snapshots/src/main/java/org/elasticsearch/xpack/searchablesnapshots/SearchableSnapshots.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.env.Environment;
 import org.elasticsearch.env.NodeEnvironment;
+import org.elasticsearch.index.IndexModule;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.engine.EngineFactory;
 import org.elasticsearch.index.engine.ReadOnlyEngine;
@@ -40,8 +41,8 @@
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.watcher.ResourceWatcherService;
-import org.elasticsearch.xpack.searchablesnapshots.action.ClearSearchableSnapshotsCacheAction;
 import org.elasticsearch.xpack.core.searchablesnapshots.MountSearchableSnapshotAction;
+import org.elasticsearch.xpack.searchablesnapshots.action.ClearSearchableSnapshotsCacheAction;
 import org.elasticsearch.xpack.searchablesnapshots.action.SearchableSnapshotsStatsAction;
 import org.elasticsearch.xpack.searchablesnapshots.action.TransportClearSearchableSnapshotsCacheAction;
 import org.elasticsearch.xpack.searchablesnapshots.action.TransportMountSearchableSnapshotAction;
@@ -125,6 +126,13 @@ public void onRepositoriesModule(RepositoriesModule repositoriesModule) {
         repositoriesService.set(repositoriesModule.getRepositoryService());
     }
 
+    @Override
+    public void onIndexModule(IndexModule indexModule) {
+        if (isSearchableSnapshotStore(indexModule.getSettings())) {
+            indexModule.addIndexEventListener(new SearchableSnapshotIndexEventListener());
+        }
+    }
+
     @Override
     public Map<String, DirectoryFactory> getDirectoryFactories() {
         return Map.of(SNAPSHOT_DIRECTORY_FACTORY_KEY, (indexSettings, shardPath) -> {
@@ -138,7 +146,7 @@ public Map<String, DirectoryFactory> getDirectoryFactories() {
 
     @Override
     public Optional<EngineFactory> getEngineFactory(IndexSettings indexSettings) {
-        if (SNAPSHOT_DIRECTORY_FACTORY_KEY.equals(INDEX_STORE_TYPE_SETTING.get(indexSettings.getSettings()))
+        if (isSearchableSnapshotStore(indexSettings.getSettings())
                 && indexSettings.getSettings().getAsBoolean("index.frozen", false) == false) {
             return Optional.of(engineConfig -> new ReadOnlyEngine(engineConfig, null, new TranslogStats(), false, Function.identity()));
         }
@@ -169,5 +177,9 @@ public List<RestHandler> getRestHandlers(Settings settings, RestController restC
     public Map<String, ExistingShardsAllocator> getExistingShardsAllocators() {
         return Collections.singletonMap(SearchableSnapshotAllocator.ALLOCATOR_NAME, new SearchableSnapshotAllocator());
     }
+
+    static boolean isSearchableSnapshotStore(Settings indexSettings) {
+        return SNAPSHOT_DIRECTORY_FACTORY_KEY.equals(INDEX_STORE_TYPE_SETTING.get(indexSettings));
+    }
 }