Fix for failure IT's

Signed-off-by: Sandeep Kumawat <skumwt@amazon.com>

Author: Sandeep Kumawat

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationBaseIT.java

@@ -42,7 +42,6 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
@@ -244,12 +243,4 @@ protected SegmentInfos getLatestSegmentInfos(IndexShard shard) throws IOExceptio
             return closeable.get();
         }
     }
-
-    protected boolean warmIndexSegmentReplicationEnabled() {
-        return Objects.equals(
-            IndexModule.INDEX_STORE_LOCALITY_SETTING.get(indexSettings()).toString(),
-            IndexModule.DataLocalityType.PARTIAL.name()
-        );
-    }
-
 }

server/src/internalClusterTest/java/org/opensearch/indices/replication/WarmIndexRemoteStoreSegmentReplicationIT.java renamed to server/src/internalClusterTest/java/org/opensearch/indices/replication/WarmIndexSegmentReplicationIT.java

@@ -348,7 +348,7 @@ private void unassignIgnoredRemoteShards(RoutingAllocation routingAllocation) {
                 // to re-fetch any shard blocks from the repository.
                 if (shard.primary()) {
                     if (RecoverySource.Type.SNAPSHOT.equals(shard.recoverySource().getType()) == false
-                        && !isPartialIndex(allocation.metadata().getIndexSafe(shard.index()))) {
+                        && isPartialIndex(allocation.metadata().getIndexSafe(shard.index())) == false) {
                         unassignedShard = shard.updateUnassigned(shard.unassignedInfo(), RecoverySource.EmptyStoreRecoverySource.INSTANCE);
                     }
                 }

server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/RemoteShardsBalancer.java

@@ -370,7 +370,7 @@ public void flush(boolean force, boolean waitIfOngoing) throws EngineException {
         ensureOpen();
         // Skip flushing for indices with partial locality (warm indices)
         // For these indices, we don't need to commit as we will sync from the remote store on re-open
-        if (engineConfig.getIndexSettings().isStoreLocalityPartial() == false) {
+        if (engineConfig.getIndexSettings().isStoreLocalityPartial()) {
             return;
         }
         // readLock is held here to wait/block any concurrent close that acquires the writeLock.
@@ -447,7 +447,9 @@ protected final void closeNoLock(String reason, CountDownLatch closedLatch) {
                     latestSegmentInfos.changed();
                 }
                 try {
-                    commitSegmentInfos(latestSegmentInfos);
+                    if (engineConfig.getIndexSettings().isStoreLocalityPartial() == false) {
+                        commitSegmentInfos(latestSegmentInfos);
+                    }
                 } catch (IOException e) {
                     // mark the store corrupted unless we are closing as result of engine failure.
                     // in this case Engine#failShard will handle store corruption.

server/src/main/java/org/opensearch/index/engine/NRTReplicationEngine.java

@@ -5110,8 +5110,6 @@ public void syncSegmentsFromRemoteSegmentStore(boolean overrideLocal) throws IOE
      */
     public void syncSegmentsFromRemoteSegmentStore(boolean overrideLocal, final Runnable onFileSync) throws IOException {
         boolean syncSegmentSuccess = false;
-        // For warm indices, it doesn't make sense to check for local existence of files.
-        overrideLocal = overrideLocal && indexSettings.isStoreLocalityPartial() == false;
         long startTimeMs = System.currentTimeMillis();
         assert indexSettings.isRemoteStoreEnabled() || this.isRemoteSeeded();
         logger.trace("Downloading segments from remote segment store");
@@ -5161,7 +5159,7 @@ public void syncSegmentsFromRemoteSegmentStore(boolean overrideLocal, final Runn
                     }
                 }
                 assert Arrays.stream(store.directory().listAll()).filter(f -> f.startsWith(IndexFileNames.SEGMENTS)).findAny().isEmpty()
-                    : "There should not be any segments file in the dir";
+                    || indexSettings.isStoreLocalityPartial() : "There should not be any segments file in the dir";
                 store.commitSegmentInfos(infosSnapshot, processedLocalCheckpoint, processedLocalCheckpoint);
             }
             syncSegmentSuccess = true;

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

@@ -34,8 +34,12 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.apache.lucene.index.IndexFileNames.SEGMENTS;
 
 /**
  * Composite Directory will contain both local and remote directory
@@ -74,22 +78,47 @@ public CompositeDirectory(Directory localDirectory, Directory remoteDirectory, F
         );
     }
 
+    /**
+     * Returns names of all files stored in local directory
+     * @throws IOException in case of I/O error
+     */
+    private String[] listLocalFiles() throws IOException {
+        ensureOpen();
+        logger.trace("Composite Directory[{}]: listLocalOnly() called", this::toString);
+        return localDirectory.listAll();
+    }
+
+    /**
+     * Returns a list of names of all block files stored in the local directory for a given file,
+     * including the original file itself if present.
+     *
+     * @param fileName The name of the file to search for, along with its associated block files.
+     * @return A list of file names, including the original file (if present) and all its block files.
+     * @throws IOException in case of I/O error while listing files.
+     */
+    private List<String> listBlockFiles(String fileName) throws IOException {
+        return Stream.of(listLocalFiles())
+            .filter(file -> file.equals(fileName) || file.startsWith(fileName + FileTypeUtils.BLOCK_FILE_IDENTIFIER))
+            .collect(Collectors.toList());
+    }
+
     /**
      * Returns names of all files stored in this directory in sorted order
      * Does not include locally stored block files (having _block_ in their names) and files pending deletion
      *
      * @throws IOException in case of I/O error
      */
+    // TODO: https://github.com/opensearch-project/OpenSearch/issues/17527
     @Override
     public String[] listAll() throws IOException {
         ensureOpen();
         logger.trace("Composite Directory[{}]: listAll() called", this::toString);
         String[] localFiles = localDirectory.listAll();
         Set<String> allFiles = new HashSet<>(Arrays.asList(localFiles));
-        // String[] remoteFiles = getRemoteFiles();
-        // allFiles.addAll(Arrays.asList(remoteFiles));
+        String[] remoteFiles = getRemoteFiles();
+        allFiles.addAll(Arrays.asList(remoteFiles));
         logger.trace("Composite Directory[{}]: Local Directory files - {}", this::toString, () -> Arrays.toString(localFiles));
-        // logger.trace("Composite Directory[{}]: Remote Directory files - {}", this::toString, () -> Arrays.toString(remoteFiles));
+        logger.trace("Composite Directory[{}]: Remote Directory files - {}", this::toString, () -> Arrays.toString(remoteFiles));
         Set<String> nonBlockLuceneFiles = allFiles.stream()
             .filter(file -> !FileTypeUtils.isBlockFile(file))
             .collect(Collectors.toUnmodifiableSet());
@@ -105,6 +134,7 @@ public String[] listAll() throws IOException {
      * Currently deleting only from local directory as files from remote should not be deleted as that is taken care by garbage collection logic of remote directory
      * @param name the name of an existing file.
      * @throws IOException in case of I/O error
+     * @throws NoSuchFileException when file does not exist in the directory
      */
     @Override
     public void deleteFile(String name) throws IOException {
@@ -113,12 +143,22 @@ public void deleteFile(String name) throws IOException {
         if (FileTypeUtils.isTempFile(name)) {
             localDirectory.deleteFile(name);
         } else if (Arrays.asList(listAll()).contains(name) == false) {
-            logger.debug("The file [{}] does not exist", name);
-            // we should not fail here as localDirectory might not contain this file.
-            // throw new NoSuchFileException("File " + name + " not found in directory");
+            throw new NoSuchFileException("File " + name + " not found in directory");
         } else {
-            localDirectory.deleteFile(name);
-            fileCache.remove(getFilePath(name));
+            List<String> blockFiles = listBlockFiles(name);
+            if (blockFiles.isEmpty()) {
+                logger.debug("The file [{}] or its block files do not exist in local directory", name);
+                // we should not throw exception in this case as localDirectory might not contain this file.
+            } else {
+                for (String blockFile : blockFiles) {
+                    if (fileCache.get(getFilePath(blockFile)) == null) {
+                        logger.debug("The file [{}] exists in local but not part of FileCache, deleting it from local", blockFile);
+                        localDirectory.deleteFile(blockFile);
+                    } else {
+                        fileCache.remove(getFilePath(blockFile));
+                    }
+                }
+            }
         }
     }
 
@@ -257,6 +297,15 @@ public IndexInput openInput(String name, IOContext context) throws IOException {
     public void close() throws IOException {
         ensureOpen();
         logger.trace("Composite Directory[{}]: close() called", this::toString);
+        String[] localFiles = listLocalFiles();
+        for (String localFile : localFiles) {
+            // Delete the segments_N file that gets created with ref count 1 when index got created. Ideally this should not get created on
+            // replica
+            if (localFile.startsWith(SEGMENTS)) {
+                fileCache.remove(getFilePath(localFile));
+            }
+        }
+        fileCache.prune();
         localDirectory.close();
     }
 

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

@@ -18,12 +18,14 @@
 @ExperimentalApi
 public class FileTypeUtils {
 
+    public static String BLOCK_FILE_IDENTIFIER = "_block_";
+
     public static boolean isTempFile(String name) {
         return name.endsWith(".tmp");
     }
 
     public static boolean isBlockFile(String name) {
-        return name.contains("_block_");
+        return name.contains(BLOCK_FILE_IDENTIFIER);
     }
 
     public static boolean isExtraFSFile(String name) {

server/src/main/java/org/opensearch/index/store/remote/utils/FileTypeUtils.java

@@ -310,14 +310,19 @@ public CacheStats stats() {
     public void logCurrentState() {
         lock.lock();
         try {
-            String allFiles = "\n";
+            final StringBuilder allFiles = new StringBuilder("\n");
             for (Map.Entry<K, Node<K, V>> entry : data.entrySet()) {
                 String path = entry.getKey().toString();
                 String file = path.substring(path.lastIndexOf('/'));
-                allFiles += file + " [RefCount: " + entry.getValue().refCount + " , Weight: " + entry.getValue().weight + " ]\n";
+                allFiles.append(file)
+                    .append(" [RefCount: ")
+                    .append(entry.getValue().refCount)
+                    .append(" , Weight: ")
+                    .append(entry.getValue().weight)
+                    .append(" ]\n");
             }
-            if (allFiles.equals("\n") == false) {
-                logger.debug("Cache entries : {}", allFiles);
+            if (allFiles.length() > 1) {
+                logger.trace(() -> "Cache entries : " + allFiles);
             }
         } finally {
             lock.unlock();

server/src/main/java/org/opensearch/index/store/remote/utils/cache/LRUCache.java

@@ -191,7 +191,8 @@ public void logCurrentState() {
         int i = 0;
         for (RefCountedCache<K, V> cache : table) {
             if (cache.size() > 0) {
-                logger.trace("SegmentedCache " + i);
+                final int segmentIndex = i;
+                logger.trace(() -> "SegmentedCache " + segmentIndex);
                 ((LRUCache<K, V>) cache).logCurrentState();
             }
             i++;

server/src/main/java/org/opensearch/index/store/remote/utils/cache/SegmentedCache.java

@@ -202,8 +202,8 @@ private List<StoreFileMetadata> getFiles(CheckpointInfoResponse checkpointInfo)
         state.setStage(SegmentReplicationState.Stage.FILE_DIFF);
 
         // Return an empty list for warm indices, In this case, replica shards don't require downloading files from remote storage
-        // as they can will sync all files from remote in case failure.
-        if (indexShard.indexSettings().isStoreLocalityPartial() == true) {
+        // as replicas will sync all files from remote in case of failure.
+        if (indexShard.indexSettings().isStoreLocalityPartial()) {
             return Collections.emptyList();
         }
         final Store.RecoveryDiff diff = Store.segmentReplicationDiff(checkpointInfo.getMetadataMap(), indexShard.getSegmentMetadataMap());

server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTarget.java

@@ -39,8 +39,16 @@ public class CompositeDirectoryTests extends BaseRemoteSegmentStoreDirectoryTest
     private FSDirectory localDirectory;
     private CompositeDirectory compositeDirectory;
 
-    private final static String[] LOCAL_FILES = new String[] { "_1.cfe", "_2.cfe", "_0.cfe_block_7", "_0.cfs_block_7", "temp_file.tmp" };
+    private final static String[] LOCAL_FILES = new String[] {
+        "_1.cfe",
+        "_1.cfe_block_0",
+        "_1.cfe_block_1",
+        "_2.cfe",
+        "_0.cfe_block_7",
+        "_0.cfs_block_7",
+        "temp_file.tmp" };
     private final static String FILE_PRESENT_LOCALLY = "_1.cfe";
+    private final static String BLOCK_FILE_PRESENT_LOCALLY = "_1.cfe_block_0";
     private final static String FILE_PRESENT_IN_REMOTE_ONLY = "_0.si";
     private final static String NON_EXISTENT_FILE = "non_existent_file";
     private final static String NEW_FILE = "new_file";
@@ -67,9 +75,11 @@ public void testListAll() throws IOException {
 
     public void testDeleteFile() throws IOException {
         assertTrue(existsInCompositeDirectory(FILE_PRESENT_LOCALLY));
+        assertTrue(existsInLocalDirectory(BLOCK_FILE_PRESENT_LOCALLY));
         // Delete the file and assert that it no more is a part of the directory
         compositeDirectory.deleteFile(FILE_PRESENT_LOCALLY);
         assertFalse(existsInCompositeDirectory(FILE_PRESENT_LOCALLY));
+        assertFalse(existsInCompositeDirectory(BLOCK_FILE_PRESENT_LOCALLY));
         // Reading deleted file from directory should result in NoSuchFileException
         assertThrows(NoSuchFileException.class, () -> compositeDirectory.openInput(FILE_PRESENT_LOCALLY, IOContext.DEFAULT));
     }