Addressed Todos

Author: Aditi Goyal

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

@@ -8,15 +8,22 @@
 
 package org.opensearch.index.shard;
 
+import org.opensearch.common.util.UploadListener;
 import org.opensearch.core.action.ActionListener;
 
 import java.util.Collection;
 import java.util.Map;
+import java.util.function.Function;
 
 /**
  * Interface to handle the functionality for upload data in the remote store
  */
 public interface RemoteStoreUploader {
 
-    void syncAndUploadNewSegments(Collection<String> localSegments, Map<String, Long> localSegmentsSizeMap, ActionListener<Void> listener);
+    void uploadSegments(
+        Collection<String> localSegments,
+        Map<String, Long> localSegmentsSizeMap,
+        ActionListener<Void> listener,
+        Function<Map<String, Long>, UploadListener> uploadListenerFunction
+    );
 }

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

@@ -10,26 +10,21 @@
 
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
-import org.apache.lucene.codecs.CodecUtil;
 import org.apache.lucene.index.CorruptIndexException;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.FilterDirectory;
 import org.apache.lucene.store.IOContext;
-import org.apache.lucene.store.IndexInput;
 import org.opensearch.action.support.GroupedActionListener;
 import org.opensearch.cluster.routing.RecoverySource;
 import org.opensearch.common.logging.Loggers;
 import org.opensearch.common.util.UploadListener;
 import org.opensearch.core.action.ActionListener;
-import org.opensearch.index.remote.RemoteSegmentTransferTracker;
 import org.opensearch.index.store.CompositeDirectory;
 import org.opensearch.index.store.RemoteSegmentStoreDirectory;
 
