Use pluggable DirectoryFactory in shard store validation (draft)

server/src/main/java/org/opensearch/gateway/TransportNodesGatewayStartedShardHelper.java

@@ -20,13 +20,16 @@
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.env.NodeEnvironment;
+import org.opensearch.index.IndexModule;
+import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.shard.ShardPath;
 import org.opensearch.index.shard.ShardStateMetadata;
 import org.opensearch.index.store.Store;
 import org.opensearch.indices.IndicesService;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
+import org.opensearch.plugins.IndexStorePlugin;
 
 import java.io.IOException;
 import java.util.Objects;
@@ -84,7 +87,14 @@ public static GatewayStartedShard getShardInfoOnLocalNode(
                     if (shardPath == null) {
                         throw new IllegalStateException(shardId + " no shard path found");
                     }
-                    Store.tryOpenIndex(shardPath.resolveIndex(), shardId, nodeEnv::shardLock, logger);
+                    Store.tryOpenIndex(
+                        shardPath.resolveIndex(),
+                        shardId,
+                        nodeEnv::shardLock,
+                        logger,
+                        resolveDirectoryFactory(shardId, indicesService, settings, clusterService),
+                        resolveIndexSettings(shardId, indicesService, settings, clusterService)
+                    );
                 } catch (Exception exception) {
                     final ShardPath finalShardPath = shardPath;
                     logger.trace(
@@ -114,6 +124,47 @@ public static GatewayStartedShard getShardInfoOnLocalNode(
         return new GatewayStartedShard(null, false, null);
     }
 
+    /**
+     * Resolves the {@link IndexStorePlugin.DirectoryFactory} for the given shard's index.
+     * Returns null if the index uses the default store type.
+     */
+    private static IndexStorePlugin.DirectoryFactory resolveDirectoryFactory(
+        ShardId shardId,
+        IndicesService indicesService,
+        Settings settings,
+        ClusterService clusterService
+    ) {
+        IndexSettings indexSettings = resolveIndexSettings(shardId, indicesService, settings, clusterService);
+        if (indexSettings == null) {
+            return null;
+        }
+        String storeType = IndexModule.INDEX_STORE_TYPE_SETTING.get(indexSettings.getSettings());
+        if (storeType.isEmpty()) {
+            return null;
+        }
+        return indicesService.getDirectoryFactories().get(storeType);
+    }
+
+    /**
+     * Resolves the {@link IndexSettings} for the given shard's index.
+     */
+    private static IndexSettings resolveIndexSettings(
+        ShardId shardId,
+        IndicesService indicesService,
+        Settings settings,
+        ClusterService clusterService
+    ) {
+        IndexService indexService = indicesService.indexService(shardId.getIndex());
+        if (indexService != null) {
+            return indexService.getIndexSettings();
+        }
+        IndexMetadata metadata = clusterService.state().metadata().index(shardId.getIndex());
+        if (metadata != null) {
+            return new IndexSettings(metadata, settings);
+        }
+        return null;
+    }
+
     /**
      * This class encapsulates the metadata about a started shard that needs to be persisted or sent between nodes.
      * This is used in {@link TransportNodesListGatewayStartedShardsBatch.NodeGatewayStartedShardsBatch} to construct the response for each node, instead of

server/src/main/java/org/opensearch/index/store/Store.java

@@ -618,10 +618,28 @@ public static MetadataSnapshot readMetadataSnapshot(
         ShardId shardId,
         NodeEnvironment.ShardLocker shardLocker,
         Logger logger
+    ) throws IOException {
+        return readMetadataSnapshot(indexLocation, shardId, shardLocker, logger, null, null);
+    }
+
+    /**
+     * Reads a MetadataSnapshot from the given index locations or returns an empty snapshot if it can't be read.
+     * When a {@link IndexStorePlugin.DirectoryFactory} and {@link IndexSettings} are provided, the directory
+     * is created using the factory (e.g. for encrypted indices). Otherwise falls back to a plain {@link NIOFSDirectory}.
+     *
+     * @throws IOException if the index we try to read is corrupted
+     */
+    public static MetadataSnapshot readMetadataSnapshot(
+        Path indexLocation,
+        ShardId shardId,
+        NodeEnvironment.ShardLocker shardLocker,
+        Logger logger,
+        IndexStorePlugin.DirectoryFactory directoryFactory,
+        IndexSettings indexSettings
     ) throws IOException {
         try (
             ShardLock lock = shardLocker.lock(shardId, "read metadata snapshot", TimeUnit.SECONDS.toMillis(5));
-            Directory dir = new NIOFSDirectory(indexLocation)
+            Directory dir = openDirectory(indexLocation, directoryFactory, indexSettings, shardId)
         ) {
             failIfCorrupted(dir);
             return new MetadataSnapshot((IndexCommit) null, dir, logger);
@@ -642,16 +660,51 @@ public static MetadataSnapshot readMetadataSnapshot(
      */
     public static void tryOpenIndex(Path indexLocation, ShardId shardId, NodeEnvironment.ShardLocker shardLocker, Logger logger)
         throws IOException, ShardLockObtainFailedException {
+        tryOpenIndex(indexLocation, shardId, shardLocker, logger, null, null);
+    }
+
+    /**
+     * Tries to open an index for the given location. This includes reading the
+     * segment infos and possible corruption markers. If the index can not
+     * be opened, an exception is thrown.
+     * When a {@link IndexStorePlugin.DirectoryFactory} and {@link IndexSettings} are provided, the directory
+     * is created using the factory (e.g. for encrypted indices). Otherwise falls back to a plain {@link NIOFSDirectory}.
+     */
+    public static void tryOpenIndex(
+        Path indexLocation,
+        ShardId shardId,
+        NodeEnvironment.ShardLocker shardLocker,
+        Logger logger,
+        IndexStorePlugin.DirectoryFactory directoryFactory,
+        IndexSettings indexSettings
+    ) throws IOException, ShardLockObtainFailedException {
         try (
             ShardLock lock = shardLocker.lock(shardId, "open index", TimeUnit.SECONDS.toMillis(5));
-            Directory dir = new NIOFSDirectory(indexLocation)
+            Directory dir = openDirectory(indexLocation, directoryFactory, indexSettings, shardId)
         ) {
             failIfCorrupted(dir);
             SegmentInfos segInfo = Lucene.readSegmentInfos(dir);
             logger.trace("{} loaded segment info [{}]", shardId, segInfo);
         }
     }
 
+    /**
+     * Opens a directory for the given index location. If a {@link IndexStorePlugin.DirectoryFactory} is provided,
+     * it is used to create the directory; otherwise a plain {@link NIOFSDirectory} is used.
+     */
+    private static Directory openDirectory(
+        Path indexLocation,
+        IndexStorePlugin.DirectoryFactory directoryFactory,
+        IndexSettings indexSettings,
+        ShardId shardId
+    ) throws IOException {
+        if (directoryFactory != null && indexSettings != null) {
+            ShardPath shardPath = new ShardPath(false, indexLocation.getParent(), indexLocation.getParent(), shardId);
+            return directoryFactory.newDirectory(indexSettings, shardPath);
+        }
+        return new NIOFSDirectory(indexLocation);
+    }
+
     /**
      * The returned IndexOutput validates the files checksum.
      * <p>

server/src/main/java/org/opensearch/indices/IndicesService.java

@@ -1418,6 +1418,13 @@ public IndicesQueryCache getIndicesQueryCache() {
         return indicesQueryCache;
     }
 
+    /**
+     * Returns the registered directory factories keyed by store type name.
+     */
+    public Map<String, IndexStorePlugin.DirectoryFactory> getDirectoryFactories() {
+        return directoryFactories;
+    }
+
     /**
      * Accumulate stats from the passed Object. Use this instead of
      * {@link #addDocStatusStats(org.opensearch.index.shard.IndexingStats.Stats.DocStatusStats)} after

server/src/main/java/org/opensearch/indices/store/TransportNodesListShardStoreMetadataHelper.java

@@ -21,6 +21,7 @@
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.env.NodeEnvironment;
+import org.opensearch.index.IndexModule;
 import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.seqno.ReplicationTracker;
@@ -30,6 +31,7 @@
 import org.opensearch.index.store.Store;
 import org.opensearch.index.store.StoreFileMetadata;
 import org.opensearch.indices.IndicesService;
+import org.opensearch.plugins.IndexStorePlugin;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -108,11 +110,20 @@ public static StoreFilesMetadata listShardMetadataInternal(
             // 1) a shard is being constructed, which means the cluster-manager will not use a copy of this replica
             // 2) A shard is shutting down and has not cleared it's content within lock timeout. In this case the cluster-manager may not
             // reuse local resources.
+            final IndexStorePlugin.DirectoryFactory directoryFactory = resolveDirectoryFactory(
+                shardId,
+                indicesService,
+                settings,
+                clusterService
+            );
+            final IndexSettings indexSettings = resolveIndexSettings(shardId, indicesService, settings, clusterService);
             final Store.MetadataSnapshot metadataSnapshot = Store.readMetadataSnapshot(
                 shardPath.resolveIndex(),
                 shardId,
                 nodeEnv::shardLock,
-                logger
+                logger,
+                directoryFactory,
+                indexSettings
             );
             // We use peer recovery retention leases from the primary for allocating replicas. We should always have retention leases when
             // we refresh shard info after the primary has started. Hence, we can ignore retention leases if there is no active shard.
@@ -127,6 +138,47 @@ public static StoreFilesMetadata listShardMetadataInternal(
         }
     }
 
+    /**
+     * Resolves the {@link IndexStorePlugin.DirectoryFactory} for the given shard's index.
+     * Returns null if the index uses the default store type.
+     */
+    private static IndexStorePlugin.DirectoryFactory resolveDirectoryFactory(
+        ShardId shardId,
+        IndicesService indicesService,
+        Settings settings,
+        ClusterService clusterService
+    ) {
+        IndexSettings indexSettings = resolveIndexSettings(shardId, indicesService, settings, clusterService);
+        if (indexSettings == null) {
+            return null;
+        }
+        String storeType = IndexModule.INDEX_STORE_TYPE_SETTING.get(indexSettings.getSettings());
+        if (storeType.isEmpty()) {
+            return null;
+        }
+        return indicesService.getDirectoryFactories().get(storeType);
+    }
+
+    /**
+     * Resolves the {@link IndexSettings} for the given shard's index.
+     */
+    private static IndexSettings resolveIndexSettings(
+        ShardId shardId,
+        IndicesService indicesService,
+        Settings settings,
+        ClusterService clusterService
+    ) {
+        IndexService indexService = indicesService.indexService(shardId.getIndex());
+        if (indexService != null) {
+            return indexService.getIndexSettings();
+        }
+        IndexMetadata metadata = clusterService.state().metadata().index(shardId.getIndex());
+        if (metadata != null) {
+            return new IndexSettings(metadata, settings);
+        }
+        return null;
+    }
+
     /**
      * Metadata for store files
      *