[RFS] Collection of Minor Changes

RFS/build-preloaded-source-image.gradle

@@ -23,8 +23,8 @@ def createNetworkTask = task createNetwork(type: Exec) {
     }
 }
 task createInitialElasticsearchContainer(type: DockerCreateContainer) {
-    dependsOn createNetwork, buildDockerImage_emptyElasticsearchSource
-    targetImageId 'migrations/empty_elasticsearch_source:latest'
+    dependsOn createNetwork, buildDockerImage_emptyElasticsearchSource_7_17
+    targetImageId 'migrations/emptyElasticsearchSource_7_17:latest'
     containerName = "elasticsearch-${uniqueId}"
     hostConfig.network = myNetworkName
     hostConfig.dns = ['elasticsearch']

RFS/build.gradle

@@ -132,9 +132,12 @@ DockerServiceProps[] dockerServices = [
                                 dockerImageName:"reindex_from_snapshot",
                                 inputDir:"./docker",
                                 taskDependencies:["copyDockerRuntimeJars"]]),
-        new DockerServiceProps([projectName:"emptyElasticsearchSource",
-                                dockerImageName:"empty_elasticsearch_source",
+        new DockerServiceProps([projectName:"emptyElasticsearchSource_7_10",
+                                dockerImageName:"empty_elasticsearch_source_7_10",
                                 inputDir:"./docker/TestSource_ES_7_10"]),
+        new DockerServiceProps([projectName:"emptyElasticsearchSource_7_17",
+                                dockerImageName:"empty_elasticsearch_source_7_17",
+                                inputDir:"./docker/TestSource_ES_7_17"]),
         new DockerServiceProps([projectName:"trafficGenerator",
                                 dockerImageName:"osb_traffic_generator",
                                 inputDir:"./docker/TrafficGenerator",

RFS/docker/TestSource_ES_7_17/Dockerfile

@@ -0,0 +1,22 @@
+FROM docker.elastic.co/elasticsearch/elasticsearch:7.17.21 AS base
+
+# Configure Elastic
+ENV ELASTIC_SEARCH_CONFIG_FILE=/usr/share/elasticsearch/config/elasticsearch.yml
+# Prevents ES from complaining about nodes count
+RUN echo "discovery.type: single-node" >> $ELASTIC_SEARCH_CONFIG_FILE
+ENV PATH=${PATH}:/usr/share/elasticsearch/jdk/bin/
+
+# Install the S3 Repo Plugin
+RUN echo y | /usr/share/elasticsearch/bin/elasticsearch-plugin install repository-s3
+
+# Install the AWS CLI for testing purposes
+RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && \
+    unzip awscliv2.zip && \
+    ./aws/install
+
+RUN mkdir /snapshots && chown elasticsearch /snapshots
+
+# Install our custom entrypoint script
+COPY ./container-start.sh /usr/share/elasticsearch/container-start.sh
+
+CMD /usr/share/elasticsearch/container-start.sh

RFS/docker/TestSource_ES_7_17/container-start.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+echo "Setting AWS Creds from ENV Variables"
+bin/elasticsearch-keystore create
+echo $AWS_ACCESS_KEY_ID | bin/elasticsearch-keystore add s3.client.default.access_key --stdin
+echo $AWS_SECRET_ACCESS_KEY | bin/elasticsearch-keystore add s3.client.default.secret_key --stdin
+
+if [ -n "$AWS_SESSION_TOKEN" ]; then
+    echo $AWS_SESSION_TOKEN | bin/elasticsearch-keystore add s3.client.default.session_token --stdin
+fi
+
+echo "Starting Elasticsearch"
+/usr/local/bin/docker-entrypoint.sh eswrapper

RFS/src/main/java/com/rfs/DemoPrintOutSnapshot.java

@@ -27,6 +27,7 @@
 import com.rfs.common.SnapshotRepo;
 import com.rfs.common.SnapshotShardUnpacker;
 import com.rfs.common.ClusterVersion;
+import com.rfs.common.DefaultSourceRepoAccessor;
 import com.rfs.common.FileSystemRepo;
 import com.rfs.version_es_7_10.*;
 
@@ -219,8 +220,12 @@
                     } else {
                         bufferSize = ElasticsearchConstants_ES_7_10.BUFFER_SIZE_IN_BYTES;
                     }
-                    SnapshotShardUnpacker unpacker = new SnapshotShardUnpacker(repo, Paths.get(luceneBasePathString), bufferSize);
-                    unpacker.unpack(shardMetadata);
+                    DefaultSourceRepoAccessor repoAccessor = new DefaultSourceRepoAccessor(repo);
+
+                    SnapshotShardUnpacker.Factory unpackerFactory = new SnapshotShardUnpacker.Factory(repoAccessor, Paths.get(luceneBasePathString), bufferSize);
+                    try (SnapshotShardUnpacker unpacker = unpackerFactory.create(shardMetadata)) {
+                        unpacker.unpack();
+                    }
 
                     // Now, read the documents back out
                     System.out.println("--- Reading docs in the shard ---");

RFS/src/main/java/com/rfs/ReindexFromSnapshot.java

@@ -329,7 +329,8 @@
                 } else {
                     bufferSize = ElasticsearchConstants_ES_7_10.BUFFER_SIZE_IN_BYTES;
                 }
-                SnapshotShardUnpacker unpacker = new SnapshotShardUnpacker(repo, luceneDirPath, bufferSize);
+                DefaultSourceRepoAccessor repoAccessor = new DefaultSourceRepoAccessor(repo);
+                SnapshotShardUnpacker.Factory unpackerFactory = new SnapshotShardUnpacker.Factory(repoAccessor,luceneDirPath, bufferSize);
 
                 for (IndexMetadata.Data indexMetadata : indexMetadatas) {
                     logger.info("Processing index: " + indexMetadata.getName());
@@ -345,7 +346,9 @@
                         }
 
                         // Unpack the shard
-                        unpacker.unpack(shardMetadata);
+                        try (SnapshotShardUnpacker unpacker = unpackerFactory.create(shardMetadata)) {
+                            unpacker.unpack();
+                        }        
                     }
                 }
 