-import java.io.IOException;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.Map;
-import java.util.Set;
+import java.util.function.Function;
 
 /**
  * The service essentially acts as a bridge between local segment storage and remote storage,
@@ -39,69 +34,56 @@ public class RemoteStoreUploaderService implements RemoteStoreUploader {
 
     private final Logger logger;
 
-    public static final Set<String> EXCLUDE_FILES = Set.of("write.lock");
-
     private final IndexShard indexShard;
     private final Directory storeDirectory;
     private final RemoteSegmentStoreDirectory remoteDirectory;
-    private final Map<String, String> localSegmentChecksumMap;
-    // todo: check if we need to create a separate segment tracker as it is said to be linked to RL
-    private final RemoteSegmentTransferTracker segmentTracker;
 
-    public RemoteStoreUploaderService(
-        IndexShard indexShard,
-        Directory storeDirectory,
-        RemoteSegmentStoreDirectory remoteDirectory,
-        RemoteSegmentTransferTracker segmentTracker
-    ) {
-        this.indexShard = indexShard;
-        // todo: check the prefix to be add for this class
+    public RemoteStoreUploaderService(IndexShard indexShard, Directory storeDirectory, RemoteSegmentStoreDirectory remoteDirectory) {
         logger = Loggers.getLogger(getClass(), indexShard.shardId());
+        this.indexShard = indexShard;
         this.storeDirectory = storeDirectory;
         this.remoteDirectory = remoteDirectory;
-        this.segmentTracker = segmentTracker;
-        this.localSegmentChecksumMap = new HashMap<>();
     }
 
     @Override
-    public void syncAndUploadNewSegments(
+    public void uploadSegments(
         Collection<String> localSegments,
         Map<String, Long> localSegmentsSizeMap,
-        ActionListener<Void> listener
+        ActionListener<Void> listener,
+        Function<Map<String, Long>, UploadListener> uploadListenerFunction
     ) {
-
-        Collection<String> filteredFiles = localSegments.stream().filter(file -> !skipUpload(file)).toList();
-        if (filteredFiles.isEmpty()) {
+        if (localSegments.isEmpty()) {
             logger.debug("No new segments to upload in uploadNewSegments");
             listener.onResponse(null);
             return;
         }
 
-        logger.debug("Effective new segments files to upload {}", filteredFiles);
+        logger.debug("Effective new segments files to upload {}", localSegments);
         ActionListener<Collection<Void>> mappedListener = ActionListener.map(listener, resp -> null);
-        GroupedActionListener<Void> batchUploadListener = new GroupedActionListener<>(mappedListener, filteredFiles.size());
+        GroupedActionListener<Void> batchUploadListener = new GroupedActionListener<>(mappedListener, localSegments.size());
         Directory directory = ((FilterDirectory) (((FilterDirectory) storeDirectory).getDelegate())).getDelegate();
 
-        for (String filteredFile : filteredFiles) {
+        for (String localSegment : localSegments) {
             // Initializing listener here to ensure that the stats increment operations are thread-safe
-            UploadListener statsListener = createUploadListener(localSegmentsSizeMap);
+            UploadListener statsListener = uploadListenerFunction.apply(localSegmentsSizeMap);
             ActionListener<Void> aggregatedListener = ActionListener.wrap(resp -> {
-                statsListener.onSuccess(filteredFile);
+                statsListener.onSuccess(localSegment);
                 batchUploadListener.onResponse(resp);
+                // Once uploaded to Remote, local files become eligible for eviction from FileCache
                 if (directory instanceof CompositeDirectory) {
-                    ((CompositeDirectory) directory).afterSyncToRemote(filteredFile);
+                    ((CompositeDirectory) directory).afterSyncToRemote(localSegment);
                 }
             }, ex -> {
                 logger.warn(() -> new ParameterizedMessage("Exception: [{}] while uploading segment files", ex), ex);
                 if (ex instanceof CorruptIndexException) {
                     indexShard.failShard(ex.getMessage(), ex);
                 }
-                statsListener.onFailure(filteredFile);
+                statsListener.onFailure(localSegment);
                 batchUploadListener.onFailure(ex);
             });
-            statsListener.beforeUpload(filteredFile);
+            statsListener.beforeUpload(localSegment);
             // Place where the actual upload is happening
-            remoteDirectory.copyFrom(storeDirectory, filteredFile, IOContext.DEFAULT, aggregatedListener, isLowPriorityUpload());
+            remoteDirectory.copyFrom(storeDirectory, localSegment, IOContext.DEFAULT, aggregatedListener, isLowPriorityUpload());
         }
     }
 
@@ -119,67 +101,4 @@ boolean isLocalOrSnapshotRecoveryOrSeeding() {
                 || indexShard.recoveryState().getRecoverySource().getType() == RecoverySource.Type.SNAPSHOT
                 || indexShard.shouldSeedRemoteStore());
     }
-
-    /**
-     * Creates an {@link UploadListener} containing the stats population logic which would be triggered before and after segment upload events
-     *
-     * @param fileSizeMap updated map of current snapshot of local segments to their sizes
-     */
-    private UploadListener createUploadListener(Map<String, Long> fileSizeMap) {
-        return new UploadListener() {
-            private long uploadStartTime = 0;
-
-            @Override
-            public void beforeUpload(String file) {
-                // Start tracking the upload bytes started
-                segmentTracker.addUploadBytesStarted(fileSizeMap.get(file));
-                uploadStartTime = System.currentTimeMillis();
-            }
-
-            @Override
-            public void onSuccess(String file) {
-                // Track upload success
-                segmentTracker.addUploadBytesSucceeded(fileSizeMap.get(file));
-                segmentTracker.addToLatestUploadedFiles(file);
-                segmentTracker.addUploadTimeInMillis(Math.max(1, System.currentTimeMillis() - uploadStartTime));
-            }
-
-            @Override
-            public void onFailure(String file) {
-                // Track upload failure
-                segmentTracker.addUploadBytesFailed(fileSizeMap.get(file));
-                segmentTracker.addUploadTimeInMillis(Math.max(1, System.currentTimeMillis() - uploadStartTime));
-            }
-        };
-    }
-
-    /**
-     * Whether to upload a file or not depending on whether file is in excluded list or has been already uploaded.
-     *
-     * @param file that needs to be uploaded.
-     * @return true if the upload has to be skipped for the file.
-     */
-    private boolean skipUpload(String file) {
-        try {
-            // Exclude files that are already uploaded and the exclude files to come up with the list of files to be uploaded.
-            // todo: Check if we need the second condition or is it just fail safe
-            return EXCLUDE_FILES.contains(file) || remoteDirectory.containsFile(file, getChecksumOfLocalFile(file));
-        } catch (IOException e) {
-            logger.error(
-                "Exception while reading checksum of local segment file: {}, ignoring the exception and re-uploading the file",
-                file
-            );
-        }
-        return false;
-    }
-
-    private String getChecksumOfLocalFile(String file) throws IOException {
-        if (!localSegmentChecksumMap.containsKey(file)) {
-            try (IndexInput indexInput = storeDirectory.openInput(file, IOContext.READONCE)) {
-                String checksum = Long.toString(CodecUtil.retrieveChecksum(indexInput));
-                localSegmentChecksumMap.put(file, checksum);
-            }
-        }
-        return localSegmentChecksumMap.get(file);
-    }
 }