Surface OnDemandBlockIndexInput blockSizeShift to RemoteSnapshotDirec…

…tory constructor Signed-off-by: Finn Carroll <[email protected]>
finnegancarroll · Jul 22, 2024 · 5e26dfe · 5e26dfe
1 parent b35690c
commit 5e26dfe
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 5 deletions.
diff --git a/...er/src/main/java/org/opensearch/index/store/remote/directory/RemoteSnapshotDirectory.java b/...er/src/main/java/org/opensearch/index/store/remote/directory/RemoteSnapshotDirectory.java
@@ -39,13 +39,24 @@
 public final class RemoteSnapshotDirectory extends Directory {
 
     public static final Version SEARCHABLE_SNAPSHOT_EXTENDED_COMPATIBILITY_MINIMUM_VERSION = LegacyESVersion.V_6_0_0;
-
     private static final String VIRTUAL_FILE_PREFIX = BlobStoreRepository.VIRTUAL_DATA_BLOB_PREFIX;
+    private static final int ON_DEMAND_DEFAULT_BLOCK_SIZE_SHIFT = 21;
 
     private final Map<String, BlobStoreIndexShardSnapshot.FileInfo> fileInfoMap;
     private final FSDirectory localStoreDir;
     private final TransferManager transferManager;
 
+    private int onDemandBlockSizeShift = ON_DEMAND_DEFAULT_BLOCK_SIZE_SHIFT;
+
+    public RemoteSnapshotDirectory(BlobStoreIndexShardSnapshot snapshot, FSDirectory localStoreDir, TransferManager transferManager, int onDemandBlockSizeShift) {
+        this(
+            snapshot,
+            localStoreDir,
+            transferManager
+        );
+        this.onDemandBlockSizeShift = onDemandBlockSizeShift;
+    }
+
     public RemoteSnapshotDirectory(BlobStoreIndexShardSnapshot snapshot, FSDirectory localStoreDir, TransferManager transferManager) {
         this.fileInfoMap = snapshot.indexFiles()
             .stream()
@@ -74,7 +85,26 @@ public IndexInput openInput(String name, IOContext context) throws IOException {
         if (fileInfo.name().startsWith(VIRTUAL_FILE_PREFIX)) {
             return new ByteArrayIndexInput(fileInfo.physicalName(), fileInfo.metadata().hash().bytes);
         }
-        return new OnDemandBlockSnapshotIndexInput(fileInfo, localStoreDir, transferManager);
+
+        String resourceDescription = "BlockedSnapshotIndexInput(path=\""
+            + localStoreDir.getDirectory().toString()
+            + "/"
+            + fileInfo.physicalName()
+            + "\", "
+            + "offset="
+            + 0
+            + ", length= "
+            + fileInfo.length()
+            + ")";
+
+        OnDemandBlockSnapshotIndexInput.Builder blockIndexInputBuilder = new OnDemandBlockSnapshotIndexInput.Builder()
+            .resourceDescription(resourceDescription)
+            .isClone(false)
+            .offset(0L)
+            .length(fileInfo.length())
+            .blockSizeShift(this.onDemandBlockSizeShift);
+
+        return new OnDemandBlockSnapshotIndexInput(blockIndexInputBuilder, fileInfo, localStoreDir, transferManager);
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/index/store/remote/file/OnDemandBlockIndexInput.java b/server/src/main/java/org/opensearch/index/store/remote/file/OnDemandBlockIndexInput.java
@@ -385,7 +385,7 @@ public static class Builder {
         private int blockSize = 1 << blockSizeShift;
         private int blockMask = blockSize - 1;
 
-        private Builder() {}
+        public Builder() {}
 
         public Builder resourceDescription(String resourceDescription) {
             this.resourceDescription = resourceDescription;
@@ -407,7 +407,7 @@ public Builder length(long length) {
             return this;
         }
 
-        Builder blockSizeShift(int blockSizeShift) {
+        public Builder blockSizeShift(int blockSizeShift) {
             assert blockSizeShift < 31 : "blockSizeShift must be < 31";
             this.blockSizeShift = blockSizeShift;
             this.blockSize = 1 << blockSizeShift;

diff --git a/...src/main/java/org/opensearch/index/store/remote/file/OnDemandBlockSnapshotIndexInput.java b/...src/main/java/org/opensearch/index/store/remote/file/OnDemandBlockSnapshotIndexInput.java
@@ -96,7 +96,7 @@ public OnDemandBlockSnapshotIndexInput(
         );
     }
 
-    OnDemandBlockSnapshotIndexInput(
+    public OnDemandBlockSnapshotIndexInput(
         OnDemandBlockIndexInput.Builder builder,
         FileInfo fileInfo,
         FSDirectory directory,