RFS/src/main/java/com/rfs/RunRfsWorker.java

@@ -21,6 +21,7 @@
 import com.rfs.cms.OpenSearchCmsClient;
 import com.rfs.common.ClusterVersion;
 import com.rfs.common.ConnectionDetails;
+import com.rfs.common.DeletingSourceRepoAccessor;
 import com.rfs.common.DocumentReindexer;
 import com.rfs.common.GlobalMetadata;
 import com.rfs.common.IndexMetadata;
@@ -100,6 +101,10 @@
         @Parameter(names = {"--component-template-allowlist"}, description = ("Optional. List of component template names to migrate"
             + " (e.g. 'posts_template1, posts_template2').  Default: empty list"), required = false)
         public List<String> componentTemplateAllowlist = List.of();
+
+        @Parameter(names = {"--max-shard-size-bytes"}, description = ("Optional. The maximum shard size, in bytes, to allow when"
+            + " performing the document migration.  Useful for preventing disk overflow.  Default: 50 * 1024 * 1024 * 1024 (50 GB)"), required = false)
+        public long maxShardSizeBytes = 50 * 1024 * 1024 * 1024L;
 
         //https://opensearch.org/docs/2.11/api-reference/cluster-api/cluster-awareness/
         @Parameter(names = {"--min-replicas"}, description = ("Optional.  The minimum number of replicas configured for migrated indices on the target."
@@ -132,6 +137,7 @@
         String targetPass = arguments.targetPass;
         List<String> indexTemplateAllowlist = arguments.indexTemplateAllowlist;
         List<String> componentTemplateAllowlist = arguments.componentTemplateAllowlist;
+        long maxShardSizeBytes = arguments.maxShardSizeBytes;
         int awarenessDimensionality = arguments.minNumberOfReplicas + 1;
         Level logLevel = arguments.logLevel;
 
@@ -165,10 +171,11 @@
             indexWorker.run();
 
             ShardMetadata.Factory shardMetadataFactory = new ShardMetadataFactory_ES_7_10(repoDataProvider);
-            SnapshotShardUnpacker unpacker = new SnapshotShardUnpacker(sourceRepo, luceneDirPath, ElasticsearchConstants_ES_7_10.BUFFER_SIZE_IN_BYTES);
+            DeletingSourceRepoAccessor repoAccessor = new DeletingSourceRepoAccessor(sourceRepo);
+            SnapshotShardUnpacker.Factory unpackerFactory = new SnapshotShardUnpacker.Factory(repoAccessor, luceneDirPath, ElasticsearchConstants_ES_7_10.BUFFER_SIZE_IN_BYTES);
             LuceneDocumentsReader reader = new LuceneDocumentsReader(luceneDirPath);
             DocumentReindexer reindexer = new DocumentReindexer(targetClient);
-            DocumentsRunner documentsWorker = new DocumentsRunner(globalState, cmsClient, snapshotName, indexMetadataFactory, shardMetadataFactory, unpacker, reader, reindexer);
+            DocumentsRunner documentsWorker = new DocumentsRunner(globalState, cmsClient, snapshotName, maxShardSizeBytes, indexMetadataFactory, shardMetadataFactory, unpackerFactory, reader, reindexer);
             documentsWorker.run();
 
         } catch (Runner.PhaseFailed e) {

RFS/src/main/java/com/rfs/common/DefaultSourceRepoAccessor.java

@@ -0,0 +1,23 @@
+package com.rfs.common;
+
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+/*
+ * Provides "simple" access to the underlying files in the source repo without any special behavior
+ */
+public class DefaultSourceRepoAccessor extends SourceRepoAccessor {
+    public DefaultSourceRepoAccessor(SourceRepo repo) {
+        super(repo);
+    }
+
+    @Override
+    protected InputStream load(Path path) {
+        try {
+            return Files.newInputStream(path);
+        } catch (Exception e) {
+            throw new CouldNotLoadRepoFile("Could not load file: " + path, e);
+        }
+    }
+}

RFS/src/main/java/com/rfs/common/DeletingSourceRepoAccessor.java

@@ -0,0 +1,46 @@
+package com.rfs.common;
+
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+
+/*
+ * Provides access to the underlying files in the source repo and deletes the files after the Stream is closed.  This
+ * is useful/interesting in the case where the files are large/numerous and you can easily re-acquire them - such as
+ * if they are being loaded from S3.
+ */
+public class DeletingSourceRepoAccessor extends SourceRepoAccessor {
+    public DeletingSourceRepoAccessor(SourceRepo repo) {
+        super(repo);
+    }
+
+    @Override
+    protected InputStream load(Path path) {
+        try {
+            return new DeletingFileInputStream(path);
+        } catch (Exception e) {
+            throw new CouldNotLoadRepoFile("Could not load file: " + path, e);
+        }
+    }
+
+    public static class DeletingFileInputStream extends FileInputStream {
+        private final Path filePath;
+
+        public DeletingFileInputStream(Path filePath) throws IOException {
+            super(filePath.toFile());
+            this.filePath = filePath;
+        }
+
+        @Override
+        public void close() throws IOException {
+            try {
+                super.close();
+            } finally {
+                Files.deleteIfExists(filePath);
+            }
+        }
+    }
+}

RFS/src/main/java/com/rfs/common/PartSliceStream.java

@@ -12,24 +12,23 @@
  */
 
 public class PartSliceStream extends InputStream {
-    private final SourceRepo repo;
+    private final SourceRepoAccessor repoAccessor;
     private final ShardMetadata.FileInfo fileMetadata;
     private final String indexId;
     private final int shardId;
     private long slice = 0;
     private InputStream currentStream;
     private boolean initialized = false;
 
-    public PartSliceStream(SourceRepo repo, ShardMetadata.FileInfo fileMetadata, String indexId, int shardId) {
-        this.repo = repo;
+    public PartSliceStream(SourceRepoAccessor repoAccessor, ShardMetadata.FileInfo fileMetadata, String indexId, int shardId) {
+        this.repoAccessor = repoAccessor;
         this.fileMetadata = fileMetadata;
         this.indexId = indexId;
         this.shardId = shardId;
     }
 
-    protected InputStream openSlice(long slice) throws IOException {
-        Path filePath = repo.getBlobFilePath(indexId, shardId, fileMetadata.partName(slice));
-        return Files.newInputStream(filePath);
+    protected InputStream openSlice(long slice) {
+        return repoAccessor.getBlobFile(indexId, shardId, fileMetadata.partName(slice));
     }
 
     private InputStream nextStream() throws IOException {

RFS/src/main/java/com/rfs/common/ShardMetadata.java

@@ -77,7 +77,7 @@ public static interface Data {
         public long getStartTime();    
         public long getTime();    
         public int getNumberOfFiles();    
-        public long getTotalSize();
+        public long getTotalSizeBytes();
         public List<FileInfo> getFiles();
     }
 

RFS/src/main/java/com/rfs/common/SnapshotShardUnpacker.java

@@ -1,6 +1,8 @@
 package com.rfs.common;
 
+import java.io.IOException;
 import java.io.InputStream;
+import java.nio.file.DirectoryStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -16,19 +18,31 @@
 import org.apache.lucene.util.BytesRef;
 
 @RequiredArgsConstructor
-public class SnapshotShardUnpacker {
+public class SnapshotShardUnpacker implements AutoCloseable {
     private static final Logger logger = LogManager.getLogger(SnapshotShardUnpacker.class);
-    protected final SourceRepo repo;
-    protected final  Path luceneFilesBasePath;
-    protected final int bufferSize;
+    private final SourceRepoAccessor repoAccessor;
+    private final Path luceneFilesBasePath;
+    private final ShardMetadata.Data shardMetadata;
+    private final int bufferSize;
 
-    public void unpack(ShardMetadata.Data shardMetadata) {
+    @RequiredArgsConstructor
+    public static class Factory {
+        private final SourceRepoAccessor repoAccessor;
+        private final Path luceneFilesBasePath;
+        private final int bufferSize;
+
+        public SnapshotShardUnpacker create(ShardMetadata.Data shardMetadata) {
+            return new SnapshotShardUnpacker(repoAccessor, luceneFilesBasePath, shardMetadata, bufferSize);
+        }
+    }
+
+    public void unpack() {
         try {
             // Some constants
             NativeFSLockFactory lockFactory = NativeFSLockFactory.INSTANCE;
 
             // Ensure the blob files are prepped, if they need to be
-            repo.prepBlobFiles(shardMetadata);
+            repoAccessor.prepBlobFiles(shardMetadata);
 
             // Create the directory for the shard's lucene files
             Path luceneIndexDir = Paths.get(luceneFilesBasePath + "/" + shardMetadata.getIndexName() + "/" + shardMetadata.getShardId());
@@ -43,7 +57,7 @@
                     final BytesRef hash = fileMetadata.getMetaHash();
                     indexOutput.writeBytes(hash.bytes, hash.offset, hash.length);
                 } else {
-                    try (InputStream stream = new PartSliceStream(repo, fileMetadata, shardMetadata.getIndexId(), shardMetadata.getShardId())) {
+                    try (InputStream stream = new PartSliceStream(repoAccessor, fileMetadata, shardMetadata.getIndexId(), shardMetadata.getShardId())) {
                         final byte[] buffer = new byte[Math.toIntExact(Math.min(bufferSize, fileMetadata.getLength()))];
                         int length;
                         while ((length = stream.read(buffer)) > 0) {
@@ -58,6 +72,36 @@
         }
     }
 
+    @Override
+    public void close() {
+        try {
+            Path luceneIndexDir = Paths.get(luceneFilesBasePath + "/" + shardMetadata.getIndexName() + "/" + shardMetadata.getShardId());
+            if (Files.exists(luceneIndexDir)) {
+                deleteRecursively(luceneIndexDir);
+            }
+
+        } catch (Exception e) {
+            throw new CouldNotCleanUpShard("Could not clean up shard: Index " + shardMetadata.getIndexId() + ", Shard " + shardMetadata.getShardId(), e);
+        }
+    }
+
+    protected void deleteRecursively(Path path) throws IOException {
+        if (Files.isDirectory(path)) {
+            try (DirectoryStream<Path> entries = Files.newDirectoryStream(path)) {
+                for (Path entry : entries) {
+                    deleteRecursively(entry);
+                }
+            }
+        }
+        Files.delete(path);
+    }
+
+    public static class CouldNotCleanUpShard extends RfsException {
+        public CouldNotCleanUpShard(String message, Exception e) {
+            super(message, e);
+        }
+    }
+
     public static class CouldNotUnpackShard extends RfsException {
         public CouldNotUnpackShard(String message, Exception e) {
             super(message, e